Merge pull request #18404 from BerriAI/litellm_ui_use_team

[Feature] Resolve Team Alias in Organization Info View + Deprecate useTeams

ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts

@@ -0,0 +1,17 @@
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { Team } from "@/components/key_team_helpers/key_list";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import { fetchTeams } from "@/app/(dashboard)/networking";
+import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory";
+
+const teamKeys = createQueryKeys("teams");
+
+export const useTeams = (): UseQueryResult<Team[]> => {
+  const { accessToken, userId: userID, userRole } = useAuthorized();
+
+  return useQuery<Team[]>({
+    queryKey: teamKeys.list({}),
+    queryFn: async () => await fetchTeams(accessToken!, userID, userRole, null),
+    enabled: Boolean(accessToken),
+  });
+};

ui/litellm-dashboard/src/app/(dashboard)/hooks/useTeams.tsx

@@ -3,6 +3,10 @@ import { Team } from "@/components/key_team_helpers/key_list";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { fetchTeams } from "@/app/(dashboard)/networking";
 
+/**
+ * @deprecated This hook is deprecated. Use the react-query implementation from `@/app/(dashboard)/hooks/teams/useTeams` instead.
+ * This version will be removed in a future release.
+ */
 const useTeams = () => {
   const [teams, setTeams] = useState<Team[] | null>([]);
   const { accessToken, userId: userID, userRole } = useAuthorized();

ui/litellm-dashboard/src/components/organization/organization_view.test.tsx

@@ -40,6 +40,24 @@ vi.mock("../mcp_server_management/MCPServerSelector", () => ({
   __esModule: true,
   default: () => null,
 }));
+const mockUseTeamsData = {
+  data: [
+    {
+      team_id: "team_123",
+      team_alias: "Engineering Team",
+    },
+    {
+      team_id: "team_456",
+      team_alias: "Marketing Team",
+    },
+  ],
+};
+
+const mockUseTeams = vi.fn(() => mockUseTeamsData);
+
+vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
+  useTeams: () => mockUseTeams(),
+}));
 
 const mockOrg = {
   organization_alias: "Acme Corp",
@@ -103,3 +121,63 @@ test("should display empty state when organization has no members", async () =>
     expect(screen.getByText("No members found")).toBeInTheDocument();
   });
 });
