feat: enhancements to agent flow on LiteLLm

ui/litellm-dashboard/src/app/page.tsx

@@ -547,7 +547,7 @@ function CreateKeyPageContent() {
                   ) : page == "policies" ? (
                     <PoliciesPanel accessToken={accessToken} userRole={userRole} />
                   ) : page == "agents" ? (
-                    <AgentsPanel accessToken={accessToken} userRole={userRole} />
+                    <AgentsPanel accessToken={accessToken} userRole={userRole} teams={teams} />
                   ) : page == "prompts" ? (
                     <PromptsPanel accessToken={accessToken} userRole={userRole} />
                   ) : page == "transform-request" ? (

ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx

@@ -125,6 +125,208 @@ describe("EntityUsage", () => {
     },
   };
 
+  const mockAgentSpendData = {
+    results: [
+      {
+        date: "2025-01-01",
+        metrics: {
+          spend: 245.8,
+          api_requests: 3200,
+          successful_requests: 3100,
+          failed_requests: 100,
+          total_tokens: 1250000,
+          prompt_tokens: 850000,
+          completion_tokens: 400000,
+          cache_read_input_tokens: 50000,
+          cache_creation_input_tokens: 10000,
+        },
+        breakdown: {
+          entities: {
+            "agent-code-review": {
+              metrics: {
+                spend: 120.4,
+                api_requests: 1500,
+                successful_requests: 1450,
+                failed_requests: 50,
+                total_tokens: 620000,
+                prompt_tokens: 420000,
+                completion_tokens: 200000,
+                cache_read_input_tokens: 30000,
+                cache_creation_input_tokens: 5000,
+              },
+              metadata: { agent_name: "Code Review Agent" },
+              api_key_breakdown: {},
+            },
+            "agent-customer-support": {
+              metrics: {
+                spend: 85.2,
+                api_requests: 1200,
+                successful_requests: 1170,
+                failed_requests: 30,
+                total_tokens: 430000,
+                prompt_tokens: 290000,
+                completion_tokens: 140000,
+                cache_read_input_tokens: 15000,
+                cache_creation_input_tokens: 3000,
+              },
+              metadata: { agent_name: "Customer Support Agent" },
+              api_key_breakdown: {},
+            },
+            "agent-data-analyst": {
+              metrics: {
+                spend: 40.2,
+                api_requests: 500,
+                successful_requests: 480,
+                failed_requests: 20,
+                total_tokens: 200000,
+                prompt_tokens: 140000,
+                completion_tokens: 60000,
+                cache_read_input_tokens: 5000,
+                cache_creation_input_tokens: 2000,
+              },
+              metadata: { agent_name: "Data Analyst Agent" },
+              api_key_breakdown: {},
+            },
+          },
+          models: {
+            "gpt-4o": {
+              metrics: {
+                spend: 180.0,
+                api_requests: 2000,
+                successful_requests: 1950,
+                failed_requests: 50,
+                total_tokens: 900000,
+                prompt_tokens: 600000,
+                completion_tokens: 300000,
+                cache_read_input_tokens: 40000,
+                cache_creation_input_tokens: 8000,
+              },
+              metadata: {},
+              api_key_breakdown: {},
+            },
+            "claude-sonnet-4-20250514": {
+              metrics: {
+                spend: 65.8,
+                api_requests: 1200,
+                successful_requests: 1150,
+                failed_requests: 50,
+                total_tokens: 350000,
+                prompt_tokens: 250000,
+                completion_tokens: 100000,
+                cache_read_input_tokens: 10000,
+                cache_creation_input_tokens: 2000,
+              },
+              metadata: {},
+              api_key_breakdown: {},
+            },
+          },
+          api_keys: {},
+          providers: {
+            openai: {
+              metrics: {
+                spend: 180.0,
+                api_requests: 2000,
+                successful_requests: 1950,
+                failed_requests: 50,
+                total_tokens: 900000,
+                prompt_tokens: 600000,
+                completion_tokens: 300000,
+                cache_read_input_tokens: 40000,
+                cache_creation_input_tokens: 8000,
+              },
+            },
+            anthropic: {
+              metrics: {
+                spend: 65.8,
+                api_requests: 1200,
+                successful_requests: 1150,
+                failed_requests: 50,
+                total_tokens: 350000,
+                prompt_tokens: 250000,
+                completion_tokens: 100000,
+                cache_read_input_tokens: 10000,
+                cache_creation_input_tokens: 2000,
+              },
+            },
+          },
+        },
+      },
+      {
+        date: "2025-01-02",
+        metrics: {
+          spend: 198.5,
+          api_requests: 2800,
+          successful_requests: 2720,
+          failed_requests: 80,
+          total_tokens: 980000,
+          prompt_tokens: 670000,
+          completion_tokens: 310000,
+          cache_read_input_tokens: 42000,
+          cache_creation_input_tokens: 9000,
+        },
+        breakdown: {
+          entities: {
+            "agent-code-review": {
+              metrics: {
+                spend: 95.3,
+                api_requests: 1300,
+                successful_requests: 1270,
+                failed_requests: 30,
+                total_tokens: 510000,
+                prompt_tokens: 350000,
+                completion_tokens: 160000,
+                cache_read_input_tokens: 25000,
+                cache_creation_input_tokens: 4000,
+              },
+              metadata: { agent_name: "Code Review Agent" },
+              api_key_breakdown: {},
+            },
+            "agent-customer-support": {
+              metrics: {
+                spend: 68.7,
+                api_requests: 1000,
+                successful_requests: 970,
+                failed_requests: 30,
+                total_tokens: 320000,
+                prompt_tokens: 220000,
+                completion_tokens: 100000,
+                cache_read_input_tokens: 12000,
+                cache_creation_input_tokens: 3000,
+              },
+              metadata: { agent_name: "Customer Support Agent" },
+              api_key_breakdown: {},
+            },
+            "agent-data-analyst": {
+              metrics: {
+                spend: 34.5,
+                api_requests: 500,
+                successful_requests: 480,
+                failed_requests: 20,
+                total_tokens: 150000,
+                prompt_tokens: 100000,
+                completion_tokens: 50000,
+                cache_read_input_tokens: 5000,
+                cache_creation_input_tokens: 2000,
+              },
+              metadata: { agent_name: "Data Analyst Agent" },
+              api_key_breakdown: {},
+            },
+          },
+          models: {},
+          api_keys: {},
+          providers: {},
+        },
+      },
+    ],
+    metadata: {
+      total_spend: 444.3,
+      total_api_requests: 6000,
+      total_successful_requests: 5820,
+      total_failed_requests: 180,
+      total_tokens: 2230000,
+    },
+  };
+
   const defaultProps = {
     accessToken: "test-token",
     entityType: "tag" as const,
@@ -153,7 +355,7 @@ describe("EntityUsage", () => {
     mockTeamDailyActivityCall.mockResolvedValue(mockSpendData);
     mockOrganizationDailyActivityCall.mockResolvedValue(mockSpendData);
     mockCustomerDailyActivityCall.mockResolvedValue(mockSpendData);
-    mockAgentDailyActivityCall.mockResolvedValue(mockSpendData);
+    mockAgentDailyActivityCall.mockResolvedValue(mockAgentSpendData);
     mockUserDailyActivityCall.mockResolvedValue(mockSpendData);
   });
 
@@ -231,7 +433,7 @@ describe("EntityUsage", () => {
     expect(screen.getByText("Agent Spend Overview")).toBeInTheDocument();
 
     await waitFor(() => {
-      const spendElements = screen.getAllByText("$100.50");
+      const spendElements = screen.getAllByText("$444.30");
       expect(spendElements.length).toBeGreaterThan(0);
     });
   });
@@ -385,6 +587,87 @@ describe("EntityUsage", () => {
     });
   });
 
