Change edit team models to match create team models

ui/litellm-dashboard/src/components/team/team_info.test.tsx

@@ -13,6 +13,7 @@ vi.mock("@/components/networking", () => ({
   getGuardrailsList: vi.fn(),
   fetchMCPAccessGroups: vi.fn(),
   getTeamPermissionsCall: vi.fn(),
+  organizationInfoCall: vi.fn(),
 }));
 
 describe("TeamInfoView", () => {
@@ -161,6 +162,116 @@ describe("TeamInfoView", () => {
 
     const allProxyModelsOption = screen.queryByText("All Proxy Models");
     expect(allProxyModelsOption).not.toBeInTheDocument();
-  }, // This is a workaround to fix the flaky test issue. TODO: Remove this once we have a better solution.
-  10000);
+  }, 10000); // This is a workaround to fix the flaky test issue. TODO: Remove this once we have a better solution.
+
+  it("should only show organization models in dropdown when team is in organization with limited models", async () => {
+    const organizationId = "org-123";
+    const organizationModels = ["gpt-4", "claude-3-opus"];
+    const userModels = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus", "claude-2"];
+
+    vi.mocked(networking.teamInfoCall).mockResolvedValue({
+      team_id: "123",
+      team_info: {
+        team_alias: "Test Team",
+        team_id: "123",
+        organization_id: organizationId,
+        admins: ["admin@test.com"],
+        members: ["user1@test.com"],
+        members_with_roles: [
+          {
+            user_id: "user1@test.com",
+            user_email: "user1@test.com",
+            role: "member",
+            spend: 0,
+            budget_id: "budget1",
+          },
+        ],
+        metadata: {},
+        tpm_limit: null,
+        rpm_limit: null,
+        max_budget: null,
+        budget_duration: null,
+        models: ["gpt-4"],
+        blocked: false,
+        spend: 0,
+        max_parallel_requests: null,
+        budget_reset_at: null,
+        model_id: null,
+        litellm_model_table: null,
+        created_at: "2024-01-01T00:00:00Z",
+        team_member_budget_table: null,
+      },
+      keys: [],
+      team_memberships: [],
+    });
+
+    vi.mocked(networking.organizationInfoCall).mockResolvedValue({
+      organization_id: organizationId,
+      organization_name: "Test Organization",
+      spend: 0,
+      max_budget: null,
+      models: organizationModels,
+      tpm_limit: null,
+      rpm_limit: null,
+      members: null,
+    });
+
+    vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] });
+    vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]);
+
+    render(
+      <TeamInfoView
+        teamId="123"
+        onUpdate={() => {}}
+        onClose={() => {}}
+        accessToken="123"
+        is_team_admin={true}
+        is_proxy_admin={true}
+        userModels={userModels}
+        editTeam={false}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getAllByText("Test Team")).not.toBeNull();
+    });
+
+    const settingsTab = screen.getByRole("tab", { name: "Settings" });
+    act(() => {
+      fireEvent.click(settingsTab);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Team Settings")).toBeInTheDocument();
+    });
+
+    const editButton = screen.getByRole("button", { name: "Edit Settings" });
+    act(() => {
+      fireEvent.click(editButton);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Models")).toBeInTheDocument();
+    });
+
+    const modelsSelect = screen.getByLabelText("Models");
+    act(() => {
+      fireEvent.mouseDown(modelsSelect);
+    });
+
+    await waitFor(() => {
+      const dropdownOptions = screen.getAllByRole("option");
+      const optionTexts = dropdownOptions.map((option) => option.textContent);
+
+      organizationModels.forEach((model) => {
+        expect(optionTexts).toContain(model);
+      });
+
+      const modelsNotInOrganization = userModels.filter((m) => !organizationModels.includes(m));
+      modelsNotInOrganization.forEach((model) => {
+        expect(optionTexts).not.toContain(model);
+      });
+    });
+  }, 10000);
 });

ui/litellm-dashboard/src/components/team/team_info.tsx