+
+test("should display team aliases when teams are available", async () => {
+  const { organizationInfoCall } = await import("../networking");
+  const orgWithTeams = {
+    ...mockOrg,
+    teams: [{ team_id: "team_123" }, { team_id: "team_456" }],
+  };
+  (organizationInfoCall as unknown as ReturnType<typeof vi.fn>).mockResolvedValueOnce(orgWithTeams);
+
+  render(
+    <OrganizationInfoView
+      organizationId="org_123"
+      onClose={() => {}}
+      accessToken="test-token"
+      is_org_admin={false}
+      is_proxy_admin={false}
+      userModels={[]}
+      editOrg={false}
+    />,
+  );
+
+  await waitFor(() => {
+    expect(screen.getByText("Engineering Team")).toBeInTheDocument();
+    expect(screen.getByText("Marketing Team")).toBeInTheDocument();
+  });
+});
+
+test("should display team ID as fallback when alias is not found", async () => {
+  const { organizationInfoCall } = await import("../networking");
+  mockUseTeams.mockReturnValueOnce({
+    data: [
+      {
+        team_id: "team_123",
+        team_alias: "Engineering Team",
+      },
+    ],
+  });
+
+  const orgWithUnknownTeam = {
+    ...mockOrg,
+    teams: [{ team_id: "team_999" }],
+  };
+  (organizationInfoCall as unknown as ReturnType<typeof vi.fn>).mockResolvedValueOnce(orgWithUnknownTeam);
+
+  render(
+    <OrganizationInfoView
+      organizationId="org_123"
+      onClose={() => {}}
+      accessToken="test-token"
+      is_org_admin={false}
+      is_proxy_admin={false}
+      userModels={[]}
+      editOrg={false}
+    />,
+  );
+
+  await waitFor(() => {
+    expect(screen.getByText("team_999")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/organization/organization_view.tsx

@@ -23,7 +23,7 @@ import {
 } from "@tremor/react";
 import { Button, Form, Input, Select } from "antd";
 import { CheckIcon, CopyIcon } from "lucide-react";
-import React, { useEffect, useState } from "react";
+import React, { useEffect, useState, useMemo } from "react";
 import UserSearchModal from "../common_components/user_search_modal";
 import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
 import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
@@ -41,6 +41,8 @@ import ObjectPermissionsView from "../object_permissions_view";
 import NumericalInput from "../shared/numerical_input";
 import MemberModal from "../team/EditMembership";
 import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
+import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
+import { createTeamAliasMap } from "@/utils/teamUtils";
 
 interface OrganizationInfoProps {
   organizationId: string;
@@ -71,6 +73,9 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
   const [copiedStates, setCopiedStates] = useState<Record<string, boolean>>({});
   const [isOrgSaving, setIsOrgSaving] = useState(false);
   const canEditOrg = is_org_admin || is_proxy_admin;
+  const { data: teams } = useTeams();
+
+  const teamAliasMap = useMemo(() => createTeamAliasMap(teams), [teams]);
 
   const fetchOrgInfo = async () => {
     try {
@@ -310,7 +315,7 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
                 <div className="mt-2 flex flex-wrap gap-2">
                   {orgData.teams?.map((team, index) => (
                     <Badge key={index} color="red">
-                      {team.team_id}
+                      {teamAliasMap[team.team_id] || team.team_id}
                     </Badge>
                   ))}
                 </div>

ui/litellm-dashboard/src/utils/teamUtils.test.ts

@@ -1,6 +1,6 @@
 import { describe, expect, it } from "vitest";
-import { resolveTeamAliasFromTeamID } from "./teamUtils";
-import type { Team } from "@/components/networking";
+import { resolveTeamAliasFromTeamID, createTeamAliasMap } from "./teamUtils";
+import type { Team } from "@/components/key_team_helpers/key_list";
 
 describe("resolveTeamAliasFromTeamID", () => {
   it("should return team alias when team is found", () => {
@@ -35,3 +35,30 @@ describe("resolveTeamAliasFromTeamID", () => {
     expect(result).toBeNull();
   });
 });
+
+describe("createTeamAliasMap", () => {
+  it("should create a map from team_id to team_alias", () => {
+    const teams = [
+      {
+        team_id: "team1",
+        team_alias: "Team One",
+      },
+      {
+        team_id: "team2",
+        team_alias: "Team Two",
+      },
+    ] as unknown as Team[];
+
+    const result = createTeamAliasMap(teams);
+    expect(result).toEqual({
+      team1: "Team One",
+      team2: "Team Two",
+    });
+  });
+
+  it("should return empty object when teams is null or undefined", () => {
+    expect(createTeamAliasMap(null)).toEqual({});
+    expect(createTeamAliasMap(undefined)).toEqual({});
+    expect(createTeamAliasMap([])).toEqual({});
+  });
+});

ui/litellm-dashboard/src/utils/teamUtils.ts

@@ -1,5 +1,27 @@
-import { Team } from "@/components/networking";
+import { Team } from "@/components/key_team_helpers/key_list";
 
+/**
+ * Creates a map from team_id to team_alias for efficient lookups.
+ * @param teams - Array of Team objects
+ * @returns Record mapping team_id to team_alias
+ */
+export const createTeamAliasMap = (teams: Team[] | null | undefined): Record<string, string> => {
+  if (!teams) return {};
+  return teams.reduce(
+    (acc, team) => {
+      acc[team.team_id] = team.team_alias;
+      return acc;
+    },
+    {} as Record<string, string>,
+  );
+};
+
+/**
+ * Resolves a team alias from a team ID.
+ * @param teamID - The team ID to look up
+ * @param teams - Array of Team objects
+ * @returns The team alias if found, null otherwise
+ */
 export const resolveTeamAliasFromTeamID = (teamID: string, teams: Team[]): string | null => {
   const team = teams.find((team) => team.team_id === teamID);
   return team ? team.team_alias : null;