+  it("should display Agent Activity tab for team entity type", async () => {
+    render(<EntityUsage {...defaultProps} entityType="team" />);
+
+    await waitFor(() => {
+      expect(mockTeamDailyActivityCall).toHaveBeenCalled();
+    });
+
+    expect(screen.getByText("Agent Activity")).toBeInTheDocument();
+  });
+
+  it("should not display Agent Activity tab for non-team entity types", async () => {
+    render(<EntityUsage {...defaultProps} entityType="tag" />);
+
+    await waitFor(() => {
+      expect(mockTagDailyActivityCall).toHaveBeenCalled();
+    });
+
+    expect(screen.queryByText("Agent Activity")).not.toBeInTheDocument();
+  });
+
+  it("should display Top Agents Driving Spend card for team entity type", async () => {
+    render(<EntityUsage {...defaultProps} entityType="team" />);
+
+    await waitFor(() => {
+      expect(mockTeamDailyActivityCall).toHaveBeenCalled();
+    });
+
+    expect(screen.getByText("Top Agents Driving Spend")).toBeInTheDocument();
+  });
+
+  it("should not display Top Agents Driving Spend card for non-team entity types", async () => {
+    render(<EntityUsage {...defaultProps} entityType="tag" />);
+
+    await waitFor(() => {
+      expect(mockTagDailyActivityCall).toHaveBeenCalled();
+    });
+
+    expect(screen.queryByText("Top Agents Driving Spend")).not.toBeInTheDocument();
+  });
+
+  it("should fetch agent activity data when entity type is team", async () => {
+    render(<EntityUsage {...defaultProps} entityType="team" />);
+
+    await waitFor(() => {
+      expect(mockAgentDailyActivityCall).toHaveBeenCalledWith(
+        "test-token",
+        expect.any(Date),
+        expect.any(Date),
+        1,
+        null,
+      );
+    });
+  });
+
+  it("should not fetch agent activity data for non-team entity types", async () => {
+    render(<EntityUsage {...defaultProps} entityType="tag" />);
+
+    await waitFor(() => {
+      expect(mockTagDailyActivityCall).toHaveBeenCalled();
+    });
+
+    expect(mockAgentDailyActivityCall).not.toHaveBeenCalled();
+  });
+
+  it("should switch to Agent Activity tab for team entity type", async () => {
+    render(<EntityUsage {...defaultProps} entityType="team" />);
+
+    await waitFor(() => {
+      expect(mockTeamDailyActivityCall).toHaveBeenCalled();
+    });
+
+    const agentActivityTab = screen.getByText("Agent Activity");
+    act(() => {
+      fireEvent.click(agentActivityTab);
+    });
+
+    await waitFor(() => {
+      expect(screen.getAllByText("Activity Metrics").length).toBeGreaterThan(0);
+    });
+  });
+
   it("should fallback to entity value when no entityList and no team_alias", async () => {
     const spendDataWithoutAlias = {
       ...mockSpendData,

ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx

@@ -100,11 +100,24 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
   });
   const { teams } = useTeams();
 
+  const [agentSpendData, setAgentSpendData] = useState<EntitySpendData>({
+    results: [],
+    metadata: {
+      total_spend: 0,
+      total_api_requests: 0,
+      total_successful_requests: 0,
+      total_failed_requests: 0,
+      total_tokens: 0,
+    },
+  });
+
   const modelMetrics = processActivityData(spendData, "models", teams || []);
   const keyMetrics = processActivityData(spendData, "api_keys", teams || []);
+  const agentMetrics = entityType === "team" ? processActivityData(agentSpendData, "entities", teams || []) : {};
   const [selectedTags, setSelectedTags] = useState<string[]>([]);
   const [topKeysLimit, setTopKeysLimit] = useState<number>(5);
   const [topModelsLimit, setTopModelsLimit] = useState<number>(5);
+  const [topAgentsLimit, setTopAgentsLimit] = useState<number>(5);
 
   const fetchSpendData = async () => {
     if (!accessToken || !dateValue.from || !dateValue.to) return;
@@ -171,8 +184,21 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
     }
   };
 
