Merge pull request #40700 from BerriAI/litellm_ui_editable_model_team_id

fix(ui): let admins change a model's team from the model edit page

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

@@ -24,9 +24,11 @@ import CacheControlInjectionPoints, {
   CACHE_CONTROL_TOOLTIP,
   type CacheControlInjectionPoint,
 } from "./add_model/cache_control_settings";
+import type { Team } from "./key_team_helpers/key_list";
 import type { CredentialItem } from "./networking";
 import NumericalInput from "./shared/numerical_input";
 import type { Tag } from "./tag_management/types";
+import { ModelTeamSelect } from "./view_model/ModelTeamSelect";
 import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
 import { formatPtuUtcDisplay, utcIsoToPickerValue } from "../utils/ptuDatetime";
 import { isMaskedSecret } from "../utils/maskedSecretUtils";
@@ -103,6 +105,7 @@ export interface ModelEditFormValues {
   litellm_credential_name?: string;
   litellm_extra_params?: string;
   model_info?: string;
+  team_id?: string;
 }
 
 type ModelEditFieldName = keyof ModelEditFormValues;
@@ -139,6 +142,7 @@ const modelEditShape = {
   litellm_credential_name: textish,
   litellm_extra_params: textish,
   model_info: textish,
+  team_id: textish,
 };
 
 const isJson = (value: string): boolean => {
@@ -260,6 +264,7 @@ export const toModelEditFormValues = (localModelData: any, isWildcardModel: bool
     null,
     2,
   ),
+  team_id: localModelData.model_info?.team_id ?? undefined,
 });
 
 const displayCost = (localModelData: any, field: TouchedPricingField): string => {
@@ -286,6 +291,7 @@ interface ModelInfoEditFormProps {
   tagsList: Record<string, Tag>;
   credentialsList: CredentialItem[];
   healthCheckModelOptions: { value: string; label: string }[];
+  teams: Team[] | null;
 }
 
 const Display: React.FC<{ children: React.ReactNode }> = ({ children }) => (
@@ -357,6 +363,7 @@ const ModelInfoEditForm: React.FC<ModelInfoEditFormProps> = ({
   tagsList,
   credentialsList,
   healthCheckModelOptions,
+  teams,
 }) => {
   // Neither RHF's blur-based touchedFields nor its resettable dirtyFields matches antd's touched-on-change.
   const touchedRef = React.useRef<ReadonlySet<string>>(new Set<string>());
@@ -802,11 +809,19 @@ const ModelInfoEditForm: React.FC<ModelInfoEditFormProps> = ({
 
             <div>
               <FieldLabel>Team</FieldLabel>
-              <Display>
-                {teamAlias
-                  ? `${teamAlias} (${modelData.model_info.team_id})`
-                  : modelData.model_info.team_id || "Not Set"}
-              </Display>
+              {isEditing ? (
+                <FormField control={form.control} name="team_id">
+                  {({ id, value, onChange, onBlur }) => (
+                    <ModelTeamSelect id={id} value={value} onChange={onChange} onBlur={onBlur} teams={teams} />
+                  )}
+                </FormField>
+              ) : (
+                <Display>
+                  {teamAlias
+                    ? `${teamAlias} (${localModelData.model_info?.team_id})`
+                    : localModelData.model_info?.team_id || "Not Set"}
+                </Display>
+              )}
             </div>
           </div>
 

ui/litellm-dashboard/src/components/model_info_view.test.tsx

@@ -1509,6 +1509,11 @@ describe("ModelInfoView", () => {
       expect(await screen.findByRole("button", { name: /save changes/i })).toBeInTheDocument();
     };
 
+    const openSelect = async (user: ReturnType<typeof userEvent.setup>, triggerText: string) => {
+      await user.click(await screen.findByText(triggerText));
+      await screen.findByRole("combobox", { expanded: true });
+    };
+
     const save = async (user: ReturnType<typeof userEvent.setup>) => {
       await user.click(screen.getByRole("button", { name: /save changes/i }));
       await waitFor(() => expect(mockModelPatchUpdateCall).toHaveBeenCalled());
@@ -1642,6 +1647,97 @@ describe("ModelInfoView", () => {
       expect(payload.model_info).toMatchObject({ team_id: "team-7" });
     });
 
+    it("sends the team picked in the Team ID selector", async () => {
+      mockUseTeams.mockReturnValue({
+        data: [
+          { team_id: "team-1", team_alias: "alpha" },
+          { team_id: "team-2", team_alias: "beta" },
+        ],
+        isLoading: false,
+        error: null,
+      });
+      const teamModel = {
+        ...defaultModelData,
+        model_info: { ...defaultModelData.model_info, team_id: "team-1" },
+      };
+      mockUseModelsInfo.mockReturnValue({ data: { data: [teamModel] }, isLoading: false, error: null });
+      mockModelInfoV1Call.mockResolvedValue({ data: [teamModel] });
+      const user = userEvent.setup();
+      await enterEditMode(user);
+
+      await openSelect(user, "alpha (team-1)");
+      await user.click(await screen.findByText("beta (team-2)"));
+
+      const payload = await save(user);
+
+      expect(payload.model_info.team_id).toBe("team-2");
+    });
+
+    it("shows the picked team in read mode right after saving", async () => {
+      mockUseTeams.mockReturnValue({
+        data: [
+          { team_id: "team-1", team_alias: "alpha" },
+          { team_id: "team-2", team_alias: "beta" },
+        ],
+        isLoading: false,
+        error: null,
+      });
+      const teamModel = {
+        ...defaultModelData,
+        model_info: { ...defaultModelData.model_info, team_id: "team-1" },
+      };
+      mockUseModelsInfo.mockReturnValue({ data: { data: [teamModel] }, isLoading: false, error: null });
+      mockModelInfoV1Call.mockResolvedValue({ data: [teamModel] });
+      const user = userEvent.setup();
+      await enterEditMode(user);
+
+      await openSelect(user, "alpha (team-1)");
+      await user.click(await screen.findByText("beta (team-2)"));
+      await save(user);
+
+      expect(await screen.findByRole("button", { name: /edit settings/i })).toBeInTheDocument();
+      expect(screen.getByText("beta (team-2)")).toBeInTheDocument();
+      expect(screen.queryByText("alpha (team-1)")).not.toBeInTheDocument();
+    });
+
+    it("shows the Team ID placeholder for a model with no team", async () => {
+      mockUseTeams.mockReturnValue({
+        data: [{ team_id: "team-1", team_alias: "alpha" }],
+        isLoading: false,
+        error: null,
+      });
+      const user = userEvent.setup();
+      await enterEditMode(user);
+
+      expect(screen.getByText("Select a team")).toBeInTheDocument();
+    });
+
+    it.each(["Internal User", "Org Admin"])("only offers a %s the teams they administer", async (userRole) => {
+      mockUseTeams.mockReturnValue({
+        data: [
+          { team_id: "team-1", team_alias: "alpha", members_with_roles: [{ user_id: "123", role: "admin" }] },
+          { team_id: "team-2", team_alias: "beta", members_with_roles: [{ user_id: "123", role: "user" }] },
+          { team_id: "team-3", team_alias: "gamma", members_with_roles: [{ user_id: "123", role: "admin" }] },
+        ],
+        isLoading: false,
+        error: null,
+      });
+      const teamModel = {
+        ...defaultModelData,
+        model_info: { ...defaultModelData.model_info, team_id: "team-1" },
+      };
+      mockUseModelsInfo.mockReturnValue({ data: { data: [teamModel] }, isLoading: false, error: null });
+      mockModelInfoV1Call.mockResolvedValue({ data: [teamModel] });
+      const user = userEvent.setup();
+      render(<ModelInfoView {...DEFAULT_ADMIN_PROPS} userRole={userRole} />, { wrapper });
+      await user.click(await screen.findByRole("button", { name: /edit settings/i }));
+
+      await user.click(await screen.findByText("alpha (team-1)"));
+
+      expect(await screen.findByRole("option", { name: "gamma (team-3)" })).toBeInTheDocument();
+      expect(screen.queryByRole("option", { name: "beta (team-2)" })).not.toBeInTheDocument();
+    });
+
     it("sends the edited LiteLLM extra params", async () => {
       const user = userEvent.setup();
       await enterEditMode(user);
@@ -1670,7 +1766,7 @@ describe("ModelInfoView", () => {
       const user = userEvent.setup();
       await enterEditMode(user);
 
-      await user.click(await screen.findByText("selected-credential"));
+      await openSelect(user, "selected-credential");
       await user.click(await screen.findByText("other-credential"));
 
       const payload = await save(user);
@@ -1706,7 +1802,7 @@ describe("ModelInfoView", () => {
       const user = userEvent.setup();
       await enterEditMode(user);
 
-      await user.click(screen.getByText("Select existing health check model"));
+      await openSelect(user, "Select existing health check model");
       await user.click(await screen.findByText("openai/gpt-4o"));
 
       const payload = await save(user);

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

@@ -22,6 +22,7 @@ import {
   isComplexityRouter as isComplexityRouterParams,
 } from "./add_model/auto_router_strategies";
 import { canEditAutoRouter, canModifyModel } from "@/utils/modelPermissions";
+import { teamsUserCanAssign } from "@/utils/roles";
 import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
 import DeleteResourceModal from "./common_components/DeleteResourceModal";
 import EditAutoRouterModal from "./edit_auto_router/edit_auto_router_modal";
@@ -114,7 +115,9 @@ export default function ModelInfoView({
   // Keep modelData variable name for backwards compatibility
   const modelData = transformedModelData;
 
-  const teamAlias = teams?.find((team) => team.team_id === modelData?.model_info?.team_id)?.team_alias || null;
+  const aliasForTeam = (teamId: string | null | undefined): string | null =>
+    teams?.find((team) => team.team_id === teamId)?.team_alias || null;
+  const teamAlias = aliasForTeam(modelData?.model_info?.team_id);
   const rawModelInfoEntries = Object.entries(modelData?.model_info ?? {}).flatMap((entry) =>
     entry[0] === "team_id" && teamAlias ? [entry, ["team_alias", teamAlias]] : [entry],
   );
@@ -130,6 +133,7 @@ export default function ModelInfoView({
   };
   const canEditModel = canModifyModel(actor, teams ?? null, origin);
   const canEditRouter = canEditAutoRouter(actor, teams ?? null, origin);
+  const assignableTeams = useMemo(() => teamsUserCanAssign(teams ?? null, userRole, userID), [teams, userRole, userID]);
   // Editor-aware on purpose: an adaptive or quality router must not offer Edit Auto Router.
   const isAutoRouterModel = hasAutoRouterEditor(modelData?.litellm_params);
   // Broader than the editor check: adaptive and quality routers equally have no upstream
@@ -379,6 +383,7 @@ export default function ModelInfoView({
             health_check_model: values.health_check_model,
           };
         }
+        if (values.team_id) updatedModelInfo = { ...updatedModelInfo, team_id: values.team_id };
         updatedModelInfo = applyPtuModelInfo(updatedModelInfo, values, ptuCostAttributionEnabled);
       } catch (e) {
         toast.fromError("Invalid JSON in Model Info");
@@ -724,7 +729,7 @@ export default function ModelInfoView({
                 <ModelInfoEditForm
                   localModelData={localModelData}
                   modelData={modelData}
-                  teamAlias={teamAlias}
+                  teamAlias={aliasForTeam(localModelData.model_info?.team_id)}
                   accessToken={accessToken}
                   isEditing={isEditing}
                   isSaving={isSaving}
@@ -739,6 +744,7 @@ export default function ModelInfoView({
                   tagsList={tagsList}
                   credentialsList={credentialsList}
                   healthCheckModelOptions={healthCheckModelOptions}
+                  teams={assignableTeams}
                 />
               ) : (
                 <p className="text-sm">Loading...</p>

ui/litellm-dashboard/src/components/view_model/ModelTeamSelect.tsx

@@ -0,0 +1,31 @@
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import type { Team } from "../key_team_helpers/key_list";
+
+interface ModelTeamSelectProps {
+  id: string;
+  value: string | undefined;
+  onChange: (teamId: string) => void;
+  onBlur: () => void;
+  teams: Team[] | null;
+}
+
+export const ModelTeamSelect: React.FC<ModelTeamSelectProps> = ({ id, value, onChange, onBlur, teams }) => {
+  const items = (teams ?? []).map((team) => ({
+    value: team.team_id,
+    label: team.team_alias ? `${team.team_alias} (${team.team_id})` : team.team_id,
+  }));
+  return (
+    <Select items={items} value={value || null} onValueChange={(selected: string | null) => onChange(selected ?? "")}>
+      <SelectTrigger id={id} className="w-full" onBlur={onBlur}>
+        <SelectValue placeholder="Select a team" />
+      </SelectTrigger>
+      <SelectContent>
+        {items.map((item) => (
+          <SelectItem key={item.value} value={item.value}>
+            {item.label}
+          </SelectItem>
+        ))}
+      </SelectContent>
+    </Select>
+  );
+};

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

@@ -43,6 +43,17 @@ export const isUserTeamAdminForSingleTeam = (teamMemberWithRoles: Member[] | nul
   return teamMemberWithRoles.some((member) => member.user_id === userID && member.role === "admin");
 };
 
+export const teamsUserCanAssign = (
+  teams: Team[] | null,
+  userRole: string | null,
+  userID: string | null,
+): Team[] | null => {
+  if (teams == null || isProxyAdminRole(userRole ?? "")) {
+    return teams;
+  }
+  return teams.filter((team) => isUserTeamAdminForSingleTeam(team.members_with_roles, userID ?? ""));
+};
+
 export const isOrgAdminForAnyOrg = (
   organizations: Organization[] | null | undefined,
   userID: string | null | undefined,