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

Litellm ryan apr 4
This commit is contained in:
ryan-crabbe-berri
2026-04-04 16:56:03 -07:00
committed by GitHub
5 changed files with 266 additions and 71 deletions
@@ -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"
@@ -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()
@@ -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,14 +188,16 @@ 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: [
@@ -199,7 +205,7 @@ describe("EntityUsageExport utils", () => {
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,
@@ -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,
@@ -17,10 +17,10 @@ import {
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils";
import { isProxyAdminRole } from "@/utils/roles";
import { EditOutlined, InfoCircleOutlined, SaveOutlined } from "@ant-design/icons";
import { EditOutlined, InfoCircleOutlined, MinusCircleOutlined, PlusOutlined, SaveOutlined } from "@ant-design/icons";
import { ArrowLeftIcon } from "@heroicons/react/outline";
import { Badge, Card, Grid, Text, TextInput, Title } from "@tremor/react";
import { Button, Form, Input, Select, Switch, Tabs, Tooltip } from "antd";
import { Button, Form, Input, InputNumber, Select, Space, Switch, Tabs, Tooltip } from "antd";
import MessageManager from "@/components/molecules/message_manager";
import { CheckIcon, CopyIcon } from "lucide-react";
import React, { useEffect, useMemo, useState } from "react";
@@ -199,6 +199,17 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
return org?.members?.some((m: any) => m.user_id === userId && m.user_role === "org_admin") ?? false;
}, [teamData, userOrganizations, userId]);
// Models currently selected in the team edit form, used to scope the per-model
// rate limit dropdown to models this team actually has access to.
const selectedModelsInForm = Form.useWatch("models", form) as string[] | undefined;
const availableRateLimitModels = useMemo(() => {
const selected = selectedModelsInForm ?? teamData?.team_info?.models ?? [];
if (selected.includes("all-proxy-models") || selected.includes("all-team-models")) {
return userModels;
}
return unfurlWildcardModelsInList(selected, userModels);
}, [selectedModelsInForm, teamData, userModels]);
const canEditTeam = is_team_admin || is_proxy_admin || is_org_admin || isOrgAdminForTeam;
const visibleTabs = useMemo(() => getTeamInfoVisibleTabs(canEditTeam), [canEditTeam]);
const defaultTabKey = useMemo(
@@ -471,12 +482,23 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
return v;
};
const modelTpmLimit: Record<string, number> = {};
const modelRpmLimit: Record<string, number> = {};
for (const entry of (values.modelLimits ?? []) as { model?: string; tpm?: number; rpm?: number }[]) {
if (entry?.model) {
if (entry.tpm != null) modelTpmLimit[entry.model] = entry.tpm;
if (entry.rpm != null) modelRpmLimit[entry.model] = entry.rpm;
}
}
const updateData: any = {
team_id: teamId,
team_alias: values.team_alias,
models: values.models,
tpm_limit: sanitizeNumeric(values.tpm_limit),
rpm_limit: sanitizeNumeric(values.rpm_limit),
model_tpm_limit: modelTpmLimit,
model_rpm_limit: modelRpmLimit,
max_budget: values.max_budget,
soft_budget: sanitizeNumeric(values.soft_budget),
budget_duration: values.budget_duration,
@@ -659,6 +681,22 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<Text>TPM: {info.tpm_limit || "Unlimited"}</Text>
<Text>RPM: {info.rpm_limit || "Unlimited"}</Text>
{info.max_parallel_requests && <Text>Max Parallel Requests: {info.max_parallel_requests}</Text>}
{(() => {
const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record<string, number>;
const modelRpm = (info.metadata?.model_rpm_limit ?? {}) as Record<string, number>;
const models = Array.from(new Set([...Object.keys(modelTpm), ...Object.keys(modelRpm)]));
if (models.length === 0) return null;
return (
<div className="mt-3">
<Text className="text-gray-500">Per-model limits:</Text>
{models.map((m) => (
<Text key={m} className="text-xs">
{m}: TPM {modelTpm[m] ?? "—"}, RPM {modelRpm[m] ?? "—"}
</Text>
))}
</div>
);
})()}
</div>
</Card>
@@ -811,6 +849,16 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
models: info.models,
tpm_limit: info.tpm_limit,
rpm_limit: info.rpm_limit,
modelLimits: Array.from(
new Set([
...Object.keys(info.metadata?.model_tpm_limit ?? {}),
...Object.keys(info.metadata?.model_rpm_limit ?? {}),
]),
).map((model) => ({
model,
tpm: info.metadata?.model_tpm_limit?.[model],
rpm: info.metadata?.model_rpm_limit?.[model],
})),
max_budget: info.max_budget,
soft_budget: info.soft_budget,
budget_duration: info.budget_duration,
@@ -827,7 +875,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
: "",
metadata: info.metadata
? JSON.stringify(
(({ logging, secret_manager_settings, soft_budget_alerting_emails, ...rest }) => rest)(info.metadata),
(({ logging, secret_manager_settings, soft_budget_alerting_emails, model_tpm_limit, model_rpm_limit, ...rest }) => rest)(info.metadata),
null,
2,
)
@@ -953,6 +1001,91 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<NumericalInput step={1} style={{ width: "100%" }} />
</Form.Item>
<Form.Item
label="Model-Specific Rate Limits"
tooltip="Set per-model TPM/RPM limits that apply across the whole team."
>
<Form.List name="modelLimits">
{(fields, { add, remove }) => (
<>
{fields.map(({ key, name, ...restField }) => (
<Space
key={key}
style={{ display: "flex", marginBottom: 8 }}
align="baseline"
>
<Form.Item
{...restField}
name={[name, "model"]}
rules={[
{ required: true, message: "Missing model" },
{
validator: (_, value) => {
if (!value) return Promise.resolve();
const all = form.getFieldValue("modelLimits") ?? [];
const dupes = all.filter(
(entry: { model?: string }) => entry?.model === value,
);
if (dupes.length > 1) {
return Promise.reject(new Error("Duplicate model"));
}
return Promise.resolve();
},
},
]}
style={{ minWidth: 240 }}
>
<Select
showSearch
placeholder="Select model"
allowClear
options={availableRateLimitModels.map((m) => ({
value: m,
label: m,
}))}
/>
</Form.Item>
<Form.Item
{...restField}
name={[name, "tpm"]}
rules={[
{
validator: async (_, value) => {
const row = (form.getFieldValue("modelLimits") ?? [])[name] ?? {};
if (row.model && value == null && row.rpm == null) {
return Promise.reject(new Error("Set at least one of TPM or RPM"));
}
return Promise.resolve();
},
},
]}
>
<InputNumber placeholder="TPM Limit" min={0} />
</Form.Item>
<Form.Item {...restField} name={[name, "rpm"]}>
<InputNumber placeholder="RPM Limit" min={0} />
</Form.Item>
<MinusCircleOutlined
onClick={() => remove(name)}
style={{ color: "#ef4444" }}
/>
</Space>
))}
<Form.Item>
<Button
type="dashed"
onClick={() => add()}
block
icon={<PlusOutlined />}
>
Add Model Limit
</Button>
</Form.Item>
</>
)}
</Form.List>
</Form.Item>
<Form.Item
label={
<span>
@@ -1189,6 +1322,22 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<Text className="font-medium">Rate Limits</Text>
<div>TPM: {info.tpm_limit || "Unlimited"}</div>
<div>RPM: {info.rpm_limit || "Unlimited"}</div>
{(() => {
const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record<string, number>;
const modelRpm = (info.metadata?.model_rpm_limit ?? {}) as Record<string, number>;
const models = Array.from(new Set([...Object.keys(modelTpm), ...Object.keys(modelRpm)]));
if (models.length === 0) return null;
return (
<div className="mt-2">
<Text className="text-gray-500">Per-model limits:</Text>
{models.map((m) => (
<div key={m} className="text-xs ml-2">
{m}: TPM {modelTpm[m] ?? "—"}, RPM {modelRpm[m] ?? "—"}
</div>
))}
</div>
);
})()}
</div>
<div>
<Text className="font-medium">Team Budget</Text>