Merge pull request #25156 from BerriAI/litellm_ryan-apr-4

ryan-crabbe-berri · web-flow · commit 6bc4b46e5620 · 2026-04-04T16:56:03.000-07:00
Litellm ryan apr 4
diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json
@@ -406,6 +406,36 @@
         "field_type": "password",
         "options": null,
         "default_value": null
+      },
+      {
+        "key": "tenant_id",
+        "label": "Tenant ID",
+        "placeholder": "Enter your Azure AD tenant ID",
+        "tooltip": "Entra ID (Service Principal) auth. Provide tenant id, client id, and client secret together as an alternative to api key.",
+        "required": false,
+        "field_type": "text",
+        "options": null,
+        "default_value": null
+      },
+      {
+        "key": "client_id",
+        "label": "Client ID",
+        "placeholder": "Enter your Service Principal client ID",
+        "tooltip": "Entra ID (Service Principal) auth. Provide tenant id, client id, and client secret together as an alternative to api key.",
+        "required": false,
+        "field_type": "text",
+        "options": null,
+        "default_value": null
+      },
+      {
+        "key": "client_secret",
+        "label": "Client Secret",
+        "placeholder": "Enter your Service Principal client secret",
+        "tooltip": "Entra ID (Service Principal) auth. Provide tenant id, client id, and client secret together as an alternative to api key.",
+        "required": false,
+        "field_type": "password",
+        "options": null,
+        "default_value": null
       }
     ],
     "default_model_placeholder": "azure/my-deployment"
diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
@@ -110,6 +110,34 @@ def test_watsonx_provider_fields():
     assert "zen_api_key" in field_keys
 
 
+def test_azure_provider_fields_include_entra_id():
+    """Azure provider must expose Entra ID (Service Principal) credential fields so
+    the UI can input tenant_id / client_id / client_secret as an alternative to api_key."""
+    app = FastAPI()
+    app.include_router(router)
+    client = TestClient(app)
+
+    response = client.get("/public/providers/fields")
+    providers = response.json()
+
+    azure = next((p for p in providers if p["provider"] == "Azure"), None)
+    assert azure is not None
+
+    fields_by_key = {f["key"]: f for f in azure["credential_fields"]}
+    # API-key auth still supported
+    assert "api_key" in fields_by_key
+    # Entra ID fields
+    assert "tenant_id" in fields_by_key
+    assert "client_id" in fields_by_key
+    assert "client_secret" in fields_by_key
+    # client_secret must be masked in the UI
+    assert fields_by_key["client_secret"]["field_type"] == "password"
+    # Entra ID is an alternative to api_key, so none of these are individually required
+    assert fields_by_key["tenant_id"]["required"] is False
+    assert fields_by_key["client_id"]["required"] is False
+    assert fields_by_key["client_secret"]["required"] is False
+
+
 def test_public_model_hub_with_healthy_model():
     """Test that health information is populated for a healthy model"""
     app = FastAPI()
diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts
@@ -30,13 +30,17 @@ vi.mock("papaparse", () => ({
 }));
 
 describe("EntityUsageExport utils", () => {
+  // Entity keys match team_ids because that's how the backend shapes team exports
+  // (breakdown.entities is keyed by team_id). The fix under test uses the entity key
+  // directly for display, so the key_alias/team_id in api_key_breakdown metadata is
+  // no longer consulted — it's retained here only to mirror real payload shape.
   const mockSpendData: EntitySpendData = {
     results: [
       {
         date: "2025-01-01",
         breakdown: {
           entities: {
-            entity1: {
+            "team-1": {
               metrics: {
                 spend: 10.5,
                 api_requests: 100,
@@ -64,7 +68,7 @@ describe("EntityUsageExport utils", () => {
                 },
               },
             },
-            entity2: {
+            "team-2": {
               metrics: {
                 spend: 20.3,
                 api_requests: 200,
@@ -99,7 +103,7 @@ describe("EntityUsageExport utils", () => {
         date: "2025-01-02",
         breakdown: {
           entities: {
-            entity1: {
+            "team-1": {
               metrics: {
                 spend: 15.2,
                 api_requests: 150,
@@ -184,22 +188,24 @@ describe("EntityUsageExport utils", () => {
       expect(entity1?.metrics.cache_creation_input_tokens).toBe(75);
     });
 
-    it("should use key alias when available", () => {
+    it("should use entity key as alias when no team alias map is provided", () => {
+      // Non-team exports (tags, orgs, customers, …) pass no teamAliasMap.
+      // For teams, this is also the fallback when a team is missing from the map.
       const result = getEntityBreakdown(mockSpendData);
       const entity1 = result.find((e) => e.metadata.id === "team-1");
 
-      expect(entity1?.metadata.alias).toBe("alias-1");
+      expect(entity1?.metadata.alias).toBe("team-1");
     });
 
-    it("should use team alias map when key alias is not available", () => {
+    it("should use team alias map to resolve alias from entity key", () => {
       const spendDataWithoutAlias: EntitySpendData = {
         ...mockSpendData,
         results: [
           {
             date: "2025-01-01",
             breakdown: {
               entities: {
-                entity1: {
+                "team-1": {
                   metrics: {
                     spend: 10.5,
                     api_requests: 100,
@@ -299,7 +305,7 @@ describe("EntityUsageExport utils", () => {
             date: "2025-01-01",
             breakdown: {
               entities: {
-                entity1: {
+                "team-1": {
                   metrics: {
                     spend: 10.5,
                     api_requests: 100,
@@ -379,15 +385,17 @@ describe("EntityUsageExport utils", () => {
       }
     });
 
-    it("should use dash when team id is not available", () => {
-      const spendDataWithoutTeamId: EntitySpendData = {
+    it("should fall back to the entity key when there is no team alias mapping", () => {
+      // e.g. tag/org/customer exports where teamAliasMap has no entry for the entity,
+      // or a team that isn't in the alias map — the entity key itself is the label.
+      const spendDataWithoutAlias: EntitySpendData = {
         ...mockSpendData,
         results: [
           {
             date: "2025-01-01",
             breakdown: {
               entities: {
-                entity1: {
+                "my-tag": {
                   metrics: {
                     spend: 10.5,
                     api_requests: 100,
@@ -406,11 +414,11 @@ describe("EntityUsageExport utils", () => {
         metadata: mockSpendData.metadata,
       };
 
-      const result = generateDailyData(spendDataWithoutTeamId, "Team");
+      const result = generateDailyData(spendDataWithoutAlias, "Tag");
       const entry = result[0];
 
-      expect(entry["Team ID"]).toBe("-");
-      expect(entry["Team"]).toBe("-");
+      expect(entry["Tag ID"]).toBe("my-tag");
+      expect(entry["Tag"]).toBe("my-tag");
     });
 
     it("should format spend values correctly", () => {
@@ -471,7 +479,7 @@ describe("EntityUsageExport utils", () => {
           date: "2025-01-01",
           breakdown: {
             entities: {
-              entity1: {
+              "team-1": {
                 metrics: {
                   spend: 10.5,
                   api_requests: 100,
@@ -514,7 +522,7 @@ describe("EntityUsageExport utils", () => {
                   },
                 },
               },
-              entity2: {
+              "team-2": {
                 metrics: {
                   spend: 20.3,
                   api_requests: 200,
@@ -549,7 +557,7 @@ describe("EntityUsageExport utils", () => {
           date: "2025-01-02",
           breakdown: {
             entities: {
-              entity1: {
+              "team-1": {
                 metrics: {
                   spend: 15.2,
                   api_requests: 150,
@@ -979,7 +987,7 @@ describe("EntityUsageExport utils", () => {
           date: "2025-01-01",
           breakdown: {
             entities: {
-              entity1: {
+              "team-1": {
                 metrics: {
                   spend: 10.5,
                   api_requests: 100,
diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts
@@ -3,19 +3,16 @@ import type { DateRangePickerValue } from "@tremor/react";
 import Papa from "papaparse";
 import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types";
 
-// Helper function to extract team_id from api_key_breakdown
-const extractTeamIdFromApiKeyBreakdown = (apiKeyBreakdown: Record<string, any> | undefined): string | null => {
-  if (!apiKeyBreakdown) return null;
-
-  // Look through all API keys to find the first non-null team_id
-  for (const apiKeyData of Object.values(apiKeyBreakdown)) {
-    const teamId = (apiKeyData as any)?.metadata?.team_id;
-    if (teamId) {
-      return teamId;
-    }
-  }
-  return null;
-};
+// Resolve display name for an entity. For teams the teamAliasMap provides
+// a human-readable alias; for every other entity type the entity key itself
+// (tag name, org id, customer id, …) is already the correct label.
+const resolveEntityDisplay = (
+  entity: string,
+  teamAliasMap: Record<string, string>,
+): { id: string; alias: string } => ({
+  id: entity,
+  alias: teamAliasMap[entity] || entity,
+});
 
 // Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py).
 // If the backend adds a field, add it here too.
@@ -68,18 +65,7 @@ export const getEntityBreakdown = (
 
   spendData.results.forEach((day) => {
     Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => {
-      // Extract team_id from api_key_breakdown metadata (not data.metadata which is empty)
-      const teamId = extractTeamIdFromApiKeyBreakdown(data.api_key_breakdown) || entity;
-      // Extract key_alias from the first API key that has one
-      const apiKeyBreakdown = data.api_key_breakdown || {};
-      let keyAlias: string | null = null;
-      for (const apiKeyData of Object.values(apiKeyBreakdown)) {
-        const alias = (apiKeyData as any)?.metadata?.key_alias;
-        if (alias) {
-          keyAlias = alias;
-          break;
-        }
-      }
+      const { id, alias } = resolveEntityDisplay(entity, teamAliasMap);
 
       if (!entitySpend[entity]) {
         entitySpend[entity] = {
@@ -95,8 +81,8 @@ export const getEntityBreakdown = (
             cache_creation_input_tokens: 0,
           },
           metadata: {
-            alias: keyAlias || teamAliasMap[teamId] || entity,
-            id: teamId,
+            alias,
+            id,
           },
         };
       }
@@ -124,14 +110,12 @@ export const generateDailyData = (
 
   spendData.results.forEach((day) => {
     Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => {
-      // Extract team_id from api_key_breakdown metadata (not data.metadata which is empty)
-      const teamId = extractTeamIdFromApiKeyBreakdown(data.api_key_breakdown);
-      const teamAlias = teamId ? teamAliasMap[teamId] || null : null;
+      const { id, alias } = resolveEntityDisplay(entity, teamAliasMap);
 
       dailyBreakdown.push({
         Date: day.date,
-        [entityLabel]: teamAlias || "-",
-        [`${entityLabel} ID`]: teamId || "-",
+        [entityLabel]: alias,
+        [`${entityLabel} ID`]: id,
         "Spend ($)": formatNumberWithCommas(data.metrics.spend, 4),
         Requests: data.metrics.api_requests,
         "Successful Requests": data.metrics.successful_requests,
@@ -151,12 +135,12 @@ export const generateDailyWithKeysData = (
   entityLabel: string,
   teamAliasMap: Record<string, string> = {},
 ): any[] => {
-  // Aggregate by unique (Date, Team ID, Key ID) combination to prevent duplicates
+  // Aggregate by unique (Date, Entity ID, Key ID) combination to prevent duplicates
   const aggregatedData: {
     [key: string]: {
       Date: string;
-      teamId: string;
-      teamAlias: string | null;
+      entityId: string;
+      entityAlias: string;
       keyId: string;
       keyAlias: string | null;
       metrics: {
@@ -173,23 +157,22 @@ export const generateDailyWithKeysData = (
 
   spendData.results.forEach((day) => {
     Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => {
+      const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap);
       const apiKeyBreakdown = data.api_key_breakdown || {};
 
       // Iterate through each API key in the breakdown
       Object.entries(apiKeyBreakdown).forEach(([keyId, keyData]: [string, any]) => {
         const keyAlias = keyData?.metadata?.key_alias || null;
-        const teamId = keyData?.metadata?.team_id || entity;
-        const teamAlias = teamId ? teamAliasMap[teamId] || null : null;
 
-        // Create unique key for aggregation: Date_TeamID_KeyID
-        const uniqueKey = `${day.date}_${teamId}_${keyId}`;
+        // Create unique key for aggregation: Date_EntityID_KeyID
+        const uniqueKey = `${day.date}_${entityId}_${keyId}`;
 
         if (!aggregatedData[uniqueKey]) {
-          // First time seeing this (Date, Team ID, Key ID) combination
+          // First time seeing this (Date, Entity ID, Key ID) combination
           aggregatedData[uniqueKey] = {
             Date: day.date,
-            teamId,
-            teamAlias,
+            entityId,
+            entityAlias,
             keyId,
             keyAlias,
             metrics: {
@@ -219,8 +202,8 @@ export const generateDailyWithKeysData = (
   // Convert aggregated data to array format
   const dailyKeyBreakdown = Object.values(aggregatedData).map((item) => ({
     Date: item.Date,
-    [entityLabel]: item.teamAlias || "-",
-    [`${entityLabel} ID`]: item.teamId || "-",
+    [entityLabel]: item.entityAlias,
+    [`${entityLabel} ID`]: item.entityId,
     "Key Alias": item.keyAlias || "-",
     "Key ID": item.keyId,
     "Spend ($)": formatNumberWithCommas(item.metrics.spend, 4),
@@ -273,16 +256,13 @@ export const generateDailyWithModelsData = (
     });
 
     Object.entries(dailyEntityModels).forEach(([entity, models]) => {
-      const entityData = resolveEntities(day.breakdown)[entity];
-      // Extract team_id from api_key_breakdown metadata (not entityData.metadata which is empty)
-      const teamId = extractTeamIdFromApiKeyBreakdown(entityData?.api_key_breakdown);
-      const teamAlias = teamId ? teamAliasMap[teamId] || null : null;
+      const { id, alias } = resolveEntityDisplay(entity, teamAliasMap);
 
       Object.entries(models).forEach(([model, metrics]: [string, any]) => {
         dailyModelBreakdown.push({
           Date: day.date,
-          [entityLabel]: teamAlias || "-",
-          [`${entityLabel} ID`]: teamId || "-",
+          [entityLabel]: alias,
+          [`${entityLabel} ID`]: id,
           Model: model,
           "Spend ($)": formatNumberWithCommas(metrics.spend, 4),
           Requests: metrics.requests,
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