+  const fetchAgentSpendData = async () => {
+    if (!accessToken || !dateValue.from || !dateValue.to || entityType !== "team") return;
+    const startTime = new Date(dateValue.from);
+    const endTime = new Date(dateValue.to);
+    try {
+      const data = await agentDailyActivityCall(accessToken, startTime, endTime, 1, null);
+      setAgentSpendData(data);
+    } catch (e) {
+      console.error("Failed to fetch agent activity data:", e);
+    }
+  };
+
   useEffect(() => {
     fetchSpendData();
+    fetchAgentSpendData();
   }, [accessToken, dateValue, entityId, selectedTags]);
 
   const getTopModels = () => {
@@ -209,6 +235,37 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
       .slice(0, topModelsLimit);
   };
 
+  const getTopAgents = () => {
+    const agentSpend: { [key: string]: any } = {};
+    agentSpendData.results.forEach((day) => {
+      Object.entries(day.breakdown.entities || {}).forEach(([agentId, data]) => {
+        if (!agentSpend[agentId]) {
+          agentSpend[agentId] = {
+            spend: 0,
+            requests: 0,
+            successful_requests: 0,
+            failed_requests: 0,
+            tokens: 0,
+            agent_name: (data.metadata as any)?.agent_name || agentId,
+          };
+        }
+        agentSpend[agentId].spend += data.metrics.spend;
+        agentSpend[agentId].requests += data.metrics.api_requests;
+        agentSpend[agentId].successful_requests += data.metrics.successful_requests;
+        agentSpend[agentId].failed_requests += data.metrics.failed_requests;
+        agentSpend[agentId].tokens += data.metrics.total_tokens;
+      });
+    });
+
+    return Object.entries(agentSpend)
+      .map(([agentId, metrics]) => ({
+        key: metrics.agent_name,
+        ...metrics,
+      }))
+      .sort((a, b) => b.spend - a.spend)
+      .slice(0, topAgentsLimit);
+  };
+
   const getTopAPIKeys = () => {
     console.log("debugTags", { spendData });
     const keySpend: { [key: string]: KeyMetricWithMetadata } = {};
@@ -408,6 +465,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
         <TabList variant="solid" className="mt-1">
           <Tab>Cost</Tab>
           <Tab>{entityType === "agent" ? "Request / Token Consumption" : "Model Activity"}</Tab>
+          {entityType === "team" && <Tab>Agent Activity</Tab>}
           <Tab>Key Activity</Tab>
           <Tab>Endpoint Activity</Tab>
         </TabList>
@@ -621,6 +679,20 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
                 </Card>
               </Col>
 
+              {/* Top Agents - only for team entity type */}
+              {entityType === "team" && (
+                <Col numColSpan={2}>
+                  <Card>
+                    <Title>Top Agents Driving Spend</Title>
+                    <TopModelView
+                      topModels={getTopAgents()}
+                      topModelsLimit={topAgentsLimit}
+                      setTopModelsLimit={setTopAgentsLimit}
+                    />
+                  </Card>
+                </Col>
+              )}
+
               {/* Spend by Provider */}
               <Col numColSpan={2}>
                 <Card>
@@ -696,6 +768,11 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
           <TabPanel>
             <ActivityMetrics modelMetrics={modelMetrics} hidePromptCachingMetrics={entityType === "agent"} />
           </TabPanel>
+          {entityType === "team" && (
+            <TabPanel>
+              <ActivityMetrics modelMetrics={agentMetrics} />
+            </TabPanel>
+          )}
           <TabPanel>
             <ActivityMetrics modelMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} />
           </TabPanel>

ui/litellm-dashboard/src/components/activity_metrics.tsx

@@ -362,7 +362,7 @@ export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string,
 // Process data function
 export const processActivityData = (
   dailyActivity: { results: DailyData[] },
-  key: "models" | "api_keys" | "mcp_servers",
+  key: "models" | "api_keys" | "mcp_servers" | "entities",
   teams: Team[] = [],
 ): Record<string, ModelActivityData> => {
   const modelMetrics: Record<string, ModelActivityData> = {};
@@ -371,7 +371,11 @@ export const processActivityData = (
     Object.entries(day.breakdown[key] || {}).forEach(([model, modelData]) => {
       if (!modelMetrics[model]) {
         modelMetrics[model] = {
-          label: key === "api_keys" ? formatKeyLabel(modelData as KeyMetricWithMetadata, model, teams) : model,
+          label: key === "api_keys"
+            ? formatKeyLabel(modelData as KeyMetricWithMetadata, model, teams)
+            : key === "entities"
+              ? ((modelData as any).metadata?.agent_name || (modelData as any).metadata?.team_alias || model)
+              : model,
           total_requests: 0,
           total_successful_requests: 0,
           total_failed_requests: 0,

ui/litellm-dashboard/src/components/agents.tsx

@@ -19,19 +19,21 @@ import { isAdminRole } from "@/utils/roles";
 import AgentInfoView from "./agents/agent_info";
 import NotificationsManager from "./molecules/notifications_manager";
 import { Agent, AgentKeyInfo } from "./agents/types";
+import { Team } from "./key_team_helpers/key_list";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
 import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton";
 
 interface AgentsPanelProps {
   accessToken: string | null;
   userRole?: string;
+  teams?: Team[] | null;
 }
 
 interface AgentsResponse {
   agents: Agent[];
 }
 
-const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole }) => {
+const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole, teams }) => {
   const [agentsList, setAgentsList] = useState<Agent[]>([]);
   const [keyInfoMap, setKeyInfoMap] = useState<Record<string, AgentKeyInfo>>({});
   const [isAddModalVisible, setIsAddModalVisible] = useState(false);
@@ -282,6 +284,7 @@ const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole }) => {
         onClose={handleCloseModal}
         accessToken={accessToken}
         onSuccess={handleSuccess}
+        teams={teams}
       />
 
       {agentToDelete && (

ui/litellm-dashboard/src/components/agents/add_agent_form.tsx

@@ -15,11 +15,14 @@ import {
 } from "../networking";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
+import { Team } from "../key_team_helpers/key_list";
+import TeamDropdown from "../common_components/team_dropdown";
 import AgentFormFields from "./agent_form_fields";
 import DynamicAgentFormFields, { buildDynamicAgentData } from "./dynamic_agent_form_fields";
 import { getDefaultFormValues, buildAgentDataFromForm } from "./agent_config";
 import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
 import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions";
+import GuardrailSelector from "../guardrails/GuardrailSelector";
 
 const { Step } = Steps;
 
@@ -30,6 +33,7 @@ interface AddAgentFormProps {
   onClose: () => void;
   accessToken: string | null;
   onSuccess: () => void;
+  teams?: Team[] | null;
 }
 
 const AddAgentForm: React.FC<AddAgentFormProps> = ({
@@ -37,6 +41,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
   onClose,
   accessToken,
   onSuccess,
+  teams,
 }) => {
   const { userId, userRole } = useAuthorized();
   const [form] = Form.useForm();
@@ -268,6 +273,17 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
         }
       }
 
+      const selectedGuardrails = values.guardrails || [];
+      if (selectedGuardrails.length > 0) {
+        if (!agentData.litellm_params) agentData.litellm_params = {};
+        agentData.litellm_params.guardrails = selectedGuardrails;
+      }
+
+      const selectedTeamId = values.team_id || null;
+      if (selectedTeamId) {
+        agentData.team_id = selectedTeamId;
+      }
+
       const agentResponse = await createAgentCall(accessToken, agentData);
       const agentId: string = agentResponse.agent_id;
       const agentName: string = agentResponse.agent_name || values.agent_name || agentId;
@@ -279,6 +295,8 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
           agentId,
           newKeyName,
           newKeyModels,
+          undefined,
+          selectedTeamId,
         );
         setCreatedKeyValue(keyResponse.key || null);
       } else if (keyAssignOption === "existing_key") {
@@ -527,6 +545,22 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
           </div>
         </div>
       </div>
+
+      <Divider className="my-0" />
+
+      <div>
+        <h4 className="text-sm font-medium text-gray-700 mb-3">Guardrails</h4>
+        <p className="text-xs text-gray-500 mb-3">
+          Apply guardrails to this agent. Selected guardrails will run on all calls made by this agent.
+        </p>
+        <Form.Item name="guardrails" initialValue={[]}>
+          <GuardrailSelector
+            accessToken={accessToken ?? ""}
+            value={form.getFieldValue("guardrails") ?? []}
+            onChange={(selected: string[]) => form.setFieldsValue({ guardrails: selected })}
+          />
+        </Form.Item>
+      </div>
     </div>
   );
 
@@ -683,6 +717,19 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
           </Tag>
         </div>
 
+        <Form.Item
+          label={<span className="text-sm font-medium text-gray-700">Assign to Team</span>}
+          name="team_id"
+          tooltip="Optionally assign this agent to a team. The agent and its key will belong to the selected team."
+        >
+          <TeamDropdown
+            teams={teams}
+            loading={!teams}
+          />
+        </Form.Item>
+
+        <Divider className="my-4" />
+
         <div className="space-y-3">
           {/* Option: Create new key */}
           <div
@@ -848,8 +895,8 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
           layout="vertical"
           initialValues={
             agentType === "a2a"
-              ? { ...getDefaultFormValues(), allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] }, mcp_tool_permissions: {}, entitlement_models: [], entitlement_agents: [] }
-              : { allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] }, mcp_tool_permissions: {}, entitlement_models: [], entitlement_agents: [] }
+              ? { ...getDefaultFormValues(), allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] }, mcp_tool_permissions: {}, entitlement_models: [], entitlement_agents: [], guardrails: [] }
+              : { allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] }, mcp_tool_permissions: {}, entitlement_models: [], entitlement_agents: [], guardrails: [] }
           }
           className="space-y-4"
         >

ui/litellm-dashboard/src/components/networking.tsx

@@ -907,6 +907,7 @@ export const keyCreateForAgentCall = async (
   keyAlias: string,
   models: string[],
   metadata?: Record<string, any>,
+  teamId?: string | null,
 ) => {
   const url = proxyBaseUrl ? `${proxyBaseUrl}/key/generate` : `/key/generate`;
   const body: Record<string, any> = {
@@ -914,6 +915,9 @@ export const keyCreateForAgentCall = async (
     key_alias: keyAlias,
     models: models.length > 0 ? models : [],
   };
+  if (teamId) {
+    body.team_id = teamId;
+  }
   if (metadata && Object.keys(metadata).length > 0) {
     body.metadata = metadata;
   }