@@ -2,6 +2,8 @@ import UserSearchModal from "@/components/common_components/user_search_modal";
 import {
   getGuardrailsList,
   Member,
+  Organization,
+  organizationInfoCall,
   teamInfoCall,
   teamMemberAddCall,
   teamMemberDeleteCall,
@@ -28,11 +30,11 @@ import {
 } from "@tremor/react";
 import { Button, Form, Input, message, Select, Switch, Tooltip } from "antd";
 import { CheckIcon, CopyIcon } from "lucide-react";
-import React, { useEffect, useState } from "react";
+import React, { useEffect, useMemo, useState } from "react";
 import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
 import DeleteResourceModal from "../common_components/DeleteResourceModal";
 import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
-import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
+import { getModelDisplayName, unfurlWildcardModelsInList } from "../key_team_helpers/fetch_available_models_team_key";
 import LoggingSettingsView from "../logging_settings_view";
 import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
 import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions";
@@ -117,6 +119,29 @@ export interface TeamInfoProps {
   premiumUser?: boolean;
 }
 
+const getOrganizationModels = (organization: Organization | null, userModels: string[]) => {
+  let tempModelsToPick = [];
+
+  if (organization) {
+    // Check if organization has "all-proxy-models" in its models array
+    if (organization.models.includes("all-proxy-models")) {
+      // Treat as all-proxy-models (use userModels)
+      tempModelsToPick = userModels;
+    } else if (organization.models.length > 0) {
+      // Organization has specific models
+      tempModelsToPick = organization.models;
+    } else {
+      // Empty array [] is treated as all-proxy-models
+      tempModelsToPick = userModels;
+    }
+  } else {
+    // No organization, show all available models
+    tempModelsToPick = userModels;
+  }
+
+  return unfurlWildcardModelsInList(tempModelsToPick, userModels);
+};
+
 const TeamInfoView: React.FC<TeamInfoProps> = ({
   teamId,
   onClose,
@@ -143,6 +168,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
   const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
   const [isDeleting, setIsDeleting] = useState(false);
   const [isTeamSaving, setIsTeamSaving] = useState(false);
+  const [organization, setOrganization] = useState<Organization | null>(null);
 
   console.log("userModels in team info", userModels);
 
@@ -166,6 +192,31 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
     fetchTeamInfo();
   }, [teamId, accessToken]);
 
+  // Fetch organization data when team has organization_id
+  useEffect(() => {
+    const fetchOrganization = async () => {
+      if (!accessToken || !teamData?.team_info?.organization_id) {
+        setOrganization(null);
+        return;
+      }
+
+      try {
+        const orgData = await organizationInfoCall(accessToken, teamData.team_info.organization_id);
+        setOrganization(orgData);
+      } catch (error) {
+        console.error("Error fetching organization info:", error);
+        setOrganization(null);
+      }
+    };
+
+    fetchOrganization();
+  }, [accessToken, teamData?.team_info?.organization_id]);
+
+  // Compute modelsToPick based on organization and userModels
+  const modelsToPick = useMemo(() => {
+    return getOrganizationModels(organization, userModels);
+  }, [organization, userModels]);
+
   const fetchMcpAccessGroups = async () => {
     if (!accessToken) return;
     if (mcpAccessGroupsLoaded) return;
@@ -596,15 +647,41 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
                     rules={[{ required: true, message: "Please select at least one model" }]}
                   >
                     <Select mode="multiple" placeholder="Select models">
-                      {(is_proxy_admin || userModels.includes("all-proxy-models")) && (
-                        <Select.Option key="all-proxy-models" value="all-proxy-models">
-                          All Proxy Models
-                        </Select.Option>
-                      )}
-                      <Select.Option key="no-default-models" value="no-default-models">
-                        No Default Models
-                      </Select.Option>
-                      {Array.from(new Set(userModels)).map((model, idx) => (
+                      {(() => {
+                        let shouldShowAllProxyModels = false;
+
+                        if (organization) {
+                          // Team is in an organization
+                          if (organization.models.length === 0 || organization.models.includes("all-proxy-models")) {
+                            // Organization has empty array [] or "all-proxy-models"
+                            shouldShowAllProxyModels = true;
+                          }
+                          // Otherwise (organization has specific models), don't show "all-proxy-models"
+                        } else {
+                          // Team is not in an organization
+                          shouldShowAllProxyModels = is_proxy_admin || userModels.includes("all-proxy-models");
+                        }
+
+                        return shouldShowAllProxyModels ? (
+                          <Select.Option key="all-proxy-models" value="all-proxy-models">
+                            All Proxy Models
+                          </Select.Option>
+                        ) : null;
+                      })()}
+                      {(() => {
+                        // Show "no-default-models" option if:
+                        // 1. Team is not in an organization, OR
+                        // 2. Team is in an organization and organization's models include "no-default-models"
+                        const shouldShowNoDefaultModels =
+                          !organization || organization.models.includes("no-default-models");
+
+                        return shouldShowNoDefaultModels ? (
+                          <Select.Option key="no-default-models" value="no-default-models">
+                            No Default Models
+                          </Select.Option>
+                        ) : null;
+                      })()}
+                      {Array.from(new Set(modelsToPick)).map((model, idx) => (
                         <Select.Option key={idx} value={model}>
                           {getModelDisplayName(model)}
                         </Select.Option>
@@ -758,7 +835,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
                   </Form.Item>
 
                   <Form.Item label="Organization ID" name="organization_id">
-                    <Input type="" />
+                    <Input type="" disabled />
                   </Form.Item>
 
                   <Form.Item label="Logging Settings" name="logging_settings">