From c6da45795b26f432b74dc6be8b34bcdf68490972 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 17:12:30 -0700 Subject: [PATCH 01/16] temp commit --- .../OrganizationDropdown.tsx | 52 +++++++++++++++++++ .../organisms/create_key_button.tsx | 43 ++++++++++++++- .../components/templates/key_edit_view.tsx | 46 +++++++++++++++- 3 files changed, 138 insertions(+), 3 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx new file mode 100644 index 00000000000..ac93041f3d7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx @@ -0,0 +1,52 @@ +import React from "react"; +import { Select } from "antd"; +import { Organization } from "../networking"; + +interface OrganizationDropdownProps { + organizations?: Organization[] | null; + value?: string; + onChange?: (value: string) => void; + disabled?: boolean; + loading?: boolean; +} + +const OrganizationDropdown: React.FC = ({ + organizations, + value, + onChange, + disabled, + loading, +}) => { + return ( + + ); +}; + +export default OrganizationDropdown; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 0926b11fe09..71e882db3fb 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1,5 +1,6 @@ "use client"; import { keyKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -23,6 +24,7 @@ import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings" import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion"; import TeamDropdown from "../common_components/team_dropdown"; +import OrganizationDropdown from "../common_components/OrganizationDropdown"; import ProjectDropdown from "../common_components/ProjectDropdown"; import { CreateUserButton } from "../CreateUserButton"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; @@ -160,6 +162,7 @@ export const fetchUserModels = async ( const CreateKey: React.FC = ({ team, teams, data, addKey, autoOpenCreate, prefillData }) => { const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole)); + const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations(); const { data: projects, isLoading: isProjectsLoading } = useProjects(); const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); @@ -179,6 +182,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const [promptsList, setPromptsList] = useState([]); const [loggingSettings, setLoggingSettings] = useState([]); const [selectedCreateKeyTeam, setSelectedCreateKeyTeam] = useState(team); + const [selectedOrganizationId, setSelectedOrganizationId] = useState(null); const [selectedProjectId, setSelectedProjectId] = useState(null); const [isCreateUserModalVisible, setIsCreateUserModalVisible] = useState(false); const [newlyCreatedUserId, setNewlyCreatedUserId] = useState(null); @@ -207,6 +211,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); setSelectedAgentId(null); + setSelectedOrganizationId(null); setSelectedProjectId(null); }; @@ -224,6 +229,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); setSelectedAgentId(null); + setSelectedOrganizationId(null); setSelectedProjectId(null); }; @@ -752,6 +758,32 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp )} + + Organization{" "} + + + + + } + name="organization_id" + className="mt-4" + > + { + setSelectedOrganizationId(orgId || null); + // Clear team and project when org changes + setSelectedCreateKeyTeam(null); + setSelectedProjectId(null); + form.setFieldValue("team_id", undefined); + form.setFieldValue("project_id", undefined); + }} + /> + @@ -773,7 +805,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp help={keyOwner === "service_account" ? "required" : ""} > t.organization_id === selectedOrganizationId) : teams} disabled={selectedProjectId !== null} loading={!teams} onChange={(teamId) => { @@ -781,6 +813,14 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setSelectedCreateKeyTeam(selectedTeam); setSelectedProjectId(null); form.setFieldValue("project_id", undefined); + // Auto-populate org from team for non-admin users + if (selectedTeam?.organization_id) { + setSelectedOrganizationId(selectedTeam.organization_id); + form.setFieldValue("organization_id", selectedTeam.organization_id); + } else if (!teamId) { + setSelectedOrganizationId(null); + form.setFieldValue("organization_id", undefined); + } }} /> @@ -1531,6 +1571,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp excludedFields={[ "key_alias", "team_id", + "organization_id", "models", "duration", "metadata", diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index b6c00577c9b..cf431d10245 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -1,4 +1,5 @@ import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; +import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import PolicySelector from "@/components/policies/PolicySelector"; @@ -13,6 +14,7 @@ import { mapInternalToDisplayNames } from "../callback_info_helpers"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; +import OrganizationDropdown from "../common_components/OrganizationDropdown"; import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; import { KeyResponse } from "../key_team_helpers/key_list"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; @@ -96,10 +98,12 @@ export function KeyEditView({ ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) : [], ); + const [selectedOrganizationId, setSelectedOrganizationId] = useState(keyData.organization_id || null); const [autoRotationEnabled, setAutoRotationEnabled] = useState(keyData.auto_rotate || false); const [rotationInterval, setRotationInterval] = useState(keyData.rotation_interval || ""); const [neverExpire, setNeverExpire] = useState(!keyData.expires); const [isKeySaving, setIsKeySaving] = useState(false); + const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations(); const { data: projects } = useProjects(); const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); @@ -610,6 +614,28 @@ export function KeyEditView({ /> + + Organization{" "} + + + + + } + name="organization_id" + > + { + setSelectedOrganizationId(orgId || null); + form.setFieldValue("team_id", undefined); + }} + /> + + { + const selectedTeam = teams?.find((t) => t.team_id === teamId) || null; + if (selectedTeam?.organization_id) { + setSelectedOrganizationId(selectedTeam.organization_id); + form.setFieldValue("organization_id", selectedTeam.organization_id); + } else if (!teamId) { + setSelectedOrganizationId(null); + form.setFieldValue("organization_id", undefined); + } + }} filterOption={(input, option) => { - const team = teams?.find((t) => t.team_id === option?.value); + const filteredTeams = selectedOrganizationId + ? teams?.filter((t) => t.organization_id === selectedOrganizationId) + : teams; + const team = filteredTeams?.find((t) => t.team_id === option?.value); if (!team) return false; return team.team_alias?.toLowerCase().includes(input.toLowerCase()) ?? false; }} > - {teams?.map((team) => ( + {(selectedOrganizationId + ? teams?.filter((t) => t.organization_id === selectedOrganizationId) + : teams + )?.map((team) => ( {`${team.team_alias} (${team.team_id})`} From c637c93a6ab4b02209523b3f4e6a2e49e8e0a2cb Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 18:08:02 -0700 Subject: [PATCH 02/16] [Fix] Skip all-team-models sentinel in team change validation When moving a key to a different team, `validate_key_team_change` was treating "all-team-models" as a literal model name and checking if the target team could access it. This always failed because "all-team-models" is a UI/backend sentinel meaning "use whatever the team allows." Also reorder checks so the membership check runs after data validation (models, rate limits) but before permission checks, keeping the admin early-return after all validation. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index db1a089ff7b..715a395e67e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2301,23 +2301,16 @@ async def validate_key_team_change( # Check if the team has access to the key's models if len(key.models) > 0: for model in key.models: + # Skip special sentinel values — "all-team-models" means + # "use whatever the team allows", so it's always valid. + if model == SpecialModelNames.all_team_models.value: + continue await can_team_access_model( model=model, team_object=team, llm_router=llm_router, ) - # Check if the key's user_id is a member of the team - member_object = _get_user_in_team( - team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id - ) - if key.user_id is not None: - if not member_object: - raise HTTPException( - status_code=403, - detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.", - ) - # Check if the key's tpm/rpm limit is less than the team's tpm/rpm limit if key.tpm_limit is not None: if team.tpm_limit and key.tpm_limit > team.tpm_limit: @@ -2331,6 +2324,17 @@ async def validate_key_team_change( detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.", ) + # Check if the key's user_id is a member of the team + member_object = _get_user_in_team( + team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id + ) + if key.user_id is not None: + if not member_object: + raise HTTPException( + status_code=403, + detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.", + ) + # Check if the person initiating the change is a Proxy Admin or Team Admin if change_initiated_by.user_role == LitellmUserRoles.PROXY_ADMIN.value: return From bce37e28c3c03b7113ea6b3246052c0d7de3c5d5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 18:15:42 -0700 Subject: [PATCH 03/16] [Test] Add tests for organization dropdown in key create/edit - OrganizationDropdown: renders options, calls onChange on selection, applies disabled state, handles empty list - CreateKey: org dropdown renders, disabled for non-admin users, enabled for admins, form state updates on org selection - KeyEditView: org dropdown renders, disabled for non-admin, enabled for admin, initializes from keyData.organization_id Co-Authored-By: Claude Opus 4.6 --- .../OrganizationDropdown.test.tsx | 69 ++++++++++ .../organisms/create_key_button.test.tsx | 119 +++++++++++++++++- .../templates/key_edit_view.test.tsx | 97 ++++++++++++++ 3 files changed, 284 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx new file mode 100644 index 00000000000..1f6a61f39f8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx @@ -0,0 +1,69 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import OrganizationDropdown from "./OrganizationDropdown"; + +const MOCK_ORGS = [ + { + organization_id: "org-1", + organization_alias: "Engineering", + budget_id: "", + metadata: {}, + models: [], + spend: 0, + model_spend: {}, + created_at: "", + created_by: "", + updated_at: "", + }, + { + organization_id: "org-2", + organization_alias: "Sales", + budget_id: "", + metadata: {}, + models: [], + spend: 0, + model_spend: {}, + created_at: "", + created_by: "", + updated_at: "", + }, +]; + +describe("OrganizationDropdown", () => { + it("should render", () => { + render(); + expect(screen.getByRole("combobox")).toBeInTheDocument(); + }); + + it("should display organization options when opened", async () => { + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("combobox")); + + expect(await screen.findByText("Engineering")).toBeInTheDocument(); + expect(screen.getByText("Sales")).toBeInTheDocument(); + }); + + it("should call onChange with the org id when an organization is selected", async () => { + const onChange = vi.fn(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("Engineering")); + + expect(onChange).toHaveBeenCalledWith("org-1", expect.anything()); + }); + + it("should add ant-select-disabled class when disabled prop is true", () => { + const { container } = render(); + expect(container.querySelector(".ant-select-disabled")).toBeTruthy(); + }); + + it("should render with empty organizations list", () => { + render(); + expect(screen.getByRole("combobox")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index 3ed4c80aea5..eef7292dac1 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -213,7 +213,22 @@ vi.mock("../common_components/PassThroughRoutesSelector", () => ({ default: () = vi.mock("../common_components/PremiumLoggingSettings", () => ({ default: () => null })); vi.mock("../common_components/RateLimitTypeFormItem", () => ({ default: () => null })); vi.mock("../common_components/RouterSettingsAccordion", () => ({ default: () => null })); -vi.mock("../common_components/team_dropdown", () => ({ default: () => null })); +vi.mock("../common_components/team_dropdown", () => ({ + default: ({ teams, onChange, disabled }: { teams?: any[]; onChange?: (v: string) => void; disabled?: boolean }) => ( + + ), +})); vi.mock("../CreateUserButton", () => ({ CreateUserButton: () => null })); vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => null })); vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null })); @@ -227,6 +242,31 @@ vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }), })); +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: vi.fn().mockReturnValue({ + data: [ + { organization_id: "org-1", organization_alias: "Engineering" }, + { organization_id: "org-2", organization_alias: "Sales" }, + ], + isLoading: false, + }), +})); + +vi.mock("../common_components/OrganizationDropdown", () => ({ + default: ({ value, onChange, disabled }: { value?: string; onChange?: (v: string) => void; disabled?: boolean }) => ( + + ), +})); + vi.mock("../common_components/ProjectDropdown", () => ({ default: ({ value, onChange }: { value?: string; onChange?: (v: string) => void }) => ( { expect(setFieldsValueMock).toHaveBeenCalledWith({ key_type: "management" }); }); }); + + describe("organization dropdown", () => { + it("should render the organization dropdown when modal is open", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); + }); + }); + + it("should disable the organization dropdown for non-admin users", async () => { + authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" }; + + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).toBeDisabled(); + }); + }); + + it("should enable the organization dropdown for admin users", async () => { + authorizedState = { ...defaultAuthorizedState, userRole: "Admin" }; + + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).not.toBeDisabled(); + }); + }); + + it("should render team dropdown alongside organization dropdown", async () => { + const teamsWithOrg = [ + { team_id: "team-1", team_alias: "Team Alpha", organization_id: "org-1", models: [] }, + ]; + + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); + expect(screen.getByTestId("team-dropdown")).toBeInTheDocument(); + }); + }); + + it("should set organization_id in form state when org is selected", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); + }); + + act(() => { + fireEvent.change(screen.getByTestId("org-dropdown"), { target: { value: "org-1" } }); + }); + + expect(formStateRef.current["organization_id"]).toBe("org-1"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index b00a8d1e3f8..2e4d0d97e4c 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -53,6 +53,16 @@ vi.mock("../organisms/create_key_button", () => ({ fetchTeamModels: vi.fn().mockResolvedValue(["team-model-1", "team-model-2"]), })); +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: vi.fn().mockReturnValue({ + data: [ + { organization_id: "org-1", organization_alias: "Engineering" }, + { organization_id: "org-2", organization_alias: "Sales" }, + ], + isLoading: false, + }), +})); + vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ useAccessGroups: vi.fn().mockReturnValue({ data: [ @@ -576,4 +586,91 @@ describe("KeyEditView", () => { resolveSubmit(); } }); + + describe("organization dropdown", () => { + it("should render the organization dropdown", async () => { + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Organization")).toBeInTheDocument(); + }); + }); + + it("should disable the organization dropdown for non-admin users", async () => { + const { container } = renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Internal User" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Organization")).toBeInTheDocument(); + }); + + const orgFormItem = screen.getByText("Organization").closest(".ant-form-item"); + const disabledSelect = orgFormItem?.querySelector(".ant-select-disabled"); + expect(disabledSelect).toBeTruthy(); + }); + + it("should not disable the organization dropdown for admin users", async () => { + const { container } = renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Organization")).toBeInTheDocument(); + }); + + const orgFormItem = screen.getByText("Organization").closest(".ant-form-item"); + const disabledSelect = orgFormItem?.querySelector(".ant-select-disabled"); + expect(disabledSelect).toBeFalsy(); + }); + + it("should initialize organization from keyData", async () => { + const keyWithOrg = { + ...MOCK_KEY_DATA, + organization_id: "org-1", + }; + + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Engineering")).toBeInTheDocument(); + }); + }); + }); }); From f6a8087375aaa0eb75569aa1179afaa5ccd29ed4 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 18:17:18 -0700 Subject: [PATCH 04/16] [Test] Add test for all-team-models sentinel skip in team change validation Verifies that validate_key_team_change does not call can_team_access_model for the "all-team-models" sentinel, allowing keys with that value to be moved between teams without model validation failures. Co-Authored-By: Claude Opus 4.6 --- .../test_key_management_endpoints.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index e496cf373ea..cfc16808afb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1502,6 +1502,57 @@ async def test_validate_key_team_change_with_member_permissions(): ) +@pytest.mark.asyncio +async def test_validate_key_team_change_skips_all_team_models_sentinel(): + """ + Test that validate_key_team_change skips the 'all-team-models' sentinel + value when checking if the target team can access the key's models. + + Keys with models=["all-team-models"] mean "use whatever models the team + allows", so moving them to any team should not fail model validation. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + mock_key = MagicMock() + mock_key.user_id = "test-user-123" + mock_key.models = ["all-team-models"] + mock_key.tpm_limit = None + mock_key.rpm_limit = None + + mock_team = MagicMock() + mock_team.team_id = "test-team-456" + mock_team.models = ["gpt-4", "claude-3"] + mock_team.members_with_roles = [] + mock_team.tpm_limit = None + mock_team.rpm_limit = None + + mock_change_initiator = MagicMock() + mock_change_initiator.user_id = "test-user-123" + mock_change_initiator.user_role = LitellmUserRoles.PROXY_ADMIN.value + + mock_router = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model", + new_callable=AsyncMock, + ) as mock_can_access: + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team" + ) as mock_get_user: + mock_get_user.return_value = MagicMock() + + await validate_key_team_change( + key=mock_key, + team=mock_team, + change_initiated_by=mock_change_initiator, + llm_router=mock_router, + ) + + # can_team_access_model should NOT have been called since + # "all-team-models" is a sentinel that should be skipped + mock_can_access.assert_not_called() + + def test_key_rotation_fields_helper(): """ Test the key data update logic for rotation fields. From 4fb71e4a21192eae092dc25213d6323572dd8fe5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 21:38:20 -0700 Subject: [PATCH 05/16] [Fix] Fix tag/list 500 error from invalid Prisma group_by kwargs Use `min`/`max` instead of `_min`/`_max` for Prisma group_by input parameters. The underscore-prefixed names are output keys, not input kwargs. Co-Authored-By: Claude Opus 4.6 --- .../management_endpoints/tag_management_endpoints.py | 11 ++++------- .../test_tag_management_endpoints.py | 6 +++--- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 8e3061c2032..0e60820aab1 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -456,11 +456,8 @@ async def list_tags( dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by( by=["tag"], where={"tag": {"not": None}}, - # The old find_many(distinct=...) returned arbitrary timestamps from - # whichever row Prisma happened to pick. MIN/MAX give more meaningful - # values: earliest appearance and most recent activity. - _min={"created_at": True}, - _max={"updated_at": True}, + min={"created_at": True}, + max={"updated_at": True}, ) dynamic_tag_config = [ @@ -468,8 +465,8 @@ async def list_tags( "name": row["tag"], "description": "This is just a spend tag that was passed dynamically in a request. It does not control any LLM models.", "models": None, - "created_at": row["_min"]["created_at"].isoformat(), - "updated_at": row["_max"]["updated_at"].isoformat(), + "created_at": row["_min"]["created_at"], + "updated_at": row["_max"]["updated_at"], } for row in dynamic_tag_rows if row["tag"] not in stored_tag_names diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index f330b40282c..4b443f211ff 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -282,9 +282,9 @@ async def test_list_tags_with_dynamic_tags(): # Setup dynamic tags via group_by — includes one that overlaps with stored mock_db.litellm_dailytagspend.group_by = AsyncMock(return_value=[ - {"tag": "dynamic-tag-1", "_min": {"created_at": datetime(2025, 2, 1)}, "_max": {"updated_at": datetime(2025, 3, 1)}}, - {"tag": "dynamic-tag-2", "_min": {"created_at": datetime(2025, 2, 2)}, "_max": {"updated_at": datetime(2025, 3, 2)}}, - {"tag": "stored-tag", "_min": {"created_at": datetime(2025, 1, 1)}, "_max": {"updated_at": datetime(2025, 1, 1)}}, # duplicate, should be excluded + {"tag": "dynamic-tag-1", "_min": {"created_at": "2025-02-01T00:00:00Z"}, "_max": {"updated_at": "2025-03-01T00:00:00Z"}}, + {"tag": "dynamic-tag-2", "_min": {"created_at": "2025-02-02T00:00:00Z"}, "_max": {"updated_at": "2025-03-02T00:00:00Z"}}, + {"tag": "stored-tag", "_min": {"created_at": "2025-01-01T00:00:00Z"}, "_max": {"updated_at": "2025-01-01T00:00:00Z"}}, # duplicate, should be excluded ]) headers = {"Authorization": "Bearer sk-1234"} From 1b0c4bdbb7875d965fe51268835722b2a15d4e7a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 21:47:03 -0700 Subject: [PATCH 06/16] Add unit tests for 5 previously untested UI components Tests for HelpLink, ScoreChart, AgentCard, ToolPoliciesView, and CostBreakdownViewer (33 tests total). Co-Authored-By: Claude Opus 4.6 --- .../GuardrailsMonitor/ScoreChart.test.tsx | 54 ++++++++ .../src/components/HelpLink.test.tsx | 117 ++++++++++++++++++ .../src/components/ToolPoliciesView.test.tsx | 53 ++++++++ .../src/components/agents/agent_card.test.tsx | 99 +++++++++++++++ .../view_logs/CostBreakdownViewer.test.tsx | 117 ++++++++++++++++++ 5 files changed, 440 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/GuardrailsMonitor/ScoreChart.test.tsx create mode 100644 ui/litellm-dashboard/src/components/HelpLink.test.tsx create mode 100644 ui/litellm-dashboard/src/components/ToolPoliciesView.test.tsx create mode 100644 ui/litellm-dashboard/src/components/agents/agent_card.test.tsx create mode 100644 ui/litellm-dashboard/src/components/view_logs/CostBreakdownViewer.test.tsx diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/ScoreChart.test.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/ScoreChart.test.tsx new file mode 100644 index 00000000000..848807ce5a0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/ScoreChart.test.tsx @@ -0,0 +1,54 @@ +import React from "react"; +import { describe, it, expect, vi } from "vitest"; +import { screen } from "@testing-library/react"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { ScoreChart } from "./ScoreChart"; + +vi.mock("@tremor/react", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + BarChart: ({ data, categories }: { data: any[]; categories: string[] }) => ( +
+ {data.map((d, i) => ( + + {d.date}: {categories.map((c) => `${c}=${d[c]}`).join(", ")} + + ))} +
+ ), + }; +}); + +describe("ScoreChart", () => { + it("should render the title", () => { + renderWithProviders(); + + expect(screen.getByText("Request Outcomes Over Time")).toBeInTheDocument(); + }); + + it("should show empty state when no data is provided", () => { + renderWithProviders(); + + expect(screen.getByText("No chart data for this period")).toBeInTheDocument(); + }); + + it("should show empty state when data is an empty array", () => { + renderWithProviders(); + + expect(screen.getByText("No chart data for this period")).toBeInTheDocument(); + }); + + it("should render the chart when data is provided", () => { + const data = [ + { date: "2026-03-01", passed: 10, blocked: 2 }, + { date: "2026-03-02", passed: 15, blocked: 1 }, + ]; + + renderWithProviders(); + + expect(screen.queryByText("No chart data for this period")).not.toBeInTheDocument(); + expect(screen.getByText(/2026-03-01/)).toBeInTheDocument(); + expect(screen.getByText(/2026-03-02/)).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/HelpLink.test.tsx b/ui/litellm-dashboard/src/components/HelpLink.test.tsx new file mode 100644 index 00000000000..72033bd93a1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/HelpLink.test.tsx @@ -0,0 +1,117 @@ +import React from "react"; +import { describe, it, expect, vi } from "vitest"; +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../tests/test-utils"; +import { HelpLink, HelpIcon, DocsMenu } from "./HelpLink"; + +describe("HelpLink", () => { + it("should render with default children and open in new tab", () => { + renderWithProviders(); + + const link = screen.getByRole("link", { name: /learn more/i }); + expect(link).toHaveAttribute("href", "https://docs.example.com"); + expect(link).toHaveAttribute("target", "_blank"); + expect(link).toHaveAttribute("rel", "noopener noreferrer"); + }); + + it("should render custom children text", () => { + renderWithProviders( + Custom docs link + ); + + expect(screen.getByText("Custom docs link")).toBeInTheDocument(); + }); + + it("should include a screen-reader-only label for accessibility", () => { + renderWithProviders(); + + expect(screen.getByText("(opens in a new tab)")).toBeInTheDocument(); + }); +}); + +describe("HelpIcon", () => { + it("should render a help button with accessible label", () => { + renderWithProviders(); + + expect(screen.getByRole("button", { name: /help information/i })).toBeInTheDocument(); + }); + + it("should show tooltip content on hover", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.hover(screen.getByRole("button", { name: /help information/i })); + + expect(screen.getByText("Tooltip help text")).toBeInTheDocument(); + }); + + it("should show learn more link when learnMoreHref is provided", async () => { + const user = userEvent.setup(); + renderWithProviders( + + ); + + await user.hover(screen.getByRole("button", { name: /help information/i })); + + const link = screen.getByRole("link", { name: /read docs/i }); + expect(link).toHaveAttribute("href", "https://docs.example.com"); + }); + + it("should not show learn more link when learnMoreHref is not provided", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.hover(screen.getByRole("button", { name: /help information/i })); + + expect(screen.queryByRole("link")).not.toBeInTheDocument(); + }); +}); + +describe("DocsMenu", () => { + const items = [ + { label: "Custom pricing", href: "https://docs.example.com/pricing" }, + { label: "Cost tracking", href: "https://docs.example.com/cost" }, + ]; + + it("should render the menu button with default text", () => { + renderWithProviders(); + + expect(screen.getByRole("button", { name: /docs/i })).toBeInTheDocument(); + }); + + it("should show menu items when button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /docs/i })); + + expect(screen.getByText("Custom pricing")).toBeInTheDocument(); + expect(screen.getByText("Cost tracking")).toBeInTheDocument(); + }); + + it("should close the menu when an item is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /docs/i })); + await user.click(screen.getByText("Custom pricing")); + + expect(screen.queryByText("Cost tracking")).not.toBeInTheDocument(); + }); + + it("should set aria-expanded correctly based on menu state", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const button = screen.getByRole("button", { name: /docs/i }); + expect(button).toHaveAttribute("aria-expanded", "false"); + + await user.click(button); + expect(button).toHaveAttribute("aria-expanded", "true"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ToolPoliciesView.test.tsx b/ui/litellm-dashboard/src/components/ToolPoliciesView.test.tsx new file mode 100644 index 00000000000..8b2b1d0e4b7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ToolPoliciesView.test.tsx @@ -0,0 +1,53 @@ +import React from "react"; +import { describe, it, expect, vi } from "vitest"; +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../tests/test-utils"; +import ToolPoliciesView from "./ToolPoliciesView"; + +vi.mock("@/components/ToolDetail", () => ({ + ToolDetail: ({ toolName, onBack }: { toolName: string; onBack: () => void }) => ( +
+ Detail: {toolName} + +
+ ), +})); + +vi.mock("@/components/ToolPolicies", () => ({ + ToolPolicies: ({ onSelectTool }: { onSelectTool: (name: string) => void }) => ( +
+ Tool Policies Overview + +
+ ), +})); + +describe("ToolPoliciesView", () => { + it("should render the overview by default", () => { + renderWithProviders(); + + expect(screen.getByText("Tool Policies Overview")).toBeInTheDocument(); + }); + + it("should navigate to tool detail when a tool is selected", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /select tool/i })); + + expect(screen.getByText("Detail: my-tool")).toBeInTheDocument(); + expect(screen.queryByText("Tool Policies Overview")).not.toBeInTheDocument(); + }); + + it("should navigate back to overview when back is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /select tool/i })); + await user.click(screen.getByRole("button", { name: /back/i })); + + expect(screen.getByText("Tool Policies Overview")).toBeInTheDocument(); + expect(screen.queryByText("Detail: my-tool")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/agents/agent_card.test.tsx b/ui/litellm-dashboard/src/components/agents/agent_card.test.tsx new file mode 100644 index 00000000000..0f928866a0d --- /dev/null +++ b/ui/litellm-dashboard/src/components/agents/agent_card.test.tsx @@ -0,0 +1,99 @@ +import React from "react"; +import { describe, it, expect, vi } from "vitest"; +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../tests/test-utils"; +import AgentCard from "./agent_card"; +import type { Agent } from "./types"; + +const baseAgent: Agent = { + agent_id: "agent-123", + agent_name: "Test Agent", + litellm_params: { model: "gpt-4" }, + agent_card_params: { + description: "A test agent for unit testing", + url: "https://agent.example.com", + }, +}; + +const defaultProps = { + agent: baseAgent, + onAgentClick: vi.fn(), + accessToken: "token-123", + isAdmin: false, + onAgentUpdated: vi.fn(), +}; + +describe("AgentCard", () => { + it("should render the agent name and description", () => { + renderWithProviders(); + + expect(screen.getByText("Test Agent")).toBeInTheDocument(); + expect(screen.getByText("A test agent for unit testing")).toBeInTheDocument(); + }); + + it("should show 'No description' when agent has no description", () => { + const agent = { ...baseAgent, agent_card_params: {} }; + renderWithProviders(); + + expect(screen.getByText("No description")).toBeInTheDocument(); + }); + + it("should show the agent URL when provided", () => { + renderWithProviders(); + + expect(screen.getByText("https://agent.example.com")).toBeInTheDocument(); + }); + + it("should show 'Needs Setup' badge when agent has no key", () => { + renderWithProviders(); + + expect(screen.getByText("Needs Setup")).toBeInTheDocument(); + expect(screen.getByText("No key assigned")).toBeInTheDocument(); + }); + + it("should show 'Active' badge and key info when agent has a key", () => { + const keyInfo = { has_key: true, key_alias: "my-key" }; + renderWithProviders(); + + expect(screen.getByText("Active")).toBeInTheDocument(); + expect(screen.getByText("my-key")).toBeInTheDocument(); + }); + + it("should call onAgentClick when card is clicked", async () => { + const user = userEvent.setup(); + const onAgentClick = vi.fn(); + renderWithProviders(); + + await user.click(screen.getByText("Test Agent")); + + expect(onAgentClick).toHaveBeenCalledWith("agent-123"); + }); + + it("should show delete button only for admins", () => { + const onDeleteClick = vi.fn(); + const { unmount } = renderWithProviders( + + ); + expect(screen.queryByRole("button", { name: /delete/i })).not.toBeInTheDocument(); + + unmount(); + + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /delete/i })).toBeInTheDocument(); + }); + + it("should call onDeleteClick with agent id and name when delete is clicked", async () => { + const user = userEvent.setup(); + const onDeleteClick = vi.fn(); + renderWithProviders( + + ); + + await user.click(screen.getByRole("button", { name: /delete/i })); + + expect(onDeleteClick).toHaveBeenCalledWith("agent-123", "Test Agent"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/CostBreakdownViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/CostBreakdownViewer.test.tsx new file mode 100644 index 00000000000..4a5cff3757a --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/CostBreakdownViewer.test.tsx @@ -0,0 +1,117 @@ +import React from "react"; +import { describe, it, expect } from "vitest"; +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { CostBreakdownViewer, CostBreakdown } from "./CostBreakdownViewer"; + +describe("CostBreakdownViewer", () => { + it("should render nothing when there is no meaningful data", () => { + const { container } = renderWithProviders( + + ); + + expect(container.firstChild).toBeNull(); + }); + + it("should render nothing when costBreakdown is undefined", () => { + const { container } = renderWithProviders( + + ); + + expect(container.firstChild).toBeNull(); + }); + + it("should render the collapse header with heading and total", () => { + const breakdown: CostBreakdown = { + input_cost: 0.001, + output_cost: 0.002, + total_cost: 0.003, + }; + + renderWithProviders( + + ); + + expect(screen.getByRole("heading", { name: "Cost Breakdown" })).toBeInTheDocument(); + }); + + it("should show input and output costs when the panel is expanded", async () => { + const user = userEvent.setup(); + const breakdown: CostBreakdown = { + input_cost: 0.001, + output_cost: 0.002, + }; + + renderWithProviders( + + ); + + await user.click(screen.getByRole("heading", { name: "Cost Breakdown" })); + + expect(screen.getByText("Input Cost:")).toBeVisible(); + expect(screen.getByText("Output Cost:")).toBeVisible(); + expect(screen.getByText(/500 prompt tokens/)).toBeVisible(); + expect(screen.getByText(/200 completion tokens/)).toBeVisible(); + }); + + it("should show '(Cached)' in the header when cacheHit is true", () => { + const breakdown: CostBreakdown = { + input_cost: 0.001, + output_cost: 0.002, + total_cost: 0.003, + }; + + renderWithProviders( + + ); + + expect(screen.getByText(/\(Cached\)/)).toBeInTheDocument(); + }); + + it("should show discount label with percentage when panel is expanded", async () => { + const user = userEvent.setup(); + const breakdown: CostBreakdown = { + input_cost: 0.01, + output_cost: 0.02, + discount_percent: 0.1, + discount_amount: 0.003, + }; + + renderWithProviders( + + ); + + await user.click(screen.getByRole("heading", { name: "Cost Breakdown" })); + + expect(screen.getByText(/Discount \(10\.00%\)/)).toBeVisible(); + }); + + it("should show margin label with percentage when panel is expanded", async () => { + const user = userEvent.setup(); + const breakdown: CostBreakdown = { + input_cost: 0.01, + output_cost: 0.02, + margin_percent: 0.15, + margin_total_amount: 0.005, + }; + + renderWithProviders( + + ); + + await user.click(screen.getByRole("heading", { name: "Cost Breakdown" })); + + expect(screen.getByText(/Margin \(15\.00%\)/)).toBeVisible(); + expect(screen.getByText("Final Calculated Cost:")).toBeVisible(); + }); +}); From f0d283cf9f4b59ad732c6b8cfcba252e1c8c9c59 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Mar 2026 00:01:25 -0700 Subject: [PATCH 07/16] [Feature] UI - Default Team Settings: Modernize page and fix defaults application Rewrite Default Team Settings UI from Tremor to antd with hardcoded fields, fix default team params not applying during team creation or persisting across proxy restarts, remove dead code, and add comprehensive tests. Co-Authored-By: Claude Opus 4.6 --- litellm/constants.py | 1 + .../management_endpoints/team_endpoints.py | 37 +- .../team_member_permission_checks.py | 20 +- .../proxy_setting_endpoints.py | 41 +- .../proxy/management_endpoints/ui_sso.py | 6 +- .../scim/test_scim_v2_endpoints.py | 1 - .../test_team_default_params.py | 487 ++++++++++++++ .../test_team_member_permission_checks.py | 190 ++++++ .../test_proxy_setting_endpoints.py | 50 ++ .../src/components/TeamSSOSettings.test.tsx | 607 ++++++------------ .../src/components/TeamSSOSettings.tsx | 443 +++++++------ .../src/components/networking.tsx | 1 - ui/litellm-dashboard/tsconfig.json | 2 +- 13 files changed, 1266 insertions(+), 620 deletions(-) create mode 100644 tests/test_litellm/proxy/management_endpoints/test_team_default_params.py create mode 100644 tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py diff --git a/litellm/constants.py b/litellm/constants.py index dbc79b69a67..0cf9700cddd 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1417,6 +1417,7 @@ SECRET_MANAGER_REFRESH_INTERVAL = int( ) LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "default_internal_user_params", + "default_team_params", "public_mcp_servers", "public_agent_groups", "public_model_groups", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 9c8e6f7282b..d83ceb1b095 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -252,6 +252,28 @@ class TeamMemberBudgetHandler: data_dict.pop("team_member_tpm_limit", None) +def _get_default_team_param(field: str) -> Any: + """ + Returns a default value for the given field from litellm.default_team_params config. + Returns None if no default is configured. + + For list fields containing enums (e.g. team_member_permissions), converts enum values to strings. + """ + default_params = litellm.default_team_params + if default_params is None: + return None + if isinstance(default_params, dict): + value = default_params.get(field) + else: + value = getattr(default_params, field, None) + if value is None: + return None + # Convert enum values in lists to strings + if isinstance(value, list): + return [v.value if hasattr(v, "value") else v for v in value] + return value + + def _is_available_team(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: if litellm.default_internal_user_params is None: return False @@ -833,16 +855,23 @@ async def new_team( # noqa: PLR0915 prisma_client=prisma_client, ) - # If max_budget is not explicitly provided in the request, - # check for a default value in the proxy configuration. + # Apply defaults from litellm.default_team_params for any fields + # not explicitly provided in the request. + for field in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit", "team_member_permissions"): + if getattr(data, field, None) is None: + default_value = _get_default_team_param(field) + if default_value is not None: + setattr(data, field, default_value) + + # Legacy fallback: apply max_budget from default_team_settings (YAML config) + # if still not set after checking default_team_params. if data.max_budget is None: if ( isinstance(litellm.default_team_settings, list) and len(litellm.default_team_settings) > 0 and isinstance(litellm.default_team_settings[0], dict) ): - default_settings = litellm.default_team_settings[0] - default_budget = default_settings.get("max_budget") + default_budget = litellm.default_team_settings[0].get("max_budget") if default_budget is not None: data.max_budget = default_budget diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index 4642028b77c..7dd99d4ff18 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -16,11 +16,13 @@ from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import PrismaClient -DEFAULT_TEAM_MEMBER_PERMISSIONS = [ +BASELINE_TEAM_MEMBER_PERMISSIONS = [ KeyManagementRoutes.KEY_INFO, KeyManagementRoutes.KEY_HEALTH, ] +DEFAULT_TEAM_MEMBER_PERMISSIONS = BASELINE_TEAM_MEMBER_PERMISSIONS + class TeamMemberPermissionChecks: @staticmethod @@ -29,15 +31,23 @@ class TeamMemberPermissionChecks: team_table: LiteLLM_TeamTableCachedObj, ) -> List[KeyManagementRoutes]: """ - Returns the permissions for a team member + Returns the permissions for a team member. + + - If team has explicit permissions set (including []), use those + plus baseline permissions (/key/info, /key/health). + - If team has no permissions set (None), fall back to + DEFAULT_TEAM_MEMBER_PERMISSIONS. """ - if team_table.team_member_permissions and isinstance( + if team_table.team_member_permissions is not None and isinstance( team_table.team_member_permissions, list ): - return [ + permissions = { KeyManagementRoutes(permission) for permission in team_table.team_member_permissions - ] + } + # Always include baseline permissions + permissions.update(BASELINE_TEAM_MEMBER_PERMISSIONS) + return list(permissions) return DEFAULT_TEAM_MEMBER_PERMISSIONS diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 8df215d9980..0fa27905bab 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -328,11 +328,34 @@ async def _get_settings_with_schema( } # Add property descriptions + defs = schema.get("$defs", schema.get("definitions", {})) for field_name, field_info in schema["properties"].items(): - result["field_schema"]["properties"][field_name] = { + # For Optional fields, Pydantic v2 uses anyOf with [actual_type, null]. + # Resolve the non-null variant to get the real type and items. + resolved = field_info + if "anyOf" in field_info: + for variant in field_info["anyOf"]: + if variant.get("type") != "null": + resolved = variant + break + + prop_entry: dict = { "description": field_info.get("description", ""), - "type": field_info.get("type", "string"), + "type": resolved.get("type", "string"), } + # Pass through items info (including enum values) for array fields + # so the UI can render a multi-select dropdown + if "items" in resolved: + items = resolved["items"] + # Resolve $ref to enum definitions if needed + if "$ref" in items: + ref_name = items["$ref"].split("/")[-1] + ref_def = defs.get(ref_name, {}) + if "enum" in ref_def: + prop_entry["items"] = {"enum": ref_def["enum"]} + else: + prop_entry["items"] = items + result["field_schema"]["properties"][field_name] = prop_entry # Add nested object descriptions for def_name, def_schema in schema.get("definitions", {}).items(): @@ -427,7 +450,6 @@ async def _update_litellm_setting( DefaultInternalUserParams, DefaultTeamSSOParams, MCPSemanticFilterSettings ], settings_key: str, - in_memory_var: Any, success_message: str, ): """ @@ -436,7 +458,6 @@ async def _update_litellm_setting( Args: settings: The settings object to update settings_key: The key in litellm_settings to update - in_memory_var: The in-memory variable to update success_message: Message to return on success """ from litellm.proxy.proxy_server import proxy_config, store_model_in_db @@ -449,13 +470,16 @@ async def _update_litellm_setting( }, ) - # Update the in-memory settings in_memory_var = settings.model_dump(exclude_none=True) - setattr(litellm, settings_key, in_memory_var) - # Load existing config + # Load existing config first, then set in-memory value after, + # because get_config() may overwrite litellm. with stale DB values + # via LITELLM_SETTINGS_SAFE_DB_OVERRIDES. config = await proxy_config.get_config() + # Update the in-memory settings (after get_config to avoid stale override) + setattr(litellm, settings_key, in_memory_var) + # Update config with new settings if "litellm_settings" not in config: config["litellm_settings"] = {} @@ -495,7 +519,6 @@ async def update_internal_user_settings( return await _update_litellm_setting( settings=settings, settings_key="default_internal_user_params", - in_memory_var=litellm.default_internal_user_params, success_message="Internal user settings updated successfully", ) @@ -513,7 +536,6 @@ async def update_default_team_settings(settings: DefaultTeamSSOParams): return await _update_litellm_setting( settings=settings, settings_key="default_team_params", - in_memory_var=litellm.default_team_params, success_message="Default team settings updated successfully", ) @@ -935,7 +957,6 @@ async def update_mcp_semantic_filter_settings( result = await _update_litellm_setting( settings=settings, settings_key="mcp_semantic_tool_filter", - in_memory_var=None, success_message="MCP Semantic Filter settings updated successfully. Changes will be applied across all pods within 10 seconds.", ) try: diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 6743c4a5b9b..7d8ff0f65c1 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -3,7 +3,7 @@ from typing import Dict, List, Literal, Optional, Union from pydantic import BaseModel, Field from typing_extensions import TypedDict -from litellm.proxy._types import LitellmUserRoles +from litellm.proxy._types import KeyManagementRoutes, LitellmUserRoles from litellm.types.utils import LiteLLMPydanticObjectBase @@ -205,6 +205,10 @@ class DefaultTeamSSOParams(LiteLLMPydanticObjectBase): default=None, description="Default rpm limit for new automatically created teams", ) + team_member_permissions: Optional[List[KeyManagementRoutes]] = Field( + default=None, + description="Default permissions granted to members of newly created teams (e.g. /key/generate, /key/update, /key/delete). /key/info and /key/health are always included.", + ) class InProductNudgeResponse(BaseModel): diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 1355ca0abbe..3c4444a5efc 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -202,7 +202,6 @@ async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeyp await _update_litellm_setting( settings=settings, settings_key="default_internal_user_params", - in_memory_var=litellm.default_internal_user_params, success_message="ok", ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py b/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py new file mode 100644 index 00000000000..7fc7cb8aae2 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py @@ -0,0 +1,487 @@ +""" +Tests for applying default team params during team creation +and loading default_team_params from DB on startup. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../") +) # Adds the parent directory to the system path + +import litellm +from litellm.proxy._types import ( + NewTeamRequest, + UserAPIKeyAuth, + LitellmUserRoles, +) +from litellm.proxy.management_endpoints.team_endpoints import ( + _get_default_team_param, +) +from litellm.proxy.proxy_server import ProxyConfig + + +# --------------------------------------------------------------------------- +# _update_config_fields: default_team_params loaded from DB on startup +# --------------------------------------------------------------------------- + + +class TestConfigFieldsDefaultTeamParams: + """Tests that _update_config_fields applies default_team_params from DB.""" + + def _make_proxy_config(self) -> ProxyConfig: + return ProxyConfig() + + def test_default_team_params_applied_from_db(self, monkeypatch): + """default_team_params in DB is set on litellm module during config load.""" + monkeypatch.setattr(litellm, "default_team_params", None) + + pc = self._make_proxy_config() + db_settings = { + "default_team_params": { + "max_budget": 500.0, + "budget_duration": "30d", + "tpm_limit": 1000, + "rpm_limit": 200, + "team_member_permissions": ["/key/generate", "/key/delete"], + } + } + + pc._update_config_fields( + current_config={}, + param_name="litellm_settings", + db_param_value=db_settings, + ) + + assert litellm.default_team_params == db_settings["default_team_params"] + + def test_default_team_params_merged_into_config_dict(self): + """DB default_team_params ends up in the returned config dict.""" + pc = self._make_proxy_config() + config = {"litellm_settings": {"cache": False}} + db_settings = { + "default_team_params": { + "max_budget": 100.0, + } + } + + result = pc._update_config_fields( + current_config=config, + param_name="litellm_settings", + db_param_value=db_settings, + ) + + assert result["litellm_settings"]["default_team_params"] == {"max_budget": 100.0} + # Existing keys preserved + assert result["litellm_settings"]["cache"] is False + + def test_default_team_params_not_applied_when_absent(self, monkeypatch): + """When DB litellm_settings has no default_team_params, it stays None.""" + monkeypatch.setattr(litellm, "default_team_params", None) + + pc = self._make_proxy_config() + pc._update_config_fields( + current_config={}, + param_name="litellm_settings", + db_param_value={"cache": True}, + ) + + assert litellm.default_team_params is None + + def test_default_team_params_overrides_yaml_value(self, monkeypatch): + """DB value for default_team_params overrides YAML value via deep merge.""" + monkeypatch.setattr(litellm, "default_team_params", None) + + pc = self._make_proxy_config() + config = { + "litellm_settings": { + "default_team_params": { + "max_budget": 50.0, + "tpm_limit": 100, + } + } + } + db_settings = { + "default_team_params": { + "max_budget": 200.0, + "rpm_limit": 500, + } + } + + result = pc._update_config_fields( + current_config=config, + param_name="litellm_settings", + db_param_value=db_settings, + ) + + merged = result["litellm_settings"]["default_team_params"] + # DB value wins for max_budget + assert merged["max_budget"] == 200.0 + # DB adds rpm_limit + assert merged["rpm_limit"] == 500 + # YAML tpm_limit preserved (not in DB) + assert merged["tpm_limit"] == 100 + + # setattr should have applied the DB value + assert litellm.default_team_params == db_settings["default_team_params"] + + +# --------------------------------------------------------------------------- +# new_team: default params applied to team creation +# +# We test the defaults-application logic by calling new_team with +# prisma_client patched at the proxy_server module level (where the +# endpoint imports it from). +# --------------------------------------------------------------------------- + + +class TestNewTeamDefaultParamsApplied: + """Tests that /team/new applies defaults from litellm.default_team_params.""" + + @pytest.fixture(autouse=True) + def setup_mocks(self, monkeypatch): + """Set up common mocks for team creation tests.""" + mock_prisma = AsyncMock() + mock_prisma.insert_data = AsyncMock( + return_value=MagicMock( + team_id="test-team-id", + team_alias="test-team", + ) + ) + mock_prisma.get_generic_data = AsyncMock(return_value=None) + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_teamtable = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ) + + # Reset default_team_settings to avoid legacy fallback interference + monkeypatch.setattr(litellm, "default_team_settings", None) + + def _make_admin_auth(self) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + @pytest.mark.asyncio + async def test_all_defaults_applied_when_not_provided(self, monkeypatch): + """When no budget/rate/permission fields are in the request, all defaults apply.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr( + litellm, + "default_team_params", + { + "max_budget": 100.0, + "budget_duration": "30d", + "tpm_limit": 200, + "rpm_limit": 500, + "team_member_permissions": ["/key/generate", "/key/update"], + }, + ) + + data = NewTeamRequest(team_alias="my-team") + auth = self._make_admin_auth() + + try: + await new_team( + data=data, + user_api_key_dict=auth, + http_request=MagicMock(), + ) + except Exception: + pass # May fail on downstream mocks, that's OK + + # Verify defaults were set on the data object + assert data.max_budget == 100.0 + assert data.budget_duration == "30d" + assert data.tpm_limit == 200 + assert data.rpm_limit == 500 + assert data.team_member_permissions == ["/key/generate", "/key/update"] + + @pytest.mark.asyncio + async def test_explicit_values_not_overridden(self, monkeypatch): + """When request provides explicit values, defaults do not override them.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr( + litellm, + "default_team_params", + { + "max_budget": 100.0, + "budget_duration": "30d", + "tpm_limit": 200, + "rpm_limit": 500, + "team_member_permissions": ["/key/generate"], + }, + ) + + data = NewTeamRequest( + team_alias="my-team", + max_budget=50.0, + budget_duration="7d", + tpm_limit=999, + rpm_limit=888, + team_member_permissions=["/key/delete"], + ) + auth = self._make_admin_auth() + + try: + await new_team( + data=data, + user_api_key_dict=auth, + http_request=MagicMock(), + ) + except Exception: + pass + + # Explicit values preserved + assert data.max_budget == 50.0 + assert data.budget_duration == "7d" + assert data.tpm_limit == 999 + assert data.rpm_limit == 888 + assert data.team_member_permissions == ["/key/delete"] + + @pytest.mark.asyncio + async def test_partial_defaults_applied(self, monkeypatch): + """Only missing fields get defaults; provided fields are untouched.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr( + litellm, + "default_team_params", + { + "max_budget": 100.0, + "budget_duration": "30d", + "tpm_limit": 200, + "rpm_limit": 500, + }, + ) + + data = NewTeamRequest( + team_alias="my-team", + max_budget=75.0, # explicit + # budget_duration, tpm_limit, rpm_limit not set → defaults apply + ) + auth = self._make_admin_auth() + + try: + await new_team( + data=data, + user_api_key_dict=auth, + http_request=MagicMock(), + ) + except Exception: + pass + + assert data.max_budget == 75.0 # explicit, not overridden + assert data.budget_duration == "30d" # default applied + assert data.tpm_limit == 200 # default applied + assert data.rpm_limit == 500 # default applied + + @pytest.mark.asyncio + async def test_no_defaults_when_config_is_none(self, monkeypatch): + """When default_team_params is None, no defaults applied.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_params", None) + + data = NewTeamRequest(team_alias="my-team") + auth = self._make_admin_auth() + + try: + await new_team( + data=data, + user_api_key_dict=auth, + http_request=MagicMock(), + ) + except Exception: + pass + + assert data.max_budget is None + assert data.budget_duration is None + assert data.tpm_limit is None + assert data.rpm_limit is None + assert data.team_member_permissions is None + + @pytest.mark.asyncio + async def test_legacy_default_team_settings_fallback(self, monkeypatch): + """Legacy default_team_settings YAML config applies max_budget as fallback.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_params", None) + monkeypatch.setattr( + litellm, + "default_team_settings", + [{"team_id": "default", "max_budget": 999.0}], + ) + + data = NewTeamRequest(team_alias="my-team") + auth = self._make_admin_auth() + + try: + await new_team( + data=data, + user_api_key_dict=auth, + http_request=MagicMock(), + ) + except Exception: + pass + + assert data.max_budget == 999.0 + + @pytest.mark.asyncio + async def test_default_team_params_takes_priority_over_legacy(self, monkeypatch): + """default_team_params max_budget takes priority over legacy default_team_settings.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr( + litellm, + "default_team_params", + {"max_budget": 100.0}, + ) + monkeypatch.setattr( + litellm, + "default_team_settings", + [{"team_id": "default", "max_budget": 999.0}], + ) + + data = NewTeamRequest(team_alias="my-team") + auth = self._make_admin_auth() + + try: + await new_team( + data=data, + user_api_key_dict=auth, + http_request=MagicMock(), + ) + except Exception: + pass + + # default_team_params wins (100.0), legacy fallback (999.0) not used + assert data.max_budget == 100.0 + + +# --------------------------------------------------------------------------- +# _update_litellm_setting: setattr ordering +# --------------------------------------------------------------------------- + + +class TestUpdateLitellmSettingOrdering: + """Tests that _update_litellm_setting sets in-memory value AFTER get_config, + so stale DB values from LITELLM_SETTINGS_SAFE_DB_OVERRIDES don't overwrite it.""" + + @pytest.mark.asyncio + async def test_setattr_not_overwritten_by_get_config(self, monkeypatch): + """The new in-memory value survives get_config() which may load stale DB values.""" + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + _update_litellm_setting, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + DefaultTeamSSOParams, + ) + + # Simulate stale DB state: get_config returns old default_team_params + stale_value = {"max_budget": 50.0} + monkeypatch.setattr(litellm, "default_team_params", stale_value) + + # get_config will overwrite litellm.default_team_params with stale DB value + async def mock_get_config(): + # Simulate what _update_config_from_db does for safe overrides + litellm.default_team_params = stale_value + return { + "litellm_settings": { + "default_team_params": stale_value, + } + } + + saved_configs = [] + + async def mock_save_config(new_config=None): + saved_configs.append(new_config) + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr(proxy_config, "save_config", mock_save_config) + monkeypatch.setattr( + "litellm.proxy.proxy_server.store_model_in_db", True + ) + + # New settings to save + new_settings = DefaultTeamSSOParams( + max_budget=200.0, + budget_duration="7d", + rpm_limit=1000, + ) + + result = await _update_litellm_setting( + settings=new_settings, + settings_key="default_team_params", + success_message="Updated", + ) + + # In-memory value should be the NEW value, not the stale one + expected = new_settings.model_dump(exclude_none=True) + assert litellm.default_team_params == expected + + # Saved config should contain the new value + assert len(saved_configs) == 1 + saved_settings = saved_configs[0]["litellm_settings"]["default_team_params"] + assert saved_settings == expected + + # Return value should reflect the new settings + assert result["settings"] == expected + + @pytest.mark.asyncio + async def test_requires_store_model_in_db(self, monkeypatch): + """Raises HTTPException when store_model_in_db is not True.""" + from fastapi import HTTPException + + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + _update_litellm_setting, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + DefaultTeamSSOParams, + ) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.store_model_in_db", False + ) + + with pytest.raises(HTTPException) as exc_info: + await _update_litellm_setting( + settings=DefaultTeamSSOParams(max_budget=100.0), + settings_key="default_team_params", + success_message="Updated", + ) + + assert exc_info.value.status_code == 500 + + +# --------------------------------------------------------------------------- +# LITELLM_SETTINGS_SAFE_DB_OVERRIDES contains default_team_params +# --------------------------------------------------------------------------- + + +class TestSafeDbOverrides: + """Verify default_team_params is in the safe overrides list.""" + + def test_default_team_params_in_safe_overrides(self): + from litellm.constants import LITELLM_SETTINGS_SAFE_DB_OVERRIDES + + assert "default_team_params" in LITELLM_SETTINGS_SAFE_DB_OVERRIDES + + def test_default_internal_user_params_in_safe_overrides(self): + """Sanity: default_internal_user_params was already in the list.""" + from litellm.constants import LITELLM_SETTINGS_SAFE_DB_OVERRIDES + + assert "default_internal_user_params" in LITELLM_SETTINGS_SAFE_DB_OVERRIDES diff --git a/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py b/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py new file mode 100644 index 00000000000..6aa08dddd08 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_team_member_permission_checks.py @@ -0,0 +1,190 @@ +import os +import sys +from unittest.mock import MagicMock + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.proxy._types import KeyManagementRoutes, Member +from litellm.proxy.management_helpers.team_member_permission_checks import ( + BASELINE_TEAM_MEMBER_PERMISSIONS, + TeamMemberPermissionChecks, +) + + +def _make_team_table(team_member_permissions): + """Create a mock team table object with given permissions.""" + team = MagicMock() + team.team_member_permissions = team_member_permissions + return team + + +class TestGetPermissionsForTeamMember: + def test_none_permissions_returns_defaults(self): + """When team_member_permissions is None, return DEFAULT_TEAM_MEMBER_PERMISSIONS.""" + team = _make_team_table(None) + member = MagicMock(spec=Member) + + result = TeamMemberPermissionChecks.get_permissions_for_team_member( + team_member_object=member, team_table=team + ) + + assert set(result) == set(BASELINE_TEAM_MEMBER_PERMISSIONS) + + def test_empty_list_includes_baseline(self): + """When team_member_permissions is [], baseline permissions are still included.""" + team = _make_team_table([]) + member = MagicMock(spec=Member) + + result = TeamMemberPermissionChecks.get_permissions_for_team_member( + team_member_object=member, team_table=team + ) + + assert KeyManagementRoutes.KEY_INFO in result + assert KeyManagementRoutes.KEY_HEALTH in result + + def test_explicit_permissions_include_baseline(self): + """When explicit permissions are set, baseline is always included.""" + team = _make_team_table(["/key/generate", "/key/delete"]) + member = MagicMock(spec=Member) + + result = TeamMemberPermissionChecks.get_permissions_for_team_member( + team_member_object=member, team_table=team + ) + + assert KeyManagementRoutes.KEY_GENERATE in result + assert KeyManagementRoutes.KEY_DELETE in result + assert KeyManagementRoutes.KEY_INFO in result + assert KeyManagementRoutes.KEY_HEALTH in result + + def test_explicit_permissions_with_baseline_no_duplicates(self): + """When explicit permissions already include baseline, no duplicates.""" + team = _make_team_table(["/key/info", "/key/generate"]) + member = MagicMock(spec=Member) + + result = TeamMemberPermissionChecks.get_permissions_for_team_member( + team_member_object=member, team_table=team + ) + + # Using set ensures no duplicates from the implementation + assert KeyManagementRoutes.KEY_INFO in result + assert KeyManagementRoutes.KEY_GENERATE in result + assert KeyManagementRoutes.KEY_HEALTH in result + + +class TestGetDefaultTeamParam: + def test_returns_none_when_no_config(self, monkeypatch): + """Returns None when litellm.default_team_params is None.""" + import litellm + + from litellm.proxy.management_endpoints.team_endpoints import ( + _get_default_team_param, + ) + + monkeypatch.setattr(litellm, "default_team_params", None) + + assert _get_default_team_param("team_member_permissions") is None + assert _get_default_team_param("max_budget") is None + + def test_returns_none_when_field_not_set(self, monkeypatch): + """Returns None when default_team_params exists but the field is not set.""" + import litellm + + from litellm.proxy.management_endpoints.team_endpoints import ( + _get_default_team_param, + ) + + monkeypatch.setattr(litellm, "default_team_params", {"models": ["gpt-4"]}) + + assert _get_default_team_param("team_member_permissions") is None + assert _get_default_team_param("max_budget") is None + + def test_returns_permissions_from_dict_config(self, monkeypatch): + """Returns permissions when default_team_params is a dict.""" + import litellm + + from litellm.proxy.management_endpoints.team_endpoints import ( + _get_default_team_param, + ) + + monkeypatch.setattr( + litellm, + "default_team_params", + {"team_member_permissions": ["/key/generate", "/key/update"]}, + ) + + result = _get_default_team_param("team_member_permissions") + assert result == ["/key/generate", "/key/update"] + + def test_returns_scalar_fields_from_dict_config(self, monkeypatch): + """Returns scalar fields (max_budget, tpm_limit, etc.) from dict config.""" + import litellm + + from litellm.proxy.management_endpoints.team_endpoints import ( + _get_default_team_param, + ) + + monkeypatch.setattr( + litellm, + "default_team_params", + { + "max_budget": 100.0, + "budget_duration": "30d", + "tpm_limit": 200, + "rpm_limit": 500, + }, + ) + + assert _get_default_team_param("max_budget") == 100.0 + assert _get_default_team_param("budget_duration") == "30d" + assert _get_default_team_param("tpm_limit") == 200 + assert _get_default_team_param("rpm_limit") == 500 + + def test_returns_permissions_from_pydantic_config(self, monkeypatch): + """Returns permissions when default_team_params is a DefaultTeamSSOParams object.""" + import litellm + + from litellm.proxy.management_endpoints.team_endpoints import ( + _get_default_team_param, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + DefaultTeamSSOParams, + ) + + params = DefaultTeamSSOParams( + team_member_permissions=[ + KeyManagementRoutes.KEY_GENERATE, + KeyManagementRoutes.KEY_DELETE, + ] + ) + monkeypatch.setattr(litellm, "default_team_params", params) + + result = _get_default_team_param("team_member_permissions") + assert result == ["/key/generate", "/key/delete"] + + def test_returns_scalar_fields_from_pydantic_config(self, monkeypatch): + """Returns scalar fields from DefaultTeamSSOParams object.""" + import litellm + + from litellm.proxy.management_endpoints.team_endpoints import ( + _get_default_team_param, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + DefaultTeamSSOParams, + ) + + params = DefaultTeamSSOParams( + max_budget=250.0, + budget_duration="7d", + tpm_limit=1000, + rpm_limit=100, + ) + monkeypatch.setattr(litellm, "default_team_params", params) + + assert _get_default_team_param("max_budget") == 250.0 + assert _get_default_team_param("budget_duration") == "7d" + assert _get_default_team_param("tpm_limit") == 1000 + assert _get_default_team_param("rpm_limit") == 100 diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index f955a6134bf..7378b14cdf5 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -231,6 +231,56 @@ class TestProxySettingEndpoints: # Verify save_config was called exactly once assert mock_proxy_config["save_call_count"]() == 1 + def test_get_default_team_settings_includes_team_member_permissions_schema( + self, mock_proxy_config, mock_auth + ): + """Test that team_member_permissions field appears in schema with enum items""" + response = client.get("/get/default_team_settings") + + assert response.status_code == 200 + data = response.json() + + # Check that team_member_permissions is in the schema + props = data["field_schema"]["properties"] + assert "team_member_permissions" in props + + perm_schema = props["team_member_permissions"] + assert perm_schema["type"] == "array" + assert "items" in perm_schema + assert "enum" in perm_schema["items"] + # Verify some known enum values are present + enum_values = perm_schema["items"]["enum"] + assert "/key/generate" in enum_values + assert "/key/info" in enum_values + assert "/key/delete" in enum_values + + def test_update_default_team_settings_with_permissions( + self, mock_proxy_config, mock_auth, monkeypatch + ): + """Test updating default team settings with team_member_permissions""" + import litellm + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.setattr(litellm, "default_team_params", {}) + + new_settings = { + "models": ["gpt-4"], + "team_member_permissions": ["/key/generate", "/key/update", "/key/delete"], + } + + response = client.patch("/update/default_team_settings", json=new_settings) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + + settings = data["settings"] + assert settings["team_member_permissions"] == [ + "/key/generate", + "/key/update", + "/key/delete", + ] + def test_get_sso_settings(self, mock_proxy_config, mock_auth, monkeypatch): """Test getting the SSO settings from the dedicated database table""" from unittest.mock import AsyncMock, MagicMock diff --git a/ui/litellm-dashboard/src/components/TeamSSOSettings.test.tsx b/ui/litellm-dashboard/src/components/TeamSSOSettings.test.tsx index ae93b118799..5006afb61e2 100644 --- a/ui/litellm-dashboard/src/components/TeamSSOSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamSSOSettings.test.tsx @@ -18,14 +18,18 @@ vi.mock("./common_components/budget_duration_dropdown", () => { aria-label="Budget duration" > - - + + + ); BudgetDurationDropdown.displayName = "BudgetDurationDropdown"; return { default: BudgetDurationDropdown, - getBudgetDurationLabel: vi.fn((value: string) => `Budget: ${value}`), + getBudgetDurationLabel: vi.fn((value: string) => { + const map: Record = { "24h": "daily", "7d": "weekly", "30d": "monthly" }; + return map[value] || value; + }), }; }); @@ -56,6 +60,7 @@ vi.mock("./ModelSelect/ModelSelect", () => { vi.mock("antd", async (importOriginal) => { const actual = await importOriginal(); const React = await import("react"); + const SelectComponent = ({ value, onChange, @@ -88,37 +93,53 @@ vi.mock("antd", async (importOriginal) => { ); }; SelectComponent.displayName = "Select"; + const SelectOption = ({ value: optionValue, children: optionChildren }: { value: string; children: React.ReactNode }) => React.createElement("option", { value: optionValue }, optionChildren); SelectOption.displayName = "SelectOption"; SelectComponent.Option = SelectOption; - const Spin = ({ size }: { size?: string }) => React.createElement("div", { "data-testid": "spinner", "data-size": size }); + + const Spin = ({ size }: { size?: string }) => + React.createElement("div", { "data-testid": "spinner", "data-size": size }); Spin.displayName = "Spin"; - const Switch = ({ checked, onChange }: { checked: boolean; onChange: (checked: boolean) => void }) => + + const InputNumber = ({ + value, + onChange, + placeholder, + prefix, + }: { + value: number | null; + onChange: (value: number | null) => void; + placeholder?: string; + prefix?: string; + min?: number; + className?: string; + style?: React.CSSProperties; + }) => React.createElement("input", { - type: "checkbox", - role: "switch", - checked: checked, - onChange: (e) => onChange(e.target.checked), - "aria-label": "Toggle switch", + type: "number", + value: value ?? "", + onChange: (e: React.ChangeEvent) => { + const v = e.target.value === "" ? null : Number(e.target.value); + onChange(v); + }, + placeholder, + "data-prefix": prefix, + "aria-label": "number input", }); - Switch.displayName = "Switch"; - const Paragraph = ({ children }: { children: React.ReactNode }) => React.createElement("p", {}, children); - Paragraph.displayName = "Paragraph"; + InputNumber.displayName = "InputNumber"; + return { ...actual, Spin, - Switch, Select: SelectComponent, - Typography: { - Paragraph, - }, + InputNumber, }; }); const mockGetDefaultTeamSettings = vi.mocked(networking.getDefaultTeamSettings); const mockUpdateDefaultTeamSettings = vi.mocked(networking.updateDefaultTeamSettings); -const mockModelAvailableCall = vi.mocked(networking.modelAvailableCall); const mockNotificationsManager = vi.mocked(NotificationsManager); describe("TeamSSOSettings", () => { @@ -128,77 +149,33 @@ describe("TeamSSOSettings", () => { userRole: "admin", }; - const mockSettings = { + const mockSettingsResponse = { values: { - budget_duration: "monthly", max_budget: 1000, - enabled: true, - allowed_models: ["gpt-4", "claude-3"], + budget_duration: "30d", + tpm_limit: 500, + rpm_limit: 100, models: ["gpt-4"], - status: "active", - }, - field_schema: { - description: "Default team settings schema", - properties: { - budget_duration: { - type: "string", - description: "Budget duration setting", - }, - max_budget: { - type: "number", - description: "Maximum budget amount", - }, - enabled: { - type: "boolean", - description: "Enable feature", - }, - allowed_models: { - type: "array", - items: { - enum: ["gpt-4", "claude-3", "gpt-3.5-turbo"], - }, - description: "Allowed models", - }, - models: { - type: "array", - description: "Selected models", - }, - status: { - type: "string", - enum: ["active", "inactive", "pending"], - description: "Status", - }, - }, + team_member_permissions: ["/key/generate", "/key/update"], }, }; beforeEach(() => { vi.clearAllMocks(); - mockModelAvailableCall.mockResolvedValue({ - data: [{ id: "gpt-4" }, { id: "claude-3" }], - }); }); - it("should render", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Default Team Settings")).toBeInTheDocument(); - }); - }); + // --- Loading & Error States --- it("should show loading spinner while fetching settings", () => { - mockGetDefaultTeamSettings.mockImplementation(() => new Promise(() => { })); + mockGetDefaultTeamSettings.mockImplementation(() => new Promise(() => {})); renderWithProviders(); expect(screen.getByTestId("spinner")).toBeInTheDocument(); }); - it("should display message when no settings are available", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(null as any); + it("should display error message when fetch fails", async () => { + mockGetDefaultTeamSettings.mockRejectedValue(new Error("Fetch failed")); renderWithProviders(); @@ -207,6 +184,7 @@ describe("TeamSSOSettings", () => { screen.getByText("No team settings available or you do not have permission to view them."), ).toBeInTheDocument(); }); + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to fetch team settings"); }); it("should not fetch settings when access token is null", async () => { @@ -217,432 +195,273 @@ describe("TeamSSOSettings", () => { }); }); - it("should display settings fields with correct values", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); + // --- View Mode --- + + it("should render title and subtitle", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Default Team Settings")).toBeInTheDocument(); + expect(screen.getByText("These settings will be applied by default when creating new teams.")).toBeInTheDocument(); + }); + }); + + it("should render section headers", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Budget & Rate Limits")).toBeInTheDocument(); + expect(screen.getByText("Access & Permissions")).toBeInTheDocument(); + }); + }); + + it("should display all field labels and descriptions", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); renderWithProviders(); await waitFor(() => { - expect(screen.getByText("Budget Duration")).toBeInTheDocument(); expect(screen.getByText("Max Budget")).toBeInTheDocument(); + expect(screen.getByText("Budget Duration")).toBeInTheDocument(); + expect(screen.getByText("TPM Limit")).toBeInTheDocument(); + expect(screen.getByText("RPM Limit")).toBeInTheDocument(); + expect(screen.getByText("Models")).toBeInTheDocument(); + expect(screen.getByText("Team Member Permissions")).toBeInTheDocument(); }); - expect(screen.getByText("Budget: monthly")).toBeInTheDocument(); - expect(screen.getByText("1000")).toBeInTheDocument(); - const enabledTexts = screen.getAllByText("Enabled"); - expect(enabledTexts.length).toBeGreaterThan(0); + // Descriptions + expect(screen.getByText("Maximum budget (in USD) for new automatically created teams.")).toBeInTheDocument(); + expect(screen.getByText("How frequently the team's budget resets.")).toBeInTheDocument(); + }); + + it("should display formatted values in view mode", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); + + renderWithProviders(); + + await waitFor(() => { + // max_budget displayed with $ + expect(screen.getByText("$1,000")).toBeInTheDocument(); + // budget_duration through getBudgetDurationLabel + expect(screen.getByText("monthly")).toBeInTheDocument(); + // tpm_limit formatted + expect(screen.getByText("500")).toBeInTheDocument(); + // rpm_limit formatted + expect(screen.getByText("100")).toBeInTheDocument(); + }); + }); + + it("should display models as tags in view mode", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + }); + }); + + it("should display permissions as tags in view mode", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("/key/generate")).toBeInTheDocument(); + expect(screen.getByText("/key/update")).toBeInTheDocument(); + }); }); it("should display 'Not set' for null values", async () => { - const settingsWithNulls = { - ...mockSettings, + mockGetDefaultTeamSettings.mockResolvedValue({ values: { - ...mockSettings.values, max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + models: [], + team_member_permissions: [], }, - }; - mockGetDefaultTeamSettings.mockResolvedValue(settingsWithNulls); + }); renderWithProviders(); await waitFor(() => { - expect(screen.getByText("Not set")).toBeInTheDocument(); + const notSetElements = screen.getAllByText("Not set"); + // max_budget, budget_duration, tpm_limit, rpm_limit, models (empty), permissions (empty) + expect(notSetElements.length).toBeGreaterThanOrEqual(4); }); }); - it("should toggle edit mode when edit button is clicked", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); + // --- Edit Mode Toggle --- + + it("should toggle to edit mode when Edit Settings is clicked", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); renderWithProviders(); await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); }); - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); - expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Save Changes" })).toBeInTheDocument(); - expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Cancel/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Save Changes/i })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /Edit Settings/i })).not.toBeInTheDocument(); }); it("should cancel edit mode and reset values", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); renderWithProviders(); await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); }); - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); + await userEvent.click(screen.getByRole("button", { name: /Cancel/i })); - const cancelButton = screen.getByRole("button", { name: "Cancel" }); - await userEvent.click(cancelButton); - - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - expect(screen.queryByRole("button", { name: "Cancel" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /Cancel/i })).not.toBeInTheDocument(); }); - it("should save settings when save button is clicked", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - mockUpdateDefaultTeamSettings.mockResolvedValue({ - settings: mockSettings.values, - }); + // --- Edit Mode Fields --- + + it("should show budget duration dropdown in edit mode", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); renderWithProviders(); await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); }); - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); await waitFor(() => { - expect(screen.getByRole("button", { name: "Save Changes" })).toBeInTheDocument(); - }); - - const saveButton = screen.getByRole("button", { name: "Save Changes" }); - await userEvent.click(saveButton); - - await waitFor(() => { - expect(mockUpdateDefaultTeamSettings).toHaveBeenCalledWith("test-token", mockSettings.values); - }); - - expect(mockNotificationsManager.success).toHaveBeenCalledWith("Default team settings updated successfully"); - }); - - it("should show error notification when save fails", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - mockUpdateDefaultTeamSettings.mockRejectedValue(new Error("Save failed")); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Save Changes" })).toBeInTheDocument(); - }); - - const saveButton = screen.getByRole("button", { name: "Save Changes" }); - await userEvent.click(saveButton); - - await waitFor(() => { - expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to update team settings"); + expect(screen.getByTestId("budget-duration-dropdown")).toBeInTheDocument(); }); }); - it("should render boolean field as switch in edit mode", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); + it("should show ModelSelect in edit mode", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); renderWithProviders(); await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); }); - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); - - await waitFor(() => { - const switchElement = screen.getByRole("switch"); - expect(switchElement).toBeInTheDocument(); - expect(switchElement).toBeChecked(); - }); - }); - - it("should update boolean value when switch is toggled", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); - - await waitFor(() => { - expect(screen.getByRole("switch")).toBeInTheDocument(); - }); - - const switchElement = screen.getByRole("switch"); - await userEvent.click(switchElement); - - expect(switchElement).not.toBeChecked(); - }); - - it("should render budget duration dropdown in edit mode", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); - - await waitFor(() => { - expect(screen.getByLabelText("Budget duration")).toBeInTheDocument(); - }); - }); - - it("should update budget duration when dropdown value changes", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); - - await waitFor(() => { - expect(screen.getByLabelText("Budget duration")).toBeInTheDocument(); - }); - - const dropdown = screen.getByLabelText("Budget duration"); - await userEvent.selectOptions(dropdown, "daily"); - - expect(dropdown).toHaveValue("daily"); - }); - - it("should render text input for string fields in edit mode", async () => { - const settingsWithString = { - ...mockSettings, - field_schema: { - ...mockSettings.field_schema, - properties: { - ...mockSettings.field_schema.properties, - team_name: { - type: "string", - description: "Team name", - }, - }, - }, - values: { - ...mockSettings.values, - team_name: "Test Team", - }, - }; - mockGetDefaultTeamSettings.mockResolvedValue(settingsWithString); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); - - await waitFor(() => { - const textInput = screen.getByDisplayValue("Test Team"); - expect(textInput).toBeInTheDocument(); - }); - }); - - it("should render enum select for string enum fields in edit mode", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); - - await waitFor(() => { - const statusSelect = screen.getAllByRole("listbox")[0]; - expect(statusSelect).toBeInTheDocument(); - }); - }); - - it("should render multi-select for array enum fields in edit mode", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); - - await waitFor(() => { - const multiSelects = screen.getAllByRole("listbox"); - expect(multiSelects.length).toBeGreaterThan(0); - }); - }); - - it("should render ModelSelect for models field in edit mode", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); await waitFor(() => { expect(screen.getByTestId("model-select")).toBeInTheDocument(); }); }); - it("should display models as badges in view mode", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); + it("should show number inputs for budget and rate limits in edit mode", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); renderWithProviders(); await waitFor(() => { - const gpt4Elements = screen.getAllByText("gpt-4"); - expect(gpt4Elements.length).toBeGreaterThan(0); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); + + await waitFor(() => { + const numberInputs = screen.getAllByLabelText("number input"); + // max_budget, tpm_limit, rpm_limit + expect(numberInputs.length).toBe(3); }); }); - it("should display 'None' for empty arrays in view mode", async () => { - const settingsWithEmptyArray = { - ...mockSettings, - values: { - ...mockSettings.values, - models: [], - }, - }; - mockGetDefaultTeamSettings.mockResolvedValue(settingsWithEmptyArray); + it("should show permissions multi-select in edit mode", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); renderWithProviders(); await waitFor(() => { - const noneTexts = screen.getAllByText("None"); - expect(noneTexts.length).toBeGreaterThan(0); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); + + await waitFor(() => { + const listboxes = screen.getAllByRole("listbox"); + expect(listboxes.length).toBeGreaterThan(0); }); }); - it("should display schema description when available", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); + // --- Save --- + + it("should save settings and show success notification", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); + mockUpdateDefaultTeamSettings.mockResolvedValue({ + settings: mockSettingsResponse.values, + }); renderWithProviders(); await waitFor(() => { - expect(screen.getByText("Default team settings schema")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); + await userEvent.click(screen.getByRole("button", { name: /Save Changes/i })); + + await waitFor(() => { + expect(mockUpdateDefaultTeamSettings).toHaveBeenCalledWith("test-token", expect.any(Object)); + }); + + expect(mockNotificationsManager.success).toHaveBeenCalledWith("Default team settings updated successfully"); + + // Should exit edit mode after save + await waitFor(() => { + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); }); }); - it("should show error notification when fetching settings fails", async () => { - mockGetDefaultTeamSettings.mockRejectedValue(new Error("Fetch failed")); + it("should show error notification when save fails", async () => { + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); + mockUpdateDefaultTeamSettings.mockRejectedValue(new Error("Save failed")); renderWithProviders(); await waitFor(() => { - expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to fetch team settings"); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); }); - }); - it("should handle model fetch error gracefully", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - mockModelAvailableCall.mockRejectedValue(new Error("Model fetch failed")); - - renderWithProviders(); + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); + await userEvent.click(screen.getByRole("button", { name: /Save Changes/i })); await waitFor(() => { - expect(screen.getByText("Default Team Settings")).toBeInTheDocument(); + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to update team settings"); }); }); it("should disable cancel button while saving", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); + mockGetDefaultTeamSettings.mockResolvedValue(mockSettingsResponse); mockUpdateDefaultTeamSettings.mockImplementation( - () => new Promise((resolve) => setTimeout(() => resolve({ settings: mockSettings.values }), 100)), + () => new Promise((resolve) => setTimeout(() => resolve({ settings: mockSettingsResponse.values }), 100)), ); renderWithProviders(); await waitFor(() => { - expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); }); - const editButton = screen.getByRole("button", { name: "Edit Settings" }); - await userEvent.click(editButton); + await userEvent.click(screen.getByRole("button", { name: /Edit Settings/i })); + await userEvent.click(screen.getByRole("button", { name: /Save Changes/i })); - await waitFor(() => { - expect(screen.getByRole("button", { name: "Save Changes" })).toBeInTheDocument(); - }); - - const saveButton = screen.getByRole("button", { name: "Save Changes" }); - await userEvent.click(saveButton); - - const cancelButton = screen.getByRole("button", { name: "Cancel" }); - expect(cancelButton).toBeDisabled(); - }); - - it("should display field descriptions", async () => { - mockGetDefaultTeamSettings.mockResolvedValue(mockSettings); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Budget duration setting")).toBeInTheDocument(); - expect(screen.getByText("Maximum budget amount")).toBeInTheDocument(); - }); - }); - - it("should format field names by replacing underscores and capitalizing", async () => { - const settingsWithUnderscores = { - ...mockSettings, - field_schema: { - ...mockSettings.field_schema, - properties: { - ...mockSettings.field_schema.properties, - max_budget_per_user: { - type: "number", - description: "Max budget per user", - }, - }, - }, - values: { - ...mockSettings.values, - max_budget_per_user: 500, - }, - }; - mockGetDefaultTeamSettings.mockResolvedValue(settingsWithUnderscores); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Max Budget Per User")).toBeInTheDocument(); - }); - }); - - it("should display 'No schema information available' when schema is missing", async () => { - const settingsWithoutSchema = { - values: {}, - field_schema: null, - }; - mockGetDefaultTeamSettings.mockResolvedValue(settingsWithoutSchema); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("No schema information available")).toBeInTheDocument(); - }); + expect(screen.getByRole("button", { name: /Cancel/i })).toBeDisabled(); }); }); diff --git a/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx b/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx index 33bfc783afd..a9c07cdbcf8 100644 --- a/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx @@ -1,30 +1,96 @@ import React, { useState, useEffect } from "react"; -import { Card, Title, Text, Divider, Button, TextInput } from "@tremor/react"; -import { Typography, Spin, Switch, Select } from "antd"; -import { getDefaultTeamSettings, updateDefaultTeamSettings, modelAvailableCall } from "./networking"; +import { Card, Button, InputNumber, Typography, Spin, Select, Tag, Row, Col } from "antd"; +import { EditOutlined, SaveOutlined } from "@ant-design/icons"; +import { getDefaultTeamSettings, updateDefaultTeamSettings } from "./networking"; import BudgetDurationDropdown, { getBudgetDurationLabel } from "./common_components/budget_duration_dropdown"; import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key"; import NotificationsManager from "./molecules/notifications_manager"; import { ModelSelect } from "./ModelSelect/ModelSelect"; +const { Title, Text } = Typography; + interface TeamSSOSettingsProps { accessToken: string | null; userID: string; userRole: string; } -const TeamSSOSettings: React.FC = ({ accessToken, userID, userRole }) => { +const PERMISSION_OPTIONS = [ + "/key/generate", + "/key/update", + "/key/delete", + "/key/regenerate", + "/key/service-account/generate", + "/key/{key_id}/regenerate", + "/key/block", + "/key/unblock", + "/key/bulk_update", + "/key/{key_id}/reset_spend", +]; + +interface SettingRowProps { + label: string; + description: string; + isEditing: boolean; + viewContent: React.ReactNode; + editContent: React.ReactNode; +} + +const SettingRow: React.FC = ({ label, description, isEditing, viewContent, editContent }) => ( + + +
{label}
+
{description}
+ + +
{isEditing ? editContent : viewContent}
+ + +); + +const NotSet = () => Not set; + +const renderTags = (values: string[], displayFn?: (v: string) => string) => { + if (!values || values.length === 0) return ; + return ( +
+ {values.map((v) => ( + + {displayFn ? displayFn(v) : v} + + ))} +
+ ); +}; + +interface SettingsValues { + max_budget: number | null; + budget_duration: string | null; + tpm_limit: number | null; + rpm_limit: number | null; + models: string[]; + team_member_permissions: string[]; +} + +const DEFAULT_VALUES: SettingsValues = { + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + models: [], + team_member_permissions: [], +}; + +const TeamSSOSettings: React.FC = ({ accessToken }) => { const [loading, setLoading] = useState(true); - const [settings, setSettings] = useState(null); + const [values, setValues] = useState(DEFAULT_VALUES); const [isEditing, setIsEditing] = useState(false); - const [editedValues, setEditedValues] = useState({}); + const [editedValues, setEditedValues] = useState(DEFAULT_VALUES); const [saving, setSaving] = useState(false); - const [availableModels, setAvailableModels] = useState([]); - const { Paragraph } = Typography; - const { Option } = Select; + const [fetchError, setFetchError] = useState(false); useEffect(() => { - const fetchTeamSSOSettings = async () => { + const fetchSettings = async () => { if (!accessToken) { setLoading(false); return; @@ -32,39 +98,30 @@ const TeamSSOSettings: React.FC = ({ accessToken, userID, try { const data = await getDefaultTeamSettings(accessToken); - setSettings(data); - setEditedValues(data.values || {}); - - // Fetch available models - if (accessToken) { - try { - const modelResponse = await modelAvailableCall(accessToken, userID, userRole); - if (modelResponse && modelResponse.data) { - const modelNames = modelResponse.data.map((model: { id: string }) => model.id); - setAvailableModels(modelNames); - } - } catch (error) { - console.error("Error fetching available models:", error); - } - } + const fetched = { ...DEFAULT_VALUES, ...(data.values || {}) }; + setValues(fetched); + setEditedValues(fetched); } catch (error) { console.error("Error fetching team SSO settings:", error); + setFetchError(true); NotificationsManager.fromBackend("Failed to fetch team settings"); } finally { setLoading(false); } }; - fetchTeamSSOSettings(); + fetchSettings(); }, [accessToken]); - const handleSaveSettings = async () => { + const handleSave = async () => { if (!accessToken) return; setSaving(true); try { const updatedSettings = await updateDefaultTeamSettings(accessToken, editedValues); - setSettings({ ...settings, values: updatedSettings.settings }); + const newValues = { ...DEFAULT_VALUES, ...(updatedSettings.settings || {}) }; + setValues(newValues); + setEditedValues(newValues); setIsEditing(false); NotificationsManager.success("Default team settings updated successfully"); } catch (error) { @@ -75,129 +132,13 @@ const TeamSSOSettings: React.FC = ({ accessToken, userID, } }; - const handleTextInputChange = (key: string, value: any) => { - setEditedValues((prev: Record) => ({ - ...prev, - [key]: value, - })); + const handleCancel = () => { + setIsEditing(false); + setEditedValues(values); }; - const renderEditableField = (key: string, property: any, value: any) => { - const type = property.type; - - if (key === "budget_duration") { - return ( - handleTextInputChange(key, value)} - className="mt-2" - /> - ); - } else if (type === "boolean") { - return ( -
- handleTextInputChange(key, checked)} /> -
- ); - } else if (type === "array" && property.items?.enum) { - return ( - - ); - } else if (key === "models") { - return ( - handleTextInputChange(key, value)} - context="global" - style={{ width: "100%" }} - options={{ - includeSpecialOptions: true, - }} - /> - ); - } else if (type === "string" && property.enum) { - return ( - - ); - } else { - return ( - handleTextInputChange(key, e.target.value)} - placeholder={property.description || ""} - className="mt-2" - /> - ); - } - }; - - const renderValue = (key: string, value: any): JSX.Element => { - if (value === null || value === undefined) return Not set; - - if (key === "budget_duration") { - return {getBudgetDurationLabel(value)}; - } - - if (typeof value === "boolean") { - return {value ? "Enabled" : "Disabled"}; - } - - if (key === "models" && Array.isArray(value)) { - if (value.length === 0) return None; - - return ( -
- {value.map((model, index) => ( - - {getModelDisplayName(model)} - - ))} -
- ); - } - - if (typeof value === "object") { - if (Array.isArray(value)) { - if (value.length === 0) return None; - - return ( -
- {value.map((item, index) => ( - - {typeof item === "object" ? JSON.stringify(item) : String(item)} - - ))} -
- ); - } - - return 
{JSON.stringify(value, null, 2)}
; - } - - return {String(value)}; + const update = (key: K, value: SettingsValues[K]) => { + setEditedValues((prev) => ({ ...prev, [key]: value })); }; if (loading) { @@ -208,7 +149,7 @@ const TeamSSOSettings: React.FC = ({ accessToken, userID, ); } - if (!settings) { + if (fetchError) { return ( No team settings available or you do not have permission to view them. @@ -216,70 +157,166 @@ const TeamSSOSettings: React.FC = ({ accessToken, userID, ); } - // Dynamically render settings based on the schema - const renderSettings = () => { - const { values, field_schema } = settings; - - if (!field_schema || !field_schema.properties) { - return No schema information available; - } - - return Object.entries(field_schema.properties).map(([key, property]: [string, any]) => { - const value = values[key]; - const displayName = key.replace(/_/g, " ").replace(/\b\w/g, (l) => l.toUpperCase()); - - return ( -
- {displayName} - - {property.description || "No description available"} - - - {isEditing ? ( -
{renderEditableField(key, property, value)}
- ) : ( -
{renderValue(key, value)}
- )} -
- ); - }); - }; - return ( - -
- Default Team Settings - {!loading && - settings && - (isEditing ? ( -
- -
) : ( - - ))} + + )} +
- These settings will be applied by default when creating new teams. +
+ {/* Budget & Rate Limits */} +
+
Budget & Rate Limits
+
+ ${Number(values.max_budget).toLocaleString()} : + } + editContent={ + update("max_budget", v)} + placeholder="Not set" + prefix="$" + min={0} + /> + } + /> - {settings?.field_schema?.description && ( - {settings.field_schema.description} - )} - + {getBudgetDurationLabel(values.budget_duration)} : + } + editContent={ + update("budget_duration", v)} + style={{ maxWidth: 320 }} + /> + } + /> -
{renderSettings()}
+ {values.tpm_limit.toLocaleString()} : + } + editContent={ + update("tpm_limit", v)} + placeholder="Not set" + min={0} + /> + } + /> + + {values.rpm_limit.toLocaleString()} : + } + editContent={ + update("rpm_limit", v)} + placeholder="Not set" + min={0} + /> + } + /> +
+
+ + {/* Access & Permissions */} +
+
Access & Permissions
+
+ update("models", v)} + context="global" + style={{ width: "100%" }} + options={{ includeSpecialOptions: true }} + /> + } + /> + + update("team_member_permissions", v)} + placeholder="Select permissions" + tagRender={({ label, closable, onClose }) => ( + + {label} + + )} + > + {PERMISSION_OPTIONS.map((option) => ( + + {option} + + ))} + + } + /> +
+
+
); }; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 11a8a6f1fe3..bb9a0c1e015 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -7206,7 +7206,6 @@ export const updateDefaultTeamSettings = async (accessToken: string, settings: R const data = await response.json(); console.log("Updated default team settings:", data); - NotificationsManager.success("Default team settings updated successfully"); return data; } catch (error) { console.error("Failed to update default team settings:", error); diff --git a/ui/litellm-dashboard/tsconfig.json b/ui/litellm-dashboard/tsconfig.json index 5b0352feb98..d24bdd340f7 100644 --- a/ui/litellm-dashboard/tsconfig.json +++ b/ui/litellm-dashboard/tsconfig.json @@ -14,7 +14,7 @@ "moduleResolution": "bundler", "resolveJsonModule": true, "isolatedModules": true, - "jsx": "preserve", + "jsx": "react-jsx", "incremental": true, "plugins": [ { From 3ae9e070d97a41f5ab01da979c33ed1ffda0052b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Mar 2026 00:30:03 -0700 Subject: [PATCH 08/16] [Feature] UI - Usage: Auto-paginate daily spend data with progressive rendering Previously, EntityUsage only fetched page 1 of paginated daily spend endpoints, showing incomplete data. UsagePageView fetched all pages but blocked the UI until completion. This adds a reusable usePaginatedDailyActivity hook that fetches pages sequentially with 500ms delays, updates charts progressively, and supports cancellation on unmount or user action. Co-Authored-By: Claude Opus 4.6 --- .../components/EntityUsage/EntityUsage.tsx | 178 +++++++-------- .../UsagePage/components/UsagePageView.tsx | 177 +++++++-------- .../hooks/usePaginatedDailyActivity.ts | 209 ++++++++++++++++++ 3 files changed, 373 insertions(+), 191 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx index 92d9c25c6be..c1d5314096c 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx @@ -22,7 +22,8 @@ import { Text, Title, } from "@tremor/react"; -import React, { useEffect, useState } from "react"; +import { LoadingOutlined } from "@ant-design/icons"; +import React, { useMemo, useState } from "react"; import { ActivityMetrics, processActivityData } from "../../../activity_metrics"; import { UsageExportHeader } from "../../../EntityUsageExport"; import type { EntityType } from "../../../EntityUsageExport/types"; @@ -35,6 +36,7 @@ import { userDailyActivityCall, } from "../../../networking"; import { getProviderLogoAndName } from "../../../provider_info_helpers"; +import { usePaginatedDailyActivity } from "../../hooks/usePaginatedDailyActivity"; import { BreakdownMetrics, DailyData, EntityMetricWithMetadata, KeyMetricWithMetadata, TagUsage } from "../../types"; import { valueFormatterSpend } from "../../utils/value_formatters"; import EndpointUsage from "../EndpointUsage/EndpointUsage"; @@ -87,119 +89,64 @@ interface EntityUsageProps { dateValue: DateRangePickerValue; } +const ENTITY_FETCH_FNS: Record Promise> = { + tag: tagDailyActivityCall, + team: teamDailyActivityCall, + organization: organizationDailyActivityCall, + customer: customerDailyActivityCall, + agent: agentDailyActivityCall, + user: userDailyActivityCall, +}; + const EntityUsage: React.FC = ({ accessToken, entityType, entityId, entityList, dateValue }) => { - const [spendData, setSpendData] = useState({ - results: [], - metadata: { - total_spend: 0, - total_api_requests: 0, - total_successful_requests: 0, - total_failed_requests: 0, - total_tokens: 0, - }, - }); const { teams } = useTeams(); - - const [agentSpendData, setAgentSpendData] = useState({ - results: [], - metadata: { - total_spend: 0, - total_api_requests: 0, - total_successful_requests: 0, - total_failed_requests: 0, - total_tokens: 0, - }, - }); - - const modelMetrics = processActivityData(spendData, "models", teams || []); - const keyMetrics = processActivityData(spendData, "api_keys", teams || []); - const agentMetrics = entityType === "team" ? processActivityData(agentSpendData, "entities", teams || []) : {}; const [selectedTags, setSelectedTags] = useState([]); const [topKeysLimit, setTopKeysLimit] = useState(5); const [topModelsLimit, setTopModelsLimit] = useState(5); const [topAgentsLimit, setTopAgentsLimit] = useState(5); - const fetchSpendData = async () => { - if (!accessToken || !dateValue.from || !dateValue.to) return; - // Create new Date objects to avoid mutating the original dates - const startTime = new Date(dateValue.from); - const endTime = new Date(dateValue.to); + const startTime = useMemo(() => dateValue.from ? new Date(dateValue.from) : null, [dateValue.from]); + const endTime = useMemo(() => dateValue.to ? new Date(dateValue.to) : null, [dateValue.to]); - if (entityType === "tag") { - const data = await tagDailyActivityCall( - accessToken, - startTime, - endTime, - 1, - selectedTags.length > 0 ? selectedTags : null, - ); - setSpendData(data); - } else if (entityType === "team") { - const data = await teamDailyActivityCall( - accessToken, - startTime, - endTime, - 1, - selectedTags.length > 0 ? selectedTags : null, - ); - setSpendData(data); - } else if (entityType === "organization") { - const data = await organizationDailyActivityCall( - accessToken, - startTime, - endTime, - 1, - selectedTags.length > 0 ? selectedTags : null, - ); - setSpendData(data); - } else if (entityType === "customer") { - const data = await customerDailyActivityCall( - accessToken, - startTime, - endTime, - 1, - selectedTags.length > 0 ? selectedTags : null, - ); - setSpendData(data); - } else if (entityType === "agent") { - const data = await agentDailyActivityCall( - accessToken, - startTime, - endTime, - 1, - selectedTags.length > 0 ? selectedTags : null, - ); - setSpendData(data); - } else if (entityType === "user") { - const data = await userDailyActivityCall( - accessToken, - startTime, - endTime, - 1, - selectedTags.length > 0 ? selectedTags[0] : null, - ); - setSpendData(data); - } else { - throw new Error("Invalid entity type"); - } - }; + const entityFilterArg = useMemo(() => { + if (entityType === "user") return selectedTags.length > 0 ? selectedTags[0] : null; + return selectedTags.length > 0 ? selectedTags : null; + }, [entityType, selectedTags]); - const fetchAgentSpendData = async () => { - if (!accessToken || !dateValue.from || !dateValue.to || entityType !== "team") return; - const startTime = new Date(dateValue.from); - const endTime = new Date(dateValue.to); - try { - const data = await agentDailyActivityCall(accessToken, startTime, endTime, 1, null); - setAgentSpendData(data); - } catch (e) { - console.error("Failed to fetch agent activity data:", e); - } - }; + const fetchFn = ENTITY_FETCH_FNS[entityType]; + const enabled = !!accessToken && !!startTime && !!endTime; - useEffect(() => { - fetchSpendData(); - fetchAgentSpendData(); - }, [accessToken, dateValue, entityId, selectedTags]); + const { + data: spendDataRaw, + isFetchingMore, + progress, + cancelled, + cancel, + } = usePaginatedDailyActivity({ + fetchFn, + args: [accessToken, startTime, endTime, entityFilterArg], + enabled, + }); + + const spendData = spendDataRaw as unknown as EntitySpendData; + + const { + data: agentSpendDataRaw, + isFetchingMore: agentIsFetchingMore, + progress: agentProgress, + cancelled: agentCancelled, + cancel: agentCancel, + } = usePaginatedDailyActivity({ + fetchFn: agentDailyActivityCall, + args: [accessToken, startTime, endTime, null], + enabled: enabled && entityType === "team", + }); + + const agentSpendData = agentSpendDataRaw as unknown as EntitySpendData; + + const modelMetrics = processActivityData(spendData, "models", teams || []); + const keyMetrics = processActivityData(spendData, "api_keys", teams || []); + const agentMetrics = entityType === "team" ? processActivityData(agentSpendData, "entities", teams || []) : {}; const getTopModels = () => { const modelSpend: { [key: string]: any } = {}; @@ -448,6 +395,29 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti return (
+ {(isFetchingMore || cancelled || agentIsFetchingMore || agentCancelled) && ( +
+ {isFetchingMore && ( + <> + + Loading spend data... (page {progress.currentPage}/{progress.totalPages}) + + + )} + {cancelled && ( + + Showing partial data ({progress.currentPage}/{progress.totalPages} pages loaded) + + )} + {agentIsFetchingMore && entityType === "team" && ( + <> + + Loading agent data... (page {agentProgress.currentPage}/{agentProgress.totalPages}) + + + )} +
+ )} = ({ teams, organizations }) => { const { accessToken, userRole, userId: userID, premiumUser } = useAuthorized(); - const [userSpendData, setUserSpendData] = useState<{ - results: DailyData[]; - metadata: any; - }>({ results: [], metadata: {} }); + // Aggregated endpoint: try first, fall back to paginated if unavailable + const [aggregatedData, setAggregatedData] = useState<{ results: DailyData[]; metadata: any } | null>(null); + const [aggregatedFailed, setAggregatedFailed] = useState(false); + const [aggregatedLoading, setAggregatedLoading] = useState(false); // Separate loading states for better UX - const [loading, setLoading] = useState(false); const [isDateChanging, setIsDateChanging] = useState(false); // Create initial dates outside of state to prevent recreation @@ -173,6 +173,67 @@ const UsagePage: React.FC = ({ teams, organizations }) => { } }, [isAdmin, userID]); + // For non-admins, always pass their own user_id + const effectiveUserId = isAdmin ? selectedUserId : (userID || null); + + const startTime = useMemo(() => dateValue.from ? new Date(dateValue.from) : null, [dateValue.from]); + const endTime = useMemo(() => dateValue.to ? new Date(dateValue.to) : null, [dateValue.to]); + + // Try aggregated endpoint first, fall back to paginated on failure + const aggregatedFetchIdRef = useRef(0); + useEffect(() => { + if (!accessToken || !startTime || !endTime) return; + const fetchId = ++aggregatedFetchIdRef.current; + setAggregatedLoading(true); + setAggregatedFailed(false); + setAggregatedData(null); + + userDailyActivityAggregatedCall(accessToken, startTime, endTime, effectiveUserId) + .then((data) => { + if (aggregatedFetchIdRef.current !== fetchId) return; + setAggregatedData(data); + setAggregatedLoading(false); + setIsDateChanging(false); + }) + .catch(() => { + if (aggregatedFetchIdRef.current !== fetchId) return; + setAggregatedFailed(true); + setAggregatedLoading(false); + }); + }, [accessToken, startTime, endTime, effectiveUserId]); + + // Paginated fallback — only enabled when aggregated endpoint fails + const paginatedResult = usePaginatedDailyActivity({ + fetchFn: userDailyActivityCall, + args: [accessToken, startTime, endTime, effectiveUserId], + enabled: aggregatedFailed && !!accessToken && !!startTime && !!endTime, + }); + + // Derive userSpendData from whichever source is active + const userSpendData = useMemo(() => { + if (aggregatedData) return aggregatedData; + if (aggregatedFailed) return paginatedResult.data; + return { results: [] as DailyData[], metadata: {} as any }; + }, [aggregatedData, aggregatedFailed, paginatedResult.data]); + + const loading = aggregatedLoading || paginatedResult.loading; + + // Clear isDateChanging when paginated data starts arriving + useEffect(() => { + if (aggregatedFailed && !paginatedResult.loading && paginatedResult.data.results.length > 0) { + setIsDateChanging(false); + } + }, [aggregatedFailed, paginatedResult.loading, paginatedResult.data.results.length]); + + // Super responsive date change handler + const handleDateChange = useCallback((newValue: DateRangePickerValue) => { + // Instant visual feedback + setIsDateChanging(true); + + // Update date immediately for UI responsiveness + setDateValue(newValue); + }, []); + // Derived states from userSpendData const totalSpend = userSpendData.metadata?.total_spend || 0; @@ -362,87 +423,6 @@ const UsagePage: React.FC = ({ teams, organizations }) => { .slice(0, topKeysLimit); }, [userSpendData.results, topKeysLimit]); - const fetchUserSpendData = useCallback(async () => { - if (!accessToken || !dateValue.from || !dateValue.to) return; - - // For non-admins, always pass their own user_id - const effectiveUserId = isAdmin ? selectedUserId : (userID || null); - - setLoading(true); - - // Create new Date objects to avoid mutating the original dates - const startTime = new Date(dateValue.from); - const endTime = new Date(dateValue.to); - - try { - // Prefer aggregated endpoint to avoid many page requests - try { - const aggregated = await userDailyActivityAggregatedCall(accessToken, startTime, endTime, effectiveUserId); - setUserSpendData(aggregated); - return; - } catch (e) { - // Fallback to paginated calls if aggregated endpoint is unavailable - } - - const firstPageData = await userDailyActivityCall(accessToken, startTime, endTime, 1, effectiveUserId); - - if (firstPageData.metadata.total_pages <= 1) { - setUserSpendData(firstPageData); - return; - } - - const allResults = [...firstPageData.results]; - const aggregatedMetadata = { ...firstPageData.metadata }; - - for (let page = 2; page <= firstPageData.metadata.total_pages; page++) { - const pageData = await userDailyActivityCall(accessToken, startTime, endTime, page, effectiveUserId); - allResults.push(...pageData.results); - if (pageData.metadata) { - aggregatedMetadata.total_spend = (aggregatedMetadata.total_spend || 0) + (pageData.metadata.total_spend || 0); - aggregatedMetadata.total_api_requests = (aggregatedMetadata.total_api_requests || 0) + (pageData.metadata.total_api_requests || 0); - aggregatedMetadata.total_successful_requests = (aggregatedMetadata.total_successful_requests || 0) + (pageData.metadata.total_successful_requests || 0); - aggregatedMetadata.total_failed_requests = (aggregatedMetadata.total_failed_requests || 0) + (pageData.metadata.total_failed_requests || 0); - aggregatedMetadata.total_tokens = (aggregatedMetadata.total_tokens || 0) + (pageData.metadata.total_tokens || 0); - aggregatedMetadata.total_prompt_tokens = (aggregatedMetadata.total_prompt_tokens || 0) + (pageData.metadata.total_prompt_tokens || 0); - aggregatedMetadata.total_completion_tokens = (aggregatedMetadata.total_completion_tokens || 0) + (pageData.metadata.total_completion_tokens || 0); - aggregatedMetadata.total_cache_read_input_tokens = (aggregatedMetadata.total_cache_read_input_tokens || 0) + (pageData.metadata.total_cache_read_input_tokens || 0); - aggregatedMetadata.total_cache_creation_input_tokens = (aggregatedMetadata.total_cache_creation_input_tokens || 0) + (pageData.metadata.total_cache_creation_input_tokens || 0); - } - } - - setUserSpendData({ - results: allResults, - metadata: aggregatedMetadata, - }); - } catch (error) { - console.error("Error fetching user spend data:", error); - } finally { - setLoading(false); - setIsDateChanging(false); - } - }, [accessToken, dateValue.from, dateValue.to, selectedUserId, isAdmin, userID]); - - // Super responsive date change handler - const handleDateChange = useCallback((newValue: DateRangePickerValue) => { - // Instant visual feedback - setIsDateChanging(true); - setLoading(true); - - // Update date immediately for UI responsiveness - setDateValue(newValue); - }, []); - - // Debounced effect for data fetching with shorter delay - useEffect(() => { - if (!dateValue.from || !dateValue.to) return; - - const timeoutId = setTimeout(() => { - fetchUserSpendData(); - }, 50); // Very short debounce - - return () => clearTimeout(timeoutId); - }, [fetchUserSpendData]); - const sortedDailyResults = useMemo( () => [...userSpendData.results].sort((a, b) => new Date(a.date).getTime() - new Date(b.date).getTime()), [userSpendData.results], @@ -502,6 +482,29 @@ const UsagePage: React.FC = ({ teams, organizations }) => { />
+ {(paginatedResult.isFetchingMore || paginatedResult.cancelled) && ( +
+ {paginatedResult.isFetchingMore && ( + <> + + + Loading spend data... (page {paginatedResult.progress.currentPage}/{paginatedResult.progress.totalPages}) + + + + )} + {paginatedResult.cancelled && ( + + Showing partial data ({paginatedResult.progress.currentPage}/{paginatedResult.progress.totalPages} pages loaded) + + )} +
+ )} {/* Your Usage Panel */} {usageView === "global" && ( <> diff --git a/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts b/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts new file mode 100644 index 00000000000..19872bf31a8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts @@ -0,0 +1,209 @@ +import { useCallback, useEffect, useRef, useState } from "react"; +import { DailyData } from "../types"; + +export interface PaginationProgress { + currentPage: number; + totalPages: number; +} + +/** Delay between sequential page fetches (ms) to avoid overloading the backend. */ +const PAGE_FETCH_DELAY_MS = 500; + +/** The metadata fields returned by the daily activity API that should be summed across pages. */ +const SUMMABLE_METADATA_KEYS = [ + "total_spend", + "total_prompt_tokens", + "total_completion_tokens", + "total_tokens", + "total_api_requests", + "total_successful_requests", + "total_failed_requests", + "total_cache_read_input_tokens", + "total_cache_creation_input_tokens", +] as const; + +interface DailyActivityResponse { + results: DailyData[]; + metadata: Record; +} + +type FetchPageFn = (...args: any[]) => Promise; + +interface UsePaginatedDailyActivityParams { + /** The API call function (e.g., userDailyActivityCall). */ + fetchFn: FetchPageFn; + /** Arguments to pass to fetchFn: [accessToken, startTime, endTime, ...extraArgs]. Page is injected by the hook at index 3. */ + args: any[]; + /** Whether the hook should fetch. Set to false to disable. */ + enabled: boolean; +} + +interface UsePaginatedDailyActivityReturn { + data: DailyActivityResponse; + loading: boolean; + isFetchingMore: boolean; + progress: PaginationProgress; + cancelled: boolean; + cancel: () => void; +} + +const EMPTY_DATA: DailyActivityResponse = { + results: [], + metadata: { + total_spend: 0, + total_prompt_tokens: 0, + total_completion_tokens: 0, + total_tokens: 0, + total_api_requests: 0, + total_successful_requests: 0, + total_failed_requests: 0, + total_cache_read_input_tokens: 0, + total_cache_creation_input_tokens: 0, + total_pages: 1, + has_more: false, + page: 1, + }, +}; + +function sumMetadata( + a: Record, + b: Record, +): Record { + const result = { ...a }; + for (const key of SUMMABLE_METADATA_KEYS) { + result[key] = (a[key] || 0) + (b[key] || 0); + } + return result; +} + +/** + * Hook that auto-paginates daily activity endpoints, updating state after each + * page so charts render progressively. Cancels on unmount, param changes, or + * manual cancel(). + * + * The `args` array should contain every argument the fetchFn expects EXCEPT + * the `page` parameter. The hook injects `page` as the 4th argument (index 3), + * matching the signature of all daily activity calls: + * (accessToken, startTime, endTime, page, ...rest) + */ +export function usePaginatedDailyActivity({ + fetchFn, + args, + enabled, +}: UsePaginatedDailyActivityParams): UsePaginatedDailyActivityReturn { + const [data, setData] = useState(EMPTY_DATA); + const [loading, setLoading] = useState(false); + const [isFetchingMore, setIsFetchingMore] = useState(false); + const [progress, setProgress] = useState({ + currentPage: 0, + totalPages: 0, + }); + const [cancelled, setCancelled] = useState(false); + + const fetchIdRef = useRef(0); + const cancelledRef = useRef(false); + + const cancel = useCallback(() => { + cancelledRef.current = true; + setCancelled(true); + setIsFetchingMore(false); + }, []); + + useEffect(() => { + if (!enabled) { + setData(EMPTY_DATA); + setLoading(false); + setIsFetchingMore(false); + setProgress({ currentPage: 0, totalPages: 0 }); + setCancelled(false); + return; + } + + const currentFetchId = ++fetchIdRef.current; + cancelledRef.current = false; + setCancelled(false); + + const isStale = () => + fetchIdRef.current !== currentFetchId || cancelledRef.current; + + const run = async () => { + setLoading(true); + setIsFetchingMore(false); + setProgress({ currentPage: 1, totalPages: 1 }); + + try { + // Inject page=1 as the 4th argument. + const argsWithPage = [...args.slice(0, 3), 1, ...args.slice(3)]; + const firstPage = await fetchFn(...argsWithPage); + + if (isStale()) return; + + setData(firstPage); + + const totalPages = firstPage.metadata?.total_pages || 1; + + setProgress({ currentPage: 1, totalPages }); + + if (totalPages <= 1) { + setLoading(false); + return; + } + + // More pages — start fetching sequentially. + setLoading(false); + setIsFetchingMore(true); + + let accumulatedResults = [...firstPage.results]; + let accumulatedMetadata = { ...firstPage.metadata }; + + for (let page = 2; page <= totalPages; page++) { + if (isStale()) return; + + // Small delay to avoid overwhelming the backend. + await new Promise((resolve) => + setTimeout(resolve, PAGE_FETCH_DELAY_MS), + ); + + if (isStale()) return; + + const argsForPage = [...args.slice(0, 3), page, ...args.slice(3)]; + const pageData = await fetchFn(...argsForPage); + + if (isStale()) return; + + accumulatedResults = [...accumulatedResults, ...pageData.results]; + accumulatedMetadata = sumMetadata( + accumulatedMetadata, + pageData.metadata, + ); + accumulatedMetadata.total_pages = totalPages; + accumulatedMetadata.has_more = page < totalPages; + accumulatedMetadata.page = page; + + setData({ + results: accumulatedResults, + metadata: accumulatedMetadata, + }); + setProgress({ currentPage: page, totalPages }); + } + + setIsFetchingMore(false); + } catch (error) { + if (!isStale()) { + console.error("Error fetching daily activity:", error); + setLoading(false); + setIsFetchingMore(false); + } + } + }; + + run(); + + return () => { + fetchIdRef.current++; + }; + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [enabled, fetchFn, ...args]); + + return { data, loading, isFetchingMore, progress, cancelled, cancel }; +} From 83e6096dae150f925f695364f842fa330144292a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Mar 2026 09:37:44 -0700 Subject: [PATCH 09/16] [Feature] UI - Internal Users: Add/remove team membership from user info page Co-Authored-By: Claude Opus 4.6 --- .../view_users/user_info_view.test.tsx | 200 +++++++++-- .../components/view_users/user_info_view.tsx | 339 ++++++++++++++---- 2 files changed, 451 insertions(+), 88 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx b/ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx index 201c686de2c..481a2c6668d 100644 --- a/ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx @@ -1,32 +1,47 @@ -import { render, screen } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; +import { render, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi, beforeEach } from "vitest"; import UserInfoView from "./user_info_view"; -vi.mock("../networking", () => { - const MOCK_USER_DATA = { - user_id: "user-123", - user_email: "test@example.com", - user_alias: "Test Alias", - user_role: "admin", - spend: 0, - max_budget: 100, - models: [], - budget_duration: "30d", - budget_reset_at: null, - metadata: {}, - created_at: "2025-01-01T00:00:00.000Z", - updated_at: "2025-01-02T00:00:00.000Z", - sso_user_id: null, - teams: [], - }; +const mockTeamMemberAddCall = vi.fn(); +const mockTeamMemberDeleteCall = vi.fn(); +const mockTeamListCall = vi.fn(); +const mockUserGetInfoV2 = vi.fn(); +const mockTeamInfoCall = vi.fn(); +const MOCK_USER_DATA = { + user_id: "user-123", + user_email: "test@example.com", + user_alias: "Test Alias", + user_role: "admin", + spend: 0, + max_budget: 100, + models: [], + budget_duration: "30d", + budget_reset_at: null, + metadata: {}, + created_at: "2025-01-01T00:00:00.000Z", + updated_at: "2025-01-02T00:00:00.000Z", + sso_user_id: null, + teams: ["team-1", "team-2"], +}; + +const MOCK_USER_DATA_NO_TEAMS = { + ...MOCK_USER_DATA, + teams: [], +}; + +vi.mock("../networking", () => { return { - userGetInfoV2: vi.fn().mockResolvedValue(MOCK_USER_DATA), + userGetInfoV2: (...args: any[]) => mockUserGetInfoV2(...args), userDeleteCall: vi.fn(), userUpdateUserCall: vi.fn(), modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), invitationCreateCall: vi.fn(), - teamInfoCall: vi.fn().mockResolvedValue({ team_alias: "Test Team" }), + teamInfoCall: (...args: any[]) => mockTeamInfoCall(...args), + teamListCall: (...args: any[]) => mockTeamListCall(...args), + teamMemberAddCall: (...args: any[]) => mockTeamMemberAddCall(...args), + teamMemberDeleteCall: (...args: any[]) => mockTeamMemberDeleteCall(...args), getProxyBaseUrl: () => "https://litellm.test", }; }); @@ -36,10 +51,30 @@ describe("UserInfoView", () => { userId: "user-123", onClose: vi.fn(), accessToken: "test-token", - userRole: null, + userRole: null as string | null, possibleUIRoles: null, }; + beforeEach(() => { + vi.clearAllMocks(); + mockUserGetInfoV2.mockResolvedValue(MOCK_USER_DATA); + mockTeamInfoCall.mockImplementation((_token: string, teamId: string) => { + const teamMap: Record = { + "team-1": { team_id: "team-1", team_info: { team_alias: "Alpha Team" } }, + "team-2": { team_id: "team-2", team_info: { team_alias: "Beta Team" } }, + "team-3": { team_id: "team-3", team_info: { team_alias: "Gamma Team" } }, + }; + return Promise.resolve(teamMap[teamId] || { team_id: teamId, team_info: { team_alias: null } }); + }); + mockTeamListCall.mockResolvedValue([ + { team_id: "team-1", team_alias: "Alpha Team" }, + { team_id: "team-2", team_alias: "Beta Team" }, + { team_id: "team-3", team_alias: "Gamma Team" }, + ]); + mockTeamMemberAddCall.mockResolvedValue({}); + mockTeamMemberDeleteCall.mockResolvedValue({}); + }); + it("should render the loading state", () => { render(); @@ -60,4 +95,125 @@ describe("UserInfoView", () => { const aliases = await screen.findAllByText("Test Alias"); expect(aliases.length).toBeGreaterThan(0); }); + + it("should render teams in a table with team names", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Alpha Team")).toBeInTheDocument(); + expect(screen.getByText("Beta Team")).toBeInTheDocument(); + }); + }); + + it("should show 'No teams' when user has no teams", async () => { + mockUserGetInfoV2.mockResolvedValue(MOCK_USER_DATA_NO_TEAMS); + render(); + + await waitFor(() => { + expect(screen.getByText("No teams")).toBeInTheDocument(); + }); + }); + + it("should show Add Team button for proxy admins", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Add Team")).toBeInTheDocument(); + }); + }); + + it("should not show Add Team button for non-proxy-admins", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Alpha Team")).toBeInTheDocument(); + }); + expect(screen.queryByText("Add Team")).not.toBeInTheDocument(); + }); + + it("should show delete buttons for proxy admins", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Alpha Team")).toBeInTheDocument(); + }); + // Should have the Actions column header + expect(screen.getByText("Actions")).toBeInTheDocument(); + }); + + it("should not show delete buttons for non-proxy-admins", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Alpha Team")).toBeInTheDocument(); + }); + expect(screen.queryByText("Actions")).not.toBeInTheDocument(); + }); + + it("should open the add team modal when Add Team is clicked", async () => { + const user = userEvent.setup(); + render(); + + await waitFor(() => { + expect(screen.getByText("Add Team")).toBeInTheDocument(); + }); + + await user.click(screen.getByText("Add Team")); + + await waitFor(() => { + expect(screen.getByText("Add User to Team")).toBeInTheDocument(); + }); + expect(mockTeamListCall).toHaveBeenCalledWith("test-token", null); + }); + + it("should open remove confirmation modal when delete is clicked", async () => { + const user = userEvent.setup(); + render(); + + await waitFor(() => { + expect(screen.getByText("Alpha Team")).toBeInTheDocument(); + }); + + // Find the row with Alpha Team and click its delete button + const alphaRow = screen.getByText("Alpha Team").closest("tr")!; + const deleteButton = within(alphaRow).getByRole("button"); + await user.click(deleteButton); + + await waitFor(() => { + expect(screen.getByText("Remove from Team")).toBeInTheDocument(); + expect(screen.getByText(/Removing this user from the team will also delete any keys/)).toBeInTheDocument(); + }); + }); + + it("should call teamMemberDeleteCall when remove is confirmed", async () => { + const user = userEvent.setup(); + render(); + + await waitFor(() => { + expect(screen.getByText("Alpha Team")).toBeInTheDocument(); + }); + + // Click delete on Alpha Team + const alphaRow = screen.getByText("Alpha Team").closest("tr")!; + const deleteButton = within(alphaRow).getByRole("button"); + await user.click(deleteButton); + + // Confirm deletion + await waitFor(() => { + expect(screen.getByText("Remove from Team")).toBeInTheDocument(); + }); + + // The DeleteResourceModal's OK button has text "Delete" - find it within the modal + const modal = screen.getByText("Remove from Team").closest(".ant-modal") as HTMLElement; + const deleteConfirmButton = within(modal).getByRole("button", { name: /delete/i }); + await user.click(deleteConfirmButton); + + await waitFor(() => { + expect(mockTeamMemberDeleteCall).toHaveBeenCalledWith( + "test-token", + "team-1", + { role: "user", user_id: "user-123" } + ); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx index d0684d81e5d..e5ec60642fc 100644 --- a/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx @@ -1,6 +1,9 @@ import React, { useState } from "react"; -import { Card, Text, Button, Grid, Tab, TabList, TabGroup, TabPanel, TabPanels, Title, Badge } from "@tremor/react"; -import { ArrowLeftIcon, TrashIcon, RefreshIcon } from "@heroicons/react/outline"; +import { + Card, Text, Button, Grid, Tab, TabList, TabGroup, TabPanel, TabPanels, Title, + Table, TableHead, TableBody, TableRow, TableHeaderCell, TableCell, +} from "@tremor/react"; +import { ArrowLeftIcon, TrashIcon, RefreshIcon, PlusIcon } from "@heroicons/react/outline"; import { userGetInfoV2, UserInfoV2Response, @@ -10,8 +13,12 @@ import { invitationCreateCall, getProxyBaseUrl, teamInfoCall, + teamListCall, + teamMemberAddCall, + teamMemberDeleteCall, + Member, } from "../networking"; -import { Button as AntdButton } from "antd"; +import { Button as AntdButton, Modal, Select as AntdSelect, Form, Tooltip } from "antd"; import { rolesWithWriteAccess } from "../../utils/roles"; import { UserEditView } from "../user_edit_view"; import OnboardingModal, { InvitationLink } from "../onboarding_link"; @@ -61,6 +68,15 @@ export default function UserInfoView({ const [activeTab, setActiveTab] = useState(initialTab); const [copiedStates, setCopiedStates] = useState>({}); const [isTeamsExpanded, setIsTeamsExpanded] = useState(false); + const [isAddTeamModalOpen, setIsAddTeamModalOpen] = useState(false); + const [isRemoveTeamModalOpen, setIsRemoveTeamModalOpen] = useState(false); + const [teamToRemove, setTeamToRemove] = useState(null); + const [isAddingTeam, setIsAddingTeam] = useState(false); + const [isRemovingTeam, setIsRemovingTeam] = useState(false); + const [allTeams, setAllTeams] = useState>([]); + const [selectedTeamId, setSelectedTeamId] = useState(""); + const [selectedRole, setSelectedRole] = useState("user"); + const [isLoadingTeams, setIsLoadingTeams] = useState(false); React.useEffect(() => { setBaseUrl(getProxyBaseUrl()); @@ -82,7 +98,7 @@ export default function UserInfoView({ const teamData = await teamInfoCall(accessToken, teamId); return { team_id: teamId, - team_alias: teamData?.team_alias || null, + team_alias: teamData?.team_info?.team_alias || null, }; } catch { return { team_id: teamId, team_alias: null }; @@ -111,6 +127,118 @@ export default function UserInfoView({ fetchData(); }, [accessToken, userId, userRole]); + const isProxyAdmin = userRole === "proxy_admin" || userRole === "Admin"; + + const fetchAllTeams = async () => { + if (!accessToken) return; + setIsLoadingTeams(true); + try { + const teams = await teamListCall(accessToken, null); + setAllTeams( + (teams || []).map((t: any) => ({ + team_id: t.team_id, + team_alias: t.team_alias || t.team_id, + })) + ); + } catch (error) { + console.error("Error fetching teams:", error); + } finally { + setIsLoadingTeams(false); + } + }; + + const handleOpenAddTeamModal = () => { + setSelectedTeamId(""); + setSelectedRole("user"); + setIsAddTeamModalOpen(true); + fetchAllTeams(); + }; + + const handleAddTeamSubmit = async () => { + if (!accessToken || !selectedTeamId) return; + setIsAddingTeam(true); + try { + const member: Member = { + role: selectedRole, + user_id: userId, + }; + await teamMemberAddCall(accessToken, selectedTeamId, member); + NotificationsManager.success("User added to team successfully"); + setIsAddTeamModalOpen(false); + // Re-fetch user data to refresh teams + const data = await userGetInfoV2(accessToken, userId); + setUserData(data); + if (data.teams && data.teams.length > 0) { + const teamPromises = data.teams.map(async (teamId: string) => { + try { + const teamData = await teamInfoCall(accessToken, teamId); + return { team_id: teamId, team_alias: teamData?.team_info?.team_alias || null }; + } catch { + return { team_id: teamId, team_alias: null }; + } + }); + setTeamDetails(await Promise.all(teamPromises)); + } else { + setTeamDetails([]); + } + } catch (error: any) { + console.error("Error adding user to team:", error); + NotificationsManager.fromBackend(error?.message || "Failed to add user to team"); + } finally { + setIsAddingTeam(false); + } + }; + + const handleOpenRemoveTeamModal = (team: TeamDisplayInfo) => { + setTeamToRemove(team); + setIsRemoveTeamModalOpen(true); + }; + + const handleRemoveTeamConfirm = async () => { + if (!accessToken || !teamToRemove) return; + setIsRemovingTeam(true); + try { + const member: Member = { + role: "user", + user_id: userId, + }; + await teamMemberDeleteCall(accessToken, teamToRemove.team_id, member); + NotificationsManager.success("User removed from team successfully"); + setIsRemoveTeamModalOpen(false); + setTeamToRemove(null); + // Re-fetch user data to refresh teams + const data = await userGetInfoV2(accessToken, userId); + setUserData(data); + if (data.teams && data.teams.length > 0) { + const teamPromises = data.teams.map(async (teamId: string) => { + try { + const teamData = await teamInfoCall(accessToken, teamId); + return { team_id: teamId, team_alias: teamData?.team_info?.team_alias || null }; + } catch { + return { team_id: teamId, team_alias: null }; + } + }); + setTeamDetails(await Promise.all(teamPromises)); + } else { + setTeamDetails([]); + } + } catch (error: any) { + console.error("Error removing user from team:", error); + NotificationsManager.fromBackend(error?.message || "Failed to remove user from team"); + } finally { + setIsRemovingTeam(false); + } + }; + + const handleRemoveTeamCancel = () => { + setIsRemoveTeamModalOpen(false); + setTeamToRemove(null); + }; + + const availableTeamsForAdd = allTeams.filter( + (t) => !teamDetails.some((td) => td.team_id === t.team_id) + ); + const handleResetPassword = async () => { if (!accessToken) { NotificationsManager.fromBackend("Access token not found"); @@ -312,37 +440,72 @@ export default function UserInfoView({ - Teams +
+ Teams + {isProxyAdmin && ( + + )} +
{teamDetails.length > 0 ? ( -
- {teamDetails.slice(0, isTeamsExpanded ? teamDetails.length : 20).map((team, index) => ( - - {team.team_alias || team.team_id} - - ))} - {!isTeamsExpanded && teamDetails.length > 20 && ( - setIsTeamsExpanded(true)} - > - +{teamDetails.length - 20} more - - )} - {isTeamsExpanded && teamDetails.length > 20 && ( - setIsTeamsExpanded(false)} - > - Show Less - - )} +
+ + + + Team Name + {isProxyAdmin && Actions} + + + + {teamDetails.slice(0, isTeamsExpanded ? teamDetails.length : 20).map((team) => ( + + {team.team_alias || team.team_id} + {isProxyAdmin && ( + +
) : ( No teams )} + {!isTeamsExpanded && teamDetails.length > 20 && ( + + )} + {isTeamsExpanded && teamDetails.length > 20 && ( + + )}
@@ -434,43 +597,6 @@ export default function UserInfoView({
-
- Teams -
- {teamDetails.length > 0 ? ( - <> - {teamDetails.slice(0, isTeamsExpanded ? teamDetails.length : 20).map((team, index) => ( - - {team.team_alias || team.team_id} - - ))} - {!isTeamsExpanded && teamDetails.length > 20 && ( - setIsTeamsExpanded(true)} - > - +{teamDetails.length - 20} more - - )} - {isTeamsExpanded && teamDetails.length > 20 && ( - setIsTeamsExpanded(false)} - > - Show Less - - )} - - ) : ( - No teams - )} -
-
-
Personal Models
@@ -519,6 +645,87 @@ export default function UserInfoView({ invitationLinkData={invitationLinkData} modalType="resetPassword" /> + + {/* Delete Team Member Modal */} + + + {/* Add to Team Modal */} + setIsAddTeamModalOpen(false)} + footer={null} + width={500} + maskClosable={!isAddingTeam} + > +
+ + { + const team = availableTeamsForAdd.find((t) => t.team_id === option?.value); + if (!team) return false; + return team.team_alias.toLowerCase().includes(input.toLowerCase()); + }} + loading={isLoadingTeams} + > + {availableTeamsForAdd.map((team) => ( + + {team.team_alias} + + ))} + + + + + + + + user + - Can view team info, but not manage it + + + + + admin + - Can create team keys, add members, and manage settings + + + + + +
+ + {isAddingTeam ? "Adding..." : "Add to Team"} + +
+
+
); } From 0cd4a681579d21e6497824f91b0c9b789299846b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Mar 2026 09:47:56 -0700 Subject: [PATCH 10/16] [Fix] Add missing networking mocks to CreateKeyPage test The test's partial vi.mock of @/components/networking was missing the daily activity call exports now imported by EntityUsage via ENTITY_FETCH_FNS. Co-Authored-By: Claude Opus 4.6 --- .../tests/CreateKeyPage.expiredToken.test.tsx | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index 8b05def9ba3..1d725572a33 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -77,6 +77,14 @@ vi.mock("@/components/networking", () => { // Called when decoding a valid token setGlobalLitellmHeaderName: vi.fn(), Organization: {}, + // Daily activity calls used by UsagePage components in the render tree + tagDailyActivityCall: vi.fn().mockResolvedValue({ results: [], metadata: {} }), + teamDailyActivityCall: vi.fn().mockResolvedValue({ results: [], metadata: {} }), + organizationDailyActivityCall: vi.fn().mockResolvedValue({ results: [], metadata: {} }), + customerDailyActivityCall: vi.fn().mockResolvedValue({ results: [], metadata: {} }), + agentDailyActivityCall: vi.fn().mockResolvedValue({ results: [], metadata: {} }), + userDailyActivityCall: vi.fn().mockResolvedValue({ results: [], metadata: {} }), + userDailyActivityAggregatedCall: vi.fn().mockResolvedValue({ results: [], metadata: {} }), }; }); From db37f3109943696a2cc9be05e1c81c8d93f0fcc6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Mar 2026 09:55:57 -0700 Subject: [PATCH 11/16] [Fix] Address review feedback on paginated daily activity hook 1. Replace ...args spread in useEffect deps with JSON.stringify(args) key to prevent infinite re-renders when callers pass unstable array references. 2. Add missing agentCancelled partial-data message in EntityUsage so the outer condition no longer renders an empty div. 3. Store setTimeout ID in a ref and clearTimeout on cleanup/cancel to avoid orphaned timers under rapid re-renders. Co-Authored-By: Claude Opus 4.6 --- .../components/EntityUsage/EntityUsage.tsx | 5 +++ .../hooks/usePaginatedDailyActivity.ts | 38 ++++++++++++++++--- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx index c1d5314096c..e075e8b34eb 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx @@ -416,6 +416,11 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti )} + {agentCancelled && entityType === "team" && ( + + Showing partial agent data ({agentProgress.currentPage}/{agentProgress.totalPages} pages loaded) + + )}
)} | null>(null); + + // Keep args in a ref so the effect can always read the latest values + // without needing them in the dependency array. + const argsRef = useRef(args); + argsRef.current = args; + + // Stable serialised key so the effect only re-runs when the arg *values* change. + const argsKey = JSON.stringify(args); const cancel = useCallback(() => { cancelledRef.current = true; setCancelled(true); setIsFetchingMore(false); + if (delayTimerRef.current !== null) { + clearTimeout(delayTimerRef.current); + delayTimerRef.current = null; + } }, []); useEffect(() => { @@ -126,14 +139,24 @@ export function usePaginatedDailyActivity({ const isStale = () => fetchIdRef.current !== currentFetchId || cancelledRef.current; + /** Cancellable delay that clears itself on cleanup. */ + const delay = (ms: number) => + new Promise((resolve) => { + delayTimerRef.current = setTimeout(() => { + delayTimerRef.current = null; + resolve(); + }, ms); + }); + const run = async () => { + const currentArgs = argsRef.current; setLoading(true); setIsFetchingMore(false); setProgress({ currentPage: 1, totalPages: 1 }); try { // Inject page=1 as the 4th argument. - const argsWithPage = [...args.slice(0, 3), 1, ...args.slice(3)]; + const argsWithPage = [...currentArgs.slice(0, 3), 1, ...currentArgs.slice(3)]; const firstPage = await fetchFn(...argsWithPage); if (isStale()) return; @@ -160,13 +183,11 @@ export function usePaginatedDailyActivity({ if (isStale()) return; // Small delay to avoid overwhelming the backend. - await new Promise((resolve) => - setTimeout(resolve, PAGE_FETCH_DELAY_MS), - ); + await delay(PAGE_FETCH_DELAY_MS); if (isStale()) return; - const argsForPage = [...args.slice(0, 3), page, ...args.slice(3)]; + const argsForPage = [...currentArgs.slice(0, 3), page, ...currentArgs.slice(3)]; const pageData = await fetchFn(...argsForPage); if (isStale()) return; @@ -201,9 +222,14 @@ export function usePaginatedDailyActivity({ return () => { fetchIdRef.current++; + if (delayTimerRef.current !== null) { + clearTimeout(delayTimerRef.current); + delayTimerRef.current = null; + } }; + // argsKey is a stable JSON string so the effect only re-fires when arg values change. // eslint-disable-next-line react-hooks/exhaustive-deps - }, [enabled, fetchFn, ...args]); + }, [enabled, fetchFn, argsKey]); return { data, loading, isFetchingMore, progress, cancelled, cancel }; } From f72931a46332d6d53f63605ca584b70e328a9b67 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Mar 2026 11:09:10 -0700 Subject: [PATCH 12/16] [Feature] UI - Usage: Prominent fetch banner, batched pagination renders Replace subtle loading text with antd Alert banners that clearly communicate pagination status, and batch state flushes to reduce chart re-renders. - Replace inline loading text with warning Alert banners showing progress, "open a new tab" link with ExportOutlined icon, and primary Stop button - Batch setData calls every 5 pages instead of per-page to cut re-renders ~80% - Reduce fetch delay from 500ms to 300ms for faster data loading - Add "Charts will update periodically" messaging to set expectations - Fix pre-existing TS error: Button icon prop was using render function instead of ReactNode Co-Authored-By: Claude Opus 4.6 --- .../components/EntityUsage/EntityUsage.tsx | 99 ++- .../components/UsagePageView.test.tsx | 1 + .../UsagePage/components/UsagePageView.tsx | 785 +++++++++--------- .../hooks/usePaginatedDailyActivity.ts | 21 +- 4 files changed, 467 insertions(+), 439 deletions(-) diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx index e075e8b34eb..cb1a08d3ff1 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx @@ -22,7 +22,8 @@ import { Text, Title, } from "@tremor/react"; -import { LoadingOutlined } from "@ant-design/icons"; +import { ExportOutlined } from "@ant-design/icons"; +import { Alert, Button } from "antd"; import React, { useMemo, useState } from "react"; import { ActivityMetrics, processActivityData } from "../../../activity_metrics"; import { UsageExportHeader } from "../../../EntityUsageExport"; @@ -105,8 +106,8 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti const [topModelsLimit, setTopModelsLimit] = useState(5); const [topAgentsLimit, setTopAgentsLimit] = useState(5); - const startTime = useMemo(() => dateValue.from ? new Date(dateValue.from) : null, [dateValue.from]); - const endTime = useMemo(() => dateValue.to ? new Date(dateValue.to) : null, [dateValue.to]); + const startTime = useMemo(() => (dateValue.from ? new Date(dateValue.from) : null), [dateValue.from]); + const endTime = useMemo(() => (dateValue.to ? new Date(dateValue.to) : null), [dateValue.to]); const entityFilterArg = useMemo(() => { if (entityType === "user") return selectedTags.length > 0 ? selectedTags[0] : null; @@ -395,33 +396,75 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti return (
- {(isFetchingMore || cancelled || agentIsFetchingMore || agentCancelled) && ( -
- {isFetchingMore && ( - <> - - Loading spend data... (page {progress.currentPage}/{progress.totalPages}) - - - )} - {cancelled && ( - + {isFetchingMore && ( + + + Currently fetching spend data: fetched {progress.currentPage} / {progress.totalPages} pages. Charts will + update periodically as data loads. Moving off of this page will stop and reset this. To continue using + the UI in the meantime,{" "} + + open a new tab + + . + + +
+ } + /> + )} + {cancelled && ( + Showing partial data ({progress.currentPage}/{progress.totalPages} pages loaded) - )} - {agentIsFetchingMore && entityType === "team" && ( - <> - - Loading agent data... (page {agentProgress.currentPage}/{agentProgress.totalPages}) - - - )} - {agentCancelled && entityType === "team" && ( - + } + /> + )} + {agentIsFetchingMore && entityType === "team" && ( + + + Currently fetching agent data: fetched {agentProgress.currentPage} / {agentProgress.totalPages} pages. + Charts will update periodically as data loads. Moving off of this page will stop and reset this. To + continue using the UI in the meantime,{" "} + + open a new tab + + . + + +
+ } + /> + )} + {agentCancelled && entityType === "team" && ( + Showing partial agent data ({agentProgress.currentPage}/{agentProgress.totalPages} pages loaded) - )} - + } + /> )} = ({ accessToken, entityType, enti - ) : <>} + ) : ( + <> + )} diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx index b9fe1687e6c..bbcddd572cd 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx @@ -251,6 +251,7 @@ vi.mock("@ant-design/icons", async () => { UserOutlined: Icon, DownOutlined: Icon, RightOutlined: Icon, + ExportOutlined: Icon, LoadingOutlined, }; }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 3ebe8205052..8b111c63c90 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -6,7 +6,8 @@ * Works at 1m+ spend logs, by querying an aggregate table instead. */ -import { DownOutlined, InfoCircleOutlined, LoadingOutlined, RightOutlined } from "@ant-design/icons"; +import { DownOutlined, ExportOutlined, InfoCircleOutlined, LoadingOutlined, RightOutlined } from "@ant-design/icons"; +import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; import { BarChart, Card, @@ -19,10 +20,9 @@ import { TabPanel, TabPanels, Text, - Title + Title, } from "@tremor/react"; -import { Alert, Segmented, Select, Tooltip, Typography } from "antd"; -import { useDebouncedState } from "@tanstack/react-pacer/debouncer"; +import { Alert, Button, Segmented, Select, Tooltip, Typography } from "antd"; import React, { useCallback, useEffect, useMemo, useRef, useState, type UIEvent } from "react"; import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; @@ -31,7 +31,6 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import { formatNumberWithCommas } from "@/utils/dataUtils"; -import { Button } from "@tremor/react"; import { all_admin_roles } from "../../../utils/roles"; import { ActivityMetrics, processActivityData } from "../../activity_metrics"; import CloudZeroExportModal from "../../cloudzero_export_modal"; @@ -50,8 +49,8 @@ import EndpointUsage from "./EndpointUsage/EndpointUsage"; import EntityUsage, { EntityList } from "./EntityUsage/EntityUsage"; import SpendByProvider from "./EntityUsage/SpendByProvider"; import TopKeyView from "./EntityUsage/TopKeyView"; -import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect"; import UsageAIChatPanel from "./UsageAIChatPanel"; +import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect"; interface UsagePageProps { teams: Team[]; @@ -128,8 +127,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const handleUserPopupScroll = (e: UIEvent) => { const target = e.currentTarget; - const scrollRatio = - (target.scrollTop + target.clientHeight) / target.scrollHeight; + const scrollRatio = (target.scrollTop + target.clientHeight) / target.scrollHeight; if (scrollRatio >= 0.8 && hasNextUsersPage && !isFetchingNextUsersPage) { fetchNextUsersPage(); } @@ -137,9 +135,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { // For admins: null means global view (all users), a string means filter by that user // For non-admins: always set to their own user ID - const [selectedUserId, setSelectedUserId] = useState( - isAdmin ? null : (userID || null) - ); + const [selectedUserId, setSelectedUserId] = useState(isAdmin ? null : userID || null); const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups"); const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false); const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false); @@ -174,10 +170,10 @@ const UsagePage: React.FC = ({ teams, organizations }) => { }, [isAdmin, userID]); // For non-admins, always pass their own user_id - const effectiveUserId = isAdmin ? selectedUserId : (userID || null); + const effectiveUserId = isAdmin ? selectedUserId : userID || null; - const startTime = useMemo(() => dateValue.from ? new Date(dateValue.from) : null, [dateValue.from]); - const endTime = useMemo(() => dateValue.to ? new Date(dateValue.to) : null, [dateValue.to]); + const startTime = useMemo(() => (dateValue.from ? new Date(dateValue.from) : null), [dateValue.from]); + const endTime = useMemo(() => (dateValue.to ? new Date(dateValue.to) : null), [dateValue.to]); // Try aggregated endpoint first, fall back to paginated on failure const aggregatedFetchIdRef = useRef(0); @@ -361,7 +357,8 @@ const UsagePage: React.FC = ({ teams, organizations }) => { providerSpendMap[provider].metrics.successful_requests += metrics.metrics.successful_requests || 0; providerSpendMap[provider].metrics.failed_requests += metrics.metrics.failed_requests || 0; providerSpendMap[provider].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0; - providerSpendMap[provider].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0; + providerSpendMap[provider].metrics.cache_creation_input_tokens += + metrics.metrics.cache_creation_input_tokens || 0; }); }); @@ -429,430 +426,408 @@ const UsagePage: React.FC = ({ teams, organizations }) => { ); const modelMetrics = useMemo(() => processActivityData(userSpendData, "models", teams), [userSpendData, teams]); const keyMetrics = useMemo(() => processActivityData(userSpendData, "api_keys", teams), [userSpendData, teams]); - const mcpServerMetrics = useMemo(() => processActivityData(userSpendData, "mcp_servers", teams), [userSpendData, teams]); + const mcpServerMetrics = useMemo( + () => processActivityData(userSpendData, "mcp_servers", teams), + [userSpendData, teams], + ); return (
- {/* Export Data Button - Positioned in top right corner */} - {/* {all_admin_roles.includes(userRole || "") && ( -
- -
- )} */} - {/* Global Date Picker and Tabs - Single Row */}
- setUsageView(value)} - isAdmin={isAdmin} - /> + setUsageView(value)} isAdmin={isAdmin} />
- {(paginatedResult.isFetchingMore || paginatedResult.cancelled) && ( -
- {paginatedResult.isFetchingMore && ( - <> - + {paginatedResult.isFetchingMore && ( + - Loading spend data... (page {paginatedResult.progress.currentPage}/{paginatedResult.progress.totalPages}) + Currently fetching spend data: fetched {paginatedResult.progress.currentPage} /{" "} + {paginatedResult.progress.totalPages} pages. Charts will update periodically as data loads. Moving + off of this page will stop and reset this. To continue using the UI in the meantime,{" "} + + open a new tab + + . - - - )} - {paginatedResult.cancelled && ( - - Showing partial data ({paginatedResult.progress.currentPage}/{paginatedResult.progress.totalPages} pages loaded) + +
+ } + /> + )} + {paginatedResult.cancelled && ( + + Showing partial data ({paginatedResult.progress.currentPage}/{paginatedResult.progress.totalPages}{" "} + pages loaded) - )} -
+ } + /> )} {/* Your Usage Panel */} {usageView === "global" && ( <> - {isAdmin && ( -
- Filter by user - setSelectedUserId(value ?? null)} + filterOption={false} + onSearch={handleUserSearchChange} + searchValue={userSearchInput} + onPopupScroll={handleUserPopupScroll} + loading={isLoadingUsers} + notFoundContent={isLoadingUsers ? : "No users found"} + options={userOptions} + popupRender={(menu) => ( + <> + {menu} + {isFetchingNextUsersPage && ( +
+ +
+ )} + )} - > - Ask AI - - + />
-
- - {/* Cost Panel */} - - - {/* Total Spend Card */} - -
- - Project Spend{" "} - {dateValue.from && dateValue.to && ( - <> - {dateValue.from.toLocaleDateString("en-US", { - month: "short", - day: "numeric", - year: dateValue.from.getFullYear() !== dateValue.to.getFullYear() ? "numeric" : undefined, - })} - {" - "} - {dateValue.to.toLocaleDateString("en-US", { - month: "short", - day: "numeric", - year: "numeric", - })} - - )} - -
+ )} + +
+ + Cost + Model Activity + Key Activity + MCP Server Activity + Endpoint Activity + +
+ + +
+
+ + {/* Cost Panel */} + + + {/* Total Spend Card */} + +
+ + Project Spend{" "} + {dateValue.from && dateValue.to && ( + <> + {dateValue.from.toLocaleDateString("en-US", { + month: "short", + day: "numeric", + year: + dateValue.from.getFullYear() !== dateValue.to.getFullYear() ? "numeric" : undefined, + })} + {" - "} + {dateValue.to.toLocaleDateString("en-US", { + month: "short", + day: "numeric", + year: "numeric", + })} + + )} + +
- - + + - - - Usage Metrics - - - Total Requests - - {userSpendData.metadata?.total_api_requests?.toLocaleString() || 0} - - - - Successful Requests - - {userSpendData.metadata?.total_successful_requests?.toLocaleString() || 0} - - - -
- Failed Requests - - - -
- - {userSpendData.metadata?.total_failed_requests?.toLocaleString() || 0} - - - - Average Cost per Request - - $ - {formatNumberWithCommas( - (totalSpend || 0) / (userSpendData.metadata?.total_api_requests || 1), - 4, - )} - - - setShowTokenBreakdown(!showTokenBreakdown)} - > -
- Total Tokens - {showTokenBreakdown ? ( - - ) : ( - - )} -
- - {userSpendData.metadata?.total_tokens?.toLocaleString() || 0} - - - - {showTokenBreakdown && ( - + + + Usage Metrics + - Input Tokens - - {userSpendData.metadata?.total_prompt_tokens?.toLocaleString() || 0} + Total Requests + + {userSpendData.metadata?.total_api_requests?.toLocaleString() || 0} - Output Tokens - - {userSpendData.metadata?.total_completion_tokens?.toLocaleString() || 0} - - - - Cache Read Tokens + Successful Requests - {userSpendData.metadata?.total_cache_read_input_tokens?.toLocaleString() || 0} + {userSpendData.metadata?.total_successful_requests?.toLocaleString() || 0} - Cache Write Tokens - - {userSpendData.metadata?.total_cache_creation_input_tokens?.toLocaleString() || 0} +
+ Failed Requests + + + +
+ + {userSpendData.metadata?.total_failed_requests?.toLocaleString() || 0} + + + + Average Cost per Request + + $ + {formatNumberWithCommas( + (totalSpend || 0) / (userSpendData.metadata?.total_api_requests || 1), + 4, + )} + + + setShowTokenBreakdown(!showTokenBreakdown)} + > +
+ Total Tokens + {showTokenBreakdown ? ( + + ) : ( + + )} +
+ + {userSpendData.metadata?.total_tokens?.toLocaleString() || 0} - )} - - + {showTokenBreakdown && ( + + + Input Tokens + + {userSpendData.metadata?.total_prompt_tokens?.toLocaleString() || 0} + + + + Output Tokens + + {userSpendData.metadata?.total_completion_tokens?.toLocaleString() || 0} + + + + Cache Read Tokens + + {userSpendData.metadata?.total_cache_read_input_tokens?.toLocaleString() || 0} + + + + Cache Write Tokens + + {userSpendData.metadata?.total_cache_creation_input_tokens?.toLocaleString() || 0} + + + + )} + + - {/* Daily Spend Chart */} - - - Daily Spend - {loading ? ( - - ) : ( - { - if (!active || !payload?.[0]) return null; - const data = payload[0].payload; - return ( -
-

{data.date}

-

- Spend: ${formatNumberWithCommas(data.metrics.spend, 2)} -

-

Requests: {data.metrics.api_requests}

-

Successful: {data.metrics.successful_requests}

-

Failed: {data.metrics.failed_requests}

-

Tokens: {data.metrics.total_tokens}

-
- ); - }} + {/* Daily Spend Chart */} + + + Daily Spend + {loading ? ( + + ) : ( + { + if (!active || !payload?.[0]) return null; + const data = payload[0].payload; + return ( +
+

{data.date}

+

+ Spend: ${formatNumberWithCommas(data.metrics.spend, 2)} +

+

Requests: {data.metrics.api_requests}

+

Successful: {data.metrics.successful_requests}

+

Failed: {data.metrics.failed_requests}

+

Tokens: {data.metrics.total_tokens}

+
+ ); + }} + /> + )} + + + {/* Top API Keys */} + + + Top Virtual Keys + - )} - - - {/* Top API Keys */} - - - Top Virtual Keys - + + + {/* Top Models */} + + + {modelViewType === "groups" ? "Top Public Model Names" : "Top Litellm Models"} +
+ setTopModelsLimit(value as number)} + /> +
+ + +
+
+ {loading ? ( + + ) : ( +
+ {(() => { + const modelData = modelViewType === "groups" ? topModelGroups : topModels; + return ( + { + if (!active || !payload?.[0]) return null; + const data = payload[0].payload; + return ( +
+

{data.key}

+

+ Spend: ${formatNumberWithCommas(data.spend, 2)} +

+

+ Total Requests: {data.requests.toLocaleString()} +

+

+ Successful: {data.successful_requests.toLocaleString()} +

+

+ Failed: {data.failed_requests.toLocaleString()} +

+

Tokens: {data.tokens.toLocaleString()}

+
+ ); + }} + /> + ); + })()} +
+ )} + + + + {/* Spend by Provider */} + + - - + - {/* Top Models */} - - - {modelViewType === "groups" ? "Top Public Model Names" : "Top Litellm Models"} -
- setTopModelsLimit(value as number)} - /> -
- - -
-
- {loading ? ( - - ) : ( -
- {(() => { - const modelData = - modelViewType === "groups" - ? topModelGroups - : topModels; - return ( - { - if (!active || !payload?.[0]) return null; - const data = payload[0].payload; - return ( -
-

{data.key}

-

Spend: ${formatNumberWithCommas(data.spend, 2)}

-

- Total Requests: {data.requests.toLocaleString()} -

-

- Successful: {data.successful_requests.toLocaleString()} -

-

Failed: {data.failed_requests.toLocaleString()}

-

Tokens: {data.tokens.toLocaleString()}

-
- ); - }} - /> - ); - })()} -
- )} - - + {/* Usage Metrics */} + + - {/* Spend by Provider */} - - - - - {/* Usage Metrics */} - - - - {/* Activity Panel */} - - - - - - - - - - - - - - + {/* Activity Panel */} + + + + + + + + + + + + + + )} {/* Organization Usage Panel */} @@ -994,11 +969,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { /> {/* AI Chat Panel */} - setIsAiChatOpen(false)} - accessToken={accessToken} - /> + setIsAiChatOpen(false)} accessToken={accessToken} />
); }; diff --git a/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts b/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts index 498c4c07a2b..1716c33ed2d 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts @@ -7,7 +7,10 @@ export interface PaginationProgress { } /** Delay between sequential page fetches (ms) to avoid overloading the backend. */ -const PAGE_FETCH_DELAY_MS = 500; +const PAGE_FETCH_DELAY_MS = 300; + +/** Number of pages to accumulate before flushing to React state (reduces re-renders). */ +const RENDER_BATCH_SIZE = 5; /** The metadata fields returned by the daily activity API that should be summed across pages. */ const SUMMABLE_METADATA_KEYS = [ @@ -201,11 +204,19 @@ export function usePaginatedDailyActivity({ accumulatedMetadata.has_more = page < totalPages; accumulatedMetadata.page = page; - setData({ - results: accumulatedResults, - metadata: accumulatedMetadata, - }); + // Always update progress so the banner stays responsive. setProgress({ currentPage: page, totalPages }); + + // Flush accumulated data to React state every RENDER_BATCH_SIZE pages + // (or on the final page) to avoid expensive per-page re-renders. + const isLastPage = page === totalPages; + const isBatchBoundary = (page - 1) % RENDER_BATCH_SIZE === 0; + if (isLastPage || isBatchBoundary) { + setData({ + results: accumulatedResults, + metadata: accumulatedMetadata, + }); + } } setIsFetchingMore(false); From d26faeb844bfcce465bf84c2b749565cb255a255 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Mar 2026 11:28:15 -0700 Subject: [PATCH 13/16] [Fix] UI - Usage: Reduce batch size to 3, add loading spinner to fetch banner - Reduce RENDER_BATCH_SIZE from 5 to 3 for more frequent chart updates - Add LoadingOutlined spinner at the start of all fetching Alert banners Co-Authored-By: Claude Opus 4.6 --- .../components/EntityUsage/EntityUsage.tsx | 4 +++- .../UsagePage/components/UsagePageView.tsx | 1 + .../UsagePage/hooks/usePaginatedDailyActivity.ts | 16 ++++++++-------- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx index cb1a08d3ff1..aaeb8ebb4be 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx @@ -22,7 +22,7 @@ import { Text, Title, } from "@tremor/react"; -import { ExportOutlined } from "@ant-design/icons"; +import { ExportOutlined, LoadingOutlined } from "@ant-design/icons"; import { Alert, Button } from "antd"; import React, { useMemo, useState } from "react"; import { ActivityMetrics, processActivityData } from "../../../activity_metrics"; @@ -404,6 +404,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti message={
+ Currently fetching spend data: fetched {progress.currentPage} / {progress.totalPages} pages. Charts will update periodically as data loads. Moving off of this page will stop and reset this. To continue using the UI in the meantime,{" "} @@ -439,6 +440,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti message={
+ Currently fetching agent data: fetched {agentProgress.currentPage} / {agentProgress.totalPages} pages. Charts will update periodically as data loads. Moving off of this page will stop and reset this. To continue using the UI in the meantime,{" "} diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 8b111c63c90..1495c7d3e5b 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -448,6 +448,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { message={
+ Currently fetching spend data: fetched {paginatedResult.progress.currentPage} /{" "} {paginatedResult.progress.totalPages} pages. Charts will update periodically as data loads. Moving off of this page will stop and reset this. To continue using the UI in the meantime,{" "} diff --git a/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts b/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts index 1716c33ed2d..9a7ab22c9af 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/hooks/usePaginatedDailyActivity.ts @@ -10,7 +10,7 @@ export interface PaginationProgress { const PAGE_FETCH_DELAY_MS = 300; /** Number of pages to accumulate before flushing to React state (reduces re-renders). */ -const RENDER_BATCH_SIZE = 5; +const RENDER_BATCH_SIZE = 3; /** The metadata fields returned by the daily activity API that should be summed across pages. */ const SUMMABLE_METADATA_KEYS = [ @@ -80,8 +80,8 @@ function sumMetadata( } /** - * Hook that auto-paginates daily activity endpoints, updating state after each - * page so charts render progressively. Cancels on unmount, param changes, or + * Hook that auto-paginates daily activity endpoints, updating state in batches + * so charts render progressively. Cancels on unmount, param changes, or * manual cancel(). * * The `args` array should contain every argument the fetchFn expects EXCEPT @@ -204,11 +204,10 @@ export function usePaginatedDailyActivity({ accumulatedMetadata.has_more = page < totalPages; accumulatedMetadata.page = page; - // Always update progress so the banner stays responsive. - setProgress({ currentPage: page, totalPages }); - - // Flush accumulated data to React state every RENDER_BATCH_SIZE pages - // (or on the final page) to avoid expensive per-page re-renders. + // Flush accumulated data and progress to React state every + // RENDER_BATCH_SIZE pages (or on the final page) to avoid + // expensive per-page re-renders. Progress and data are updated + // together so the counter never appears to decrement. const isLastPage = page === totalPages; const isBatchBoundary = (page - 1) % RENDER_BATCH_SIZE === 0; if (isLastPage || isBatchBoundary) { @@ -216,6 +215,7 @@ export function usePaginatedDailyActivity({ results: accumulatedResults, metadata: accumulatedMetadata, }); + setProgress({ currentPage: page, totalPages }); } } From 374c3458d51ed6fb5c82400a55d6c38d997e101e Mon Sep 17 00:00:00 2001 From: ryanh-ai <3118399+ryanh-ai@users.noreply.github.com> Date: Sat, 14 Mar 2026 15:10:01 -0700 Subject: [PATCH 14/16] feat: add sagemaker_nova provider for Amazon Nova models on SageMaker (#21542) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add sagemaker_nova provider for Nova models on SageMaker Add support for custom/fine-tuned Amazon Nova models (Nova Micro, Nova Lite, Nova 2 Lite) deployed on SageMaker Inference real-time endpoints. Nova uses OpenAI-compatible request/response format with additional Nova-specific parameters (top_k, reasoning_effort, allowed_token_ids, truncate_prompt_tokens) and requires stream:true in the request body. Nova endpoints also reject 'model' in the request body. Changes: - New provider: sagemaker_nova/ - SagemakerNovaConfig inherits from SagemakerChatConfig - Override transform_request to strip 'model' from request body - Override supports_stream_param_in_request_body (True for Nova) - Extend get_supported_openai_params with Nova-specific params - Refactored SagemakerChatConfig to use custom_llm_provider param instead of hardcoded strings (backwards-compatible) - Consolidated main.py routing for sagemaker_chat and sagemaker_nova - 22 unit tests + 9 integration tests (skip-gated) - Documentation with SDK, streaming, multimodal, and proxy examples - All tests verified against live SageMaker Nova endpoint * fix: move integration tests to tests/local_testing/ per test directory policy * fix: remove unused module-level SagemakerNovaConfig instance The sagemaker_nova_config singleton was never imported or used — the ProviderConfigManager creates its own instance via the lambda registered in utils.py. Removing this leftover boilerplate. --------- Co-authored-by: Krish Dholakia --- .../docs/providers/aws_sagemaker.md | 95 +++++ litellm/__init__.py | 9 +- litellm/_lazy_imports_registry.py | 5 + litellm/constants.py | 1 + litellm/llms/sagemaker/chat/transformation.py | 10 +- litellm/llms/sagemaker/nova/__init__.py | 1 + litellm/llms/sagemaker/nova/transformation.py | 70 ++++ litellm/main.py | 6 +- litellm/types/utils.py | 1 + litellm/utils.py | 1 + .../test_sagemaker_nova_integration.py | 276 ++++++++++++ .../test_sagemaker_nova_transformation.py | 393 ++++++++++++++++++ 12 files changed, 857 insertions(+), 11 deletions(-) create mode 100644 litellm/llms/sagemaker/nova/__init__.py create mode 100644 litellm/llms/sagemaker/nova/transformation.py create mode 100644 tests/local_testing/test_sagemaker_nova_integration.py create mode 100644 tests/test_litellm/llms/sagemaker/test_sagemaker_nova_transformation.py diff --git a/docs/my-website/docs/providers/aws_sagemaker.md b/docs/my-website/docs/providers/aws_sagemaker.md index bab475e7305..a2440c73d7d 100644 --- a/docs/my-website/docs/providers/aws_sagemaker.md +++ b/docs/my-website/docs/providers/aws_sagemaker.md @@ -526,3 +526,98 @@ print(f"response: {response}") ``` + +## Nova Models on SageMaker + +LiteLLM supports Amazon Nova models (Nova Micro, Nova Lite, Nova 2 Lite) deployed on SageMaker Inference real-time endpoints. These custom/fine-tuned Nova models use an OpenAI-compatible API format. + +**Reference:** [AWS Blog - Amazon SageMaker Inference for Custom Amazon Nova Models](https://aws.amazon.com/blogs/aws/announcing-amazon-sagemaker-inference-for-custom-amazon-nova-models/) + +### Usage + +Use the `sagemaker_nova/` prefix with your SageMaker endpoint name: + +```python +import litellm +import os + +os.environ["AWS_ACCESS_KEY_ID"] = "" +os.environ["AWS_SECRET_ACCESS_KEY"] = "" +os.environ["AWS_REGION_NAME"] = "us-east-1" + +# Basic chat completion +response = litellm.completion( + model="sagemaker_nova/my-nova-endpoint", + messages=[{"role": "user", "content": "Hello, how are you?"}], + temperature=0.7, + max_tokens=512, +) +print(response.choices[0].message.content) +``` + +### Streaming + +```python +response = litellm.completion( + model="sagemaker_nova/my-nova-endpoint", + messages=[{"role": "user", "content": "Write a short poem"}], + stream=True, + stream_options={"include_usage": True}, +) +for chunk in response: + if chunk.choices[0].delta.content: + print(chunk.choices[0].delta.content, end="") +``` + +### Multimodal (Images) + +Nova models on SageMaker support image inputs using base64 data URIs: + +```python +response = litellm.completion( + model="sagemaker_nova/my-nova-endpoint", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}} + ] + } + ], +) +``` + +### Proxy Config + +```yaml +model_list: + - model_name: nova-micro + litellm_params: + model: sagemaker_nova/my-nova-micro-endpoint + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_region_name: us-east-1 +``` + +### Supported Parameters + +All standard OpenAI parameters are supported, plus these Nova-specific parameters: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `top_k` | integer | Limits token selection to top K most likely tokens | +| `reasoning_effort` | `"low"` \| `"high"` | Reasoning effort level (Nova 2 Lite custom models only) | +| `allowed_token_ids` | array[int] | Restrict output to specified token IDs | +| `truncate_prompt_tokens` | integer | Truncate prompt to N tokens if it exceeds limit | + +```python +response = litellm.completion( + model="sagemaker_nova/my-nova-endpoint", + messages=[{"role": "user", "content": "Think step by step: what is 2+2?"}], + top_k=40, + reasoning_effort="low", + logprobs=True, + top_logprobs=2, +) +``` diff --git a/litellm/__init__.py b/litellm/__init__.py index dcd4ce29096..b75630c19d8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1464,12 +1464,9 @@ if TYPE_CHECKING: from .llms.petals.completion.transformation import PetalsConfig as PetalsConfig from .llms.ollama.chat.transformation import OllamaChatConfig as OllamaChatConfig from .llms.ollama.completion.transformation import OllamaConfig as OllamaConfig - from .llms.sagemaker.completion.transformation import ( - SagemakerConfig as SagemakerConfig, - ) - from .llms.sagemaker.chat.transformation import ( - SagemakerChatConfig as SagemakerChatConfig, - ) + from .llms.sagemaker.completion.transformation import SagemakerConfig as SagemakerConfig + from .llms.sagemaker.chat.transformation import SagemakerChatConfig as SagemakerChatConfig + from .llms.sagemaker.nova.transformation import SagemakerNovaConfig as SagemakerNovaConfig from .llms.cohere.chat.transformation import CohereChatConfig as CohereChatConfig from .llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig as AnthropicMessagesConfig, diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index f7f56d2b889..9164a3c8ae4 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -167,6 +167,7 @@ LLM_CONFIG_NAMES = ( "OllamaConfig", "SagemakerConfig", "SagemakerChatConfig", + "SagemakerNovaConfig", "CohereChatConfig", "AnthropicMessagesConfig", "AmazonAnthropicClaudeMessagesConfig", @@ -701,6 +702,10 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.sagemaker.chat.transformation", "SagemakerChatConfig", ), + "SagemakerNovaConfig": ( + ".llms.sagemaker.nova.transformation", + "SagemakerNovaConfig", + ), "CohereChatConfig": (".llms.cohere.chat.transformation", "CohereChatConfig"), "AnthropicMessagesConfig": ( ".llms.anthropic.experimental_pass_through.messages.transformation", diff --git a/litellm/constants.py b/litellm/constants.py index 0cf9700cddd..89c59ee9326 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -505,6 +505,7 @@ LITELLM_CHAT_PROVIDERS = [ "azure_ai", "sagemaker", "sagemaker_chat", + "sagemaker_nova", "bedrock", "vllm", "nlp_cloud", diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 2b458fbc438..60e85c9f93b 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -160,7 +160,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, - custom_llm_provider="sagemaker_chat", + custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) return streaming_response @@ -180,8 +180,12 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): signed_json_body: Optional[bytes] = None, ) -> CustomStreamWrapper: if client is None or isinstance(client, HTTPHandler): + try: + llm_provider = LlmProviders(custom_llm_provider) + except ValueError: + llm_provider = LlmProviders.SAGEMAKER_CHAT client = get_async_httpx_client( - llm_provider=LlmProviders.SAGEMAKER_CHAT, params={} + llm_provider=llm_provider, params={} ) try: @@ -210,7 +214,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, - custom_llm_provider="sagemaker_chat", + custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) return streaming_response diff --git a/litellm/llms/sagemaker/nova/__init__.py b/litellm/llms/sagemaker/nova/__init__.py new file mode 100644 index 00000000000..fdebd0b0e41 --- /dev/null +++ b/litellm/llms/sagemaker/nova/__init__.py @@ -0,0 +1 @@ +from .transformation import SagemakerNovaConfig # noqa: F401 diff --git a/litellm/llms/sagemaker/nova/transformation.py b/litellm/llms/sagemaker/nova/transformation.py new file mode 100644 index 00000000000..41c20847b53 --- /dev/null +++ b/litellm/llms/sagemaker/nova/transformation.py @@ -0,0 +1,70 @@ +""" +Translate from OpenAI's `/v1/chat/completions` to SageMaker Nova Inference endpoints. + +Nova models on SageMaker use OpenAI-compatible request/response format with +additional Nova-specific parameters (top_k, reasoning_effort, etc.). + +Docs: https://docs.aws.amazon.com/nova/latest/nova2-userguide/nova-sagemaker-inference-api-reference.html +""" + +from typing import List + +from litellm.types.llms.openai import AllMessageValues + +from ..chat.transformation import SagemakerChatConfig + + +class SagemakerNovaConfig(SagemakerChatConfig): + """ + Config for Amazon Nova models deployed on SageMaker Inference endpoints. + + Nova uses OpenAI-compatible format (same as sagemaker_chat / HF Messages API) + but with additional Nova-specific parameters and requires `stream: true` in + the request body for streaming. + + Usage: + model="sagemaker_nova/" + """ + + @property + def supports_stream_param_in_request_body(self) -> bool: + """Nova expects `stream: true` in the request body for streaming.""" + return True + + def get_supported_openai_params(self, model: str) -> List: + """Extend parent params with Nova-specific parameters.""" + params = super().get_supported_openai_params(model) + nova_params = [ + "top_k", + "reasoning_effort", + "allowed_token_ids", + "truncate_prompt_tokens", + ] + for p in nova_params: + if p not in params: + params.append(p) + return params + + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Nova SageMaker endpoints do not accept 'model' in the request body. + Only supported fields: messages, max_tokens, max_completion_tokens, + temperature, top_p, top_k, stream, stream_options, logprobs, + top_logprobs, reasoning_effort, allowed_token_ids, truncate_prompt_tokens. + """ + request_body = super().transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + request_body.pop("model", None) + return request_body diff --git a/litellm/main.py b/litellm/main.py index 5a1b7f0fe16..986ca4527f2 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3712,8 +3712,10 @@ def completion( # type: ignore # noqa: PLR0915 ): return _model_response response = _model_response - elif custom_llm_provider == "sagemaker_chat": + elif custom_llm_provider in ("sagemaker_chat", "sagemaker_nova"): # boto3 reads keys from .env + # sagemaker_chat: HF Messages API endpoints + # sagemaker_nova: Nova models on SageMaker (OpenAI-compatible) model_response = base_llm_http_handler.completion( model=model, stream=stream, @@ -3723,7 +3725,7 @@ def completion( # type: ignore # noqa: PLR0915 model_response=model_response, optional_params=optional_params, litellm_params=litellm_params, - custom_llm_provider="sagemaker_chat", + custom_llm_provider=custom_llm_provider, timeout=timeout, headers=headers, encoding=_get_encoding(), diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f9a4f1429b6..892c9578b94 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3131,6 +3131,7 @@ class LlmProviders(str, Enum): AZURE_AI = "azure_ai" SAGEMAKER = "sagemaker" SAGEMAKER_CHAT = "sagemaker_chat" + SAGEMAKER_NOVA = "sagemaker_nova" BEDROCK = "bedrock" VLLM = "vllm" NLP_CLOUD = "nlp_cloud" diff --git a/litellm/utils.py b/litellm/utils.py index b4ed37f9a78..81d749ab821 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7945,6 +7945,7 @@ class ProviderConfigManager: LlmProviders.VERTEX_AI_BETA: (lambda: litellm.VertexGeminiConfig(), False), LlmProviders.CLOUDFLARE: (lambda: litellm.CloudflareChatConfig(), False), LlmProviders.SAGEMAKER_CHAT: (lambda: litellm.SagemakerChatConfig(), False), + LlmProviders.SAGEMAKER_NOVA: (lambda: litellm.SagemakerNovaConfig(), False), LlmProviders.SAGEMAKER: (lambda: litellm.SagemakerConfig(), False), LlmProviders.FIREWORKS_AI: (lambda: litellm.FireworksAIConfig(), False), LlmProviders.FRIENDLIAI: (lambda: litellm.FriendliaiChatConfig(), False), diff --git a/tests/local_testing/test_sagemaker_nova_integration.py b/tests/local_testing/test_sagemaker_nova_integration.py new file mode 100644 index 00000000000..6f55bea38b9 --- /dev/null +++ b/tests/local_testing/test_sagemaker_nova_integration.py @@ -0,0 +1,276 @@ +""" +Integration tests for SageMaker Nova provider. + +These tests require a live SageMaker Nova endpoint and AWS credentials. +They are skipped by default — run manually with: + + pytest tests/test_litellm/llms/sagemaker/test_sagemaker_nova_integration.py -v --no-header -rN + +Prerequisites: + export AWS_PROFILE= # or set AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY + export AWS_REGION_NAME=us-east-1 + export SAGEMAKER_NOVA_ENDPOINT= +""" + +import base64 +import io +import json +import os +import struct +import zlib + +import pytest + +import litellm + +ENDPOINT = os.environ.get("SAGEMAKER_NOVA_ENDPOINT", "") +MODEL = f"sagemaker_nova/{ENDPOINT}" + +skip_if_no_endpoint = pytest.mark.skipif( + not ENDPOINT, + reason="SAGEMAKER_NOVA_ENDPOINT not set — skipping live integration tests", +) + + +def _make_test_png() -> str: + """Create a minimal 4x4 PNG (red border, blue center) and return base64.""" + + def chunk(ctype, data): + c = ctype + data + return ( + struct.pack(">I", len(data)) + + c + + struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF) + ) + + width, height = 4, 4 + pixels = [] + for y in range(height): + for x in range(width): + if 1 <= x <= 2 and 1 <= y <= 2: + pixels.append((0, 0, 255)) + else: + pixels.append((255, 0, 0)) + + raw = b"" + for y in range(height): + raw += b"\x00" + for x in range(width): + raw += bytes(pixels[y * width + x]) + + png = ( + b"\x89PNG\r\n\x1a\n" + + chunk( + b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0) + ) + + chunk(b"IDAT", zlib.compress(raw)) + + chunk(b"IEND", b"") + ) + return base64.b64encode(png).decode() + + +@skip_if_no_endpoint +class TestSagemakerNovaIntegration: + """Live integration tests for sagemaker_nova provider.""" + + def test_should_complete_basic_single_turn(self): + """Basic single-turn chat completion.""" + response = litellm.completion( + model=MODEL, + messages=[{"role": "user", "content": "What is 2+2? Reply in one word."}], + max_tokens=32, + temperature=0.1, + ) + assert response.choices[0].message.content is not None + assert len(response.choices[0].message.content.strip()) > 0 + assert response.choices[0].finish_reason == "stop" + assert response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens > 0 + assert response.usage.total_tokens == ( + response.usage.prompt_tokens + response.usage.completion_tokens + ) + + def test_should_complete_multi_turn_conversation(self): + """Multi-turn conversation maintains context.""" + messages = [ + {"role": "user", "content": "My name is Alice."}, + ] + response1 = litellm.completion( + model=MODEL, + messages=messages, + max_tokens=64, + temperature=0.1, + ) + assistant_msg = response1.choices[0].message.content + assert assistant_msg is not None + + # Second turn — model should remember the name + messages.append({"role": "assistant", "content": assistant_msg}) + messages.append({"role": "user", "content": "What is my name?"}) + + response2 = litellm.completion( + model=MODEL, + messages=messages, + max_tokens=64, + temperature=0.1, + ) + answer = response2.choices[0].message.content.lower() + assert "alice" in answer, f"Expected 'alice' in response, got: {answer}" + + def test_should_stream_response(self): + """Streaming returns chunks with content and final usage.""" + response = litellm.completion( + model=MODEL, + messages=[{"role": "user", "content": "Count from 1 to 5."}], + max_tokens=64, + stream=True, + stream_options={"include_usage": True}, + ) + + chunks = [] + full_content = "" + for chunk in response: + chunks.append(chunk) + delta = chunk.choices[0].delta.content or "" + full_content += delta + + assert len(chunks) > 1, "Expected multiple streaming chunks" + assert len(full_content.strip()) > 0, "Expected non-empty streamed content" + + # Last chunk should have finish_reason + final_chunks_with_finish = [ + c for c in chunks if c.choices and c.choices[0].finish_reason is not None + ] + assert len(final_chunks_with_finish) > 0, "Expected at least one chunk with finish_reason" + + def test_should_return_logprobs(self): + """Logprobs are returned when requested.""" + response = litellm.completion( + model=MODEL, + messages=[{"role": "user", "content": "Say hello."}], + max_tokens=16, + temperature=0.1, + logprobs=True, + top_logprobs=3, + ) + lp = response.choices[0].logprobs + assert lp is not None, "Expected logprobs in response" + + content = lp.content if hasattr(lp, "content") else lp.get("content") + assert content is not None and len(content) > 0, "Expected logprobs content" + + first_token = content[0] + assert "token" in first_token or hasattr(first_token, "token") + assert "logprob" in first_token or hasattr(first_token, "logprob") + + top = first_token.get("top_logprobs") if isinstance(first_token, dict) else first_token.top_logprobs + assert top is not None and len(top) == 3, "Expected 3 top_logprobs" + + def test_should_handle_multimodal_image_input(self): + """Multimodal with base64 image in content array.""" + b64_image = _make_test_png() + response = litellm.completion( + model=MODEL, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What colors do you see in this image? List them.", + }, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{b64_image}" + }, + }, + ], + } + ], + max_tokens=128, + ) + content = response.choices[0].message.content.lower() + assert response.choices[0].message.content is not None + assert len(content) > 0 + # The image has red and blue — model should mention at least one + assert "red" in content or "blue" in content, ( + f"Expected 'red' or 'blue' in multimodal response, got: {content}" + ) + + def test_should_pass_nova_specific_params(self): + """Nova-specific parameters (top_k) are accepted.""" + response = litellm.completion( + model=MODEL, + messages=[{"role": "user", "content": "Say hello."}], + max_tokens=32, + top_k=40, + temperature=0.7, + ) + assert response.choices[0].message.content is not None + assert response.usage.total_tokens > 0 + + def test_should_respect_system_message(self): + """System message should influence the response.""" + response = litellm.completion( + model=MODEL, + messages=[ + { + "role": "system", + "content": "You are a pirate. Always respond in pirate speak.", + }, + {"role": "user", "content": "How are you today?"}, + ], + max_tokens=128, + temperature=0.7, + ) + content = response.choices[0].message.content.lower() + assert response.choices[0].message.content is not None + # Pirate-themed words likely in response + pirate_words = ["arr", "ahoy", "matey", "ye", "sail", "sea", "cap"] + assert any( + w in content for w in pirate_words + ), f"Expected pirate speak, got: {content}" + + +NOVA2_ENDPOINT = os.environ.get("SAGEMAKER_NOVA2_LITE_ENDPOINT", "") +NOVA2_MODEL = f"sagemaker_nova/{NOVA2_ENDPOINT}" + +skip_if_no_nova2_endpoint = pytest.mark.skipif( + not NOVA2_ENDPOINT, + reason="SAGEMAKER_NOVA2_LITE_ENDPOINT not set — requires Nova 2 Lite endpoint", +) + + +@skip_if_no_nova2_endpoint +class TestSagemakerNova2LiteIntegration: + """ + Integration tests requiring a Nova 2 Lite endpoint (reasoning_effort support). + + Run with: + export SAGEMAKER_NOVA2_LITE_ENDPOINT= + pytest tests/test_litellm/llms/sagemaker/test_sagemaker_nova_integration.py::TestSagemakerNova2LiteIntegration -v + """ + + def test_should_accept_reasoning_effort_low(self): + """reasoning_effort='low' should be accepted by Nova 2 Lite.""" + response = litellm.completion( + model=NOVA2_MODEL, + messages=[{"role": "user", "content": "What is 2+2?"}], + max_tokens=32, + reasoning_effort="low", + ) + assert response.choices[0].message.content is not None + assert response.usage.total_tokens > 0 + + def test_should_accept_reasoning_effort_high(self): + """reasoning_effort='high' should be accepted by Nova 2 Lite.""" + response = litellm.completion( + model=NOVA2_MODEL, + messages=[{"role": "user", "content": "Explain why the sky is blue."}], + max_tokens=256, + reasoning_effort="high", + ) + assert response.choices[0].message.content is not None + assert len(response.choices[0].message.content) > 0 + assert response.usage.completion_tokens > 0 diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_nova_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_nova_transformation.py new file mode 100644 index 00000000000..8c468a1ff66 --- /dev/null +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_nova_transformation.py @@ -0,0 +1,393 @@ +""" +Unit tests for SageMaker Nova transformation config. +""" + +import json +import pytest + +from litellm.llms.sagemaker.nova.transformation import SagemakerNovaConfig +from litellm.types.utils import ModelResponse +from litellm.utils import convert_to_model_response_object + + +class TestSagemakerNovaConfig: + def setup_method(self): + self.config = SagemakerNovaConfig() + + def test_should_support_stream_param_in_request_body(self): + """Nova requires stream: true in the request body.""" + assert self.config.supports_stream_param_in_request_body is True + + def test_should_include_nova_specific_params(self): + """Nova-specific params should be in the supported params list.""" + params = self.config.get_supported_openai_params(model="my-nova-endpoint") + assert "top_k" in params + assert "reasoning_effort" in params + assert "allowed_token_ids" in params + assert "truncate_prompt_tokens" in params + + def test_should_include_standard_openai_params(self): + """Standard OpenAI params from parent should still be present.""" + params = self.config.get_supported_openai_params(model="my-nova-endpoint") + assert "temperature" in params + assert "max_tokens" in params + assert "top_p" in params + assert "stream" in params + assert "logprobs" in params + assert "top_logprobs" in params + assert "stream_options" in params + + def test_should_map_nova_params_to_request(self): + """Nova-specific params should pass through to optional_params.""" + optional_params = self.config.map_openai_params( + non_default_params={ + "top_k": 40, + "reasoning_effort": "low", + "temperature": 0.7, + }, + optional_params={}, + model="my-nova-endpoint", + drop_params=False, + ) + assert optional_params["top_k"] == 40 + assert optional_params["reasoning_effort"] == "low" + assert optional_params["temperature"] == 0.7 + + def test_should_generate_correct_url_non_streaming(self): + """Non-streaming URL should use /invocations.""" + url = self.config.get_complete_url( + api_base=None, + api_key=None, + model="my-nova-endpoint", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + stream=False, + ) + assert url == "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/my-nova-endpoint/invocations" + + def test_should_generate_correct_url_streaming(self): + """Streaming URL should use /invocations-response-stream.""" + url = self.config.get_complete_url( + api_base=None, + api_key=None, + model="my-nova-endpoint", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + stream=True, + ) + assert url == "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/my-nova-endpoint/invocations-response-stream" + + def test_should_have_custom_stream_wrapper(self): + """Nova should use custom stream wrapper (AWS EventStream).""" + assert self.config.has_custom_stream_wrapper is True + + +class TestSagemakerNovaResponseParsing: + """Test that Nova's OpenAI-compatible responses are correctly parsed.""" + + def test_should_parse_non_streaming_response(self): + """Nova non-streaming response should be parsed into ModelResponse.""" + nova_response = { + "id": "chatcmpl-123e4567-e89b-12d3-a456-426614174000", + "object": "chat.completion", + "created": 1677652288, + "model": "nova-micro-custom", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! How can I help?", + "refusal": None, + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + } + result = convert_to_model_response_object( + response_object=nova_response, + model_response_object=ModelResponse(), + ) + assert result.id == "chatcmpl-123e4567-e89b-12d3-a456-426614174000" + assert result.choices[0].message.content == "Hello! How can I help?" + assert result.choices[0].finish_reason == "stop" + assert result.usage.prompt_tokens == 9 + assert result.usage.completion_tokens == 12 + assert result.usage.total_tokens == 21 + + def test_should_parse_response_with_reasoning_content(self): + """Nova reasoning_content should be extracted correctly.""" + nova_response = { + "id": "chatcmpl-reasoning-test", + "object": "chat.completion", + "created": 1677652288, + "model": "nova-2-lite-custom", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "The answer is 4.", + "reasoning_content": "Let me think: 2+2=4", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 15, + "completion_tokens": 20, + "total_tokens": 35, + }, + } + result = convert_to_model_response_object( + response_object=nova_response, + model_response_object=ModelResponse(), + ) + assert result.choices[0].message.content == "The answer is 4." + assert result.choices[0].message.reasoning_content == "Let me think: 2+2=4" + + def test_should_parse_response_with_logprobs(self): + """Nova logprobs should be preserved in response.""" + nova_response = { + "id": "chatcmpl-logprobs-test", + "object": "chat.completion", + "created": 1677652288, + "model": "nova-micro-custom", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello", + }, + "logprobs": { + "content": [ + { + "token": "Hello", + "logprob": -0.5, + "top_logprobs": [ + {"token": "Hello", "logprob": -0.5}, + {"token": "Hi", "logprob": -1.2}, + ], + } + ] + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + }, + } + result = convert_to_model_response_object( + response_object=nova_response, + model_response_object=ModelResponse(), + ) + assert result.choices[0].logprobs is not None + assert result.choices[0].logprobs["content"][0]["token"] == "Hello" + + def test_should_parse_response_with_cached_tokens(self): + """Nova prompt_tokens_details with cached_tokens should be parsed.""" + nova_response = { + "id": "chatcmpl-cached-test", + "object": "chat.completion", + "created": 1677652288, + "model": "nova-micro-custom", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hi", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 20, + "completion_tokens": 1, + "total_tokens": 21, + "prompt_tokens_details": {"cached_tokens": 10}, + }, + } + result = convert_to_model_response_object( + response_object=nova_response, + model_response_object=ModelResponse(), + ) + assert result.usage.prompt_tokens_details.cached_tokens == 10 + + +class TestSagemakerChatBackwardsCompatibility: + """Verify that changes to SagemakerChatConfig don't break existing sagemaker_chat callers.""" + + def setup_method(self): + from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig + self.config = SagemakerChatConfig() + + def test_should_not_support_stream_param_in_request_body(self): + """sagemaker_chat should NOT send stream in request body (unchanged behavior).""" + assert self.config.supports_stream_param_in_request_body is False + + def test_should_generate_correct_urls(self): + """sagemaker_chat URLs should be unchanged.""" + url = self.config.get_complete_url( + api_base=None, + api_key=None, + model="my-hf-endpoint", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={}, + stream=False, + ) + assert url == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/my-hf-endpoint/invocations" + + stream_url = self.config.get_complete_url( + api_base=None, + api_key=None, + model="my-hf-endpoint", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={}, + stream=True, + ) + assert stream_url == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/my-hf-endpoint/invocations-response-stream" + + def test_should_still_have_custom_stream_wrapper(self): + """sagemaker_chat should still use custom stream wrapper.""" + assert self.config.has_custom_stream_wrapper is True + + def test_should_not_include_nova_specific_params(self): + """sagemaker_chat should NOT have Nova-specific params.""" + params = self.config.get_supported_openai_params(model="my-hf-endpoint") + assert "top_k" not in params + assert "reasoning_effort" not in params + assert "allowed_token_ids" not in params + assert "truncate_prompt_tokens" not in params + + def test_should_preserve_standard_openai_params(self): + """sagemaker_chat should still support standard OpenAI params.""" + params = self.config.get_supported_openai_params(model="my-hf-endpoint") + assert "temperature" in params + assert "max_tokens" in params + assert "top_p" in params + assert "stream" in params + + def test_sync_stream_wrapper_uses_correct_provider_string(self): + """ + Verify that when get_sync_custom_stream_wrapper is called with + custom_llm_provider="sagemaker_chat", the CustomStreamWrapper + receives "sagemaker_chat" (not something else). + """ + from unittest.mock import patch, MagicMock + + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.iter_bytes.return_value = iter([]) + mock_client.post.return_value = mock_response + + with patch("litellm.llms.sagemaker.chat.transformation.CustomStreamWrapper") as mock_csw: + mock_csw.return_value = MagicMock() + self.config.get_sync_custom_stream_wrapper( + model="my-hf-endpoint", + custom_llm_provider="sagemaker_chat", + logging_obj=MagicMock(), + api_base="https://example.com", + headers={}, + data={}, + messages=[], + client=mock_client, + ) + mock_csw.assert_called_once() + call_kwargs = mock_csw.call_args[1] + assert call_kwargs["custom_llm_provider"] == "sagemaker_chat" + + def test_async_stream_wrapper_uses_correct_provider_string(self): + """ + Verify that when get_async_custom_stream_wrapper is called with + custom_llm_provider="sagemaker_chat", the CustomStreamWrapper + receives "sagemaker_chat". + """ + import asyncio + from unittest.mock import patch, MagicMock, AsyncMock + + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.status_code = 200 + + async def empty_aiter(): + return + yield # make it an async generator + + mock_response.aiter_bytes.return_value = empty_aiter() + mock_client.post.return_value = mock_response + + with patch("litellm.llms.sagemaker.chat.transformation.CustomStreamWrapper") as mock_csw: + mock_csw.return_value = MagicMock() + asyncio.run( + self.config.get_async_custom_stream_wrapper( + model="my-hf-endpoint", + custom_llm_provider="sagemaker_chat", + logging_obj=MagicMock(), + api_base="https://example.com", + headers={}, + data={}, + messages=[], + client=mock_client, + ) + ) + mock_csw.assert_called_once() + call_kwargs = mock_csw.call_args[1] + assert call_kwargs["custom_llm_provider"] == "sagemaker_chat" + + def test_async_stream_wrapper_llm_provider_enum_resolves(self): + """ + Verify LlmProviders(custom_llm_provider) resolves correctly for + "sagemaker_chat" and doesn't fall through to the ValueError fallback. + """ + from litellm.types.utils import LlmProviders + provider = LlmProviders("sagemaker_chat") + assert provider == LlmProviders.SAGEMAKER_CHAT + + +class TestSagemakerNovaTransformRequest: + """Test Nova-specific request transformation.""" + + def setup_method(self): + self.config = SagemakerNovaConfig() + + def test_should_not_include_model_in_request_body(self): + """Nova SageMaker endpoints reject 'model' in the request body.""" + request = self.config.transform_request( + model="my-nova-endpoint", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"temperature": 0.7}, + litellm_params={}, + headers={}, + ) + assert "model" not in request + assert "messages" in request + assert request["temperature"] == 0.7 + + def test_should_include_all_nova_params_in_request(self): + """Nova-specific params should appear in the request body.""" + request = self.config.transform_request( + model="my-nova-endpoint", + messages=[{"role": "user", "content": "Hello"}], + optional_params={ + "top_k": 40, + "max_tokens": 512, + "reasoning_effort": "low", + }, + litellm_params={}, + headers={}, + ) + assert "model" not in request + assert request["top_k"] == 40 + assert request["max_tokens"] == 512 + assert request["reasoning_effort"] == "low" From e0b3fcb34cf6e2153e2af88ff745736a109394af Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 14 Mar 2026 15:19:40 -0700 Subject: [PATCH 15/16] refactor: update pr template to invite users to slack oss --- .github/pull_request_template.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index bd434bea39d..d830c16dfa2 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -11,6 +11,10 @@ - [ ] My PR's scope is as isolated as possible, it only solves 1 specific problem - [ ] I have requested a Greptile review by commenting `@greptileai` and received a **Confidence Score of at least 4/5** before requesting a maintainer review +## Delays in PR merge? + +If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slack (#pr-review)](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA). + ## CI (LiteLLM team) > **CI status guideline:** From b96f033c90cd808524e51dcdbf874d37134eabf3 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 14 Mar 2026 16:54:21 -0700 Subject: [PATCH 16/16] fix: prisma migrate deploy failures on pre-existing instances (#23655) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: prisma migrate deploy failures on pre-existing instances Fixes failed migrations due to idempotent schema changes on pre-existing litellm instances. Problems: 1. P3018 recovery handler never returned True on successful resolution, causing "Database setup failed after multiple retries" even when the final recovery succeeded 2. _roll_back_migration exceptions escaped the P3018 handler, preventing _resolve_specific_migration from running 3. Migration SQL used ADD COLUMN/DROP COLUMN without IF [NOT] EXISTS, failing if schema was already modified Changes: - Add return True after successful P3018 idempotent error recovery - Wrap _roll_back_migration in try/except to allow recovery continuation even if rollback fails - Make migration.sql idempotent with IF NOT EXISTS / IF EXISTS clauses Co-Authored-By: Claude Haiku 4.5 * test: add migration SQL idempotency safety tests Adds TestMigrationSQLIdempotency test class that statically validates all migration SQL files created after 2026-03-11 use idempotent DDL: - ADD COLUMN must use IF NOT EXISTS - DROP COLUMN must use IF EXISTS - DROP INDEX must use IF EXISTS - CREATE INDEX must use IF NOT EXISTS This prevents the class of errors where prisma migrate deploy fails on pre-existing instances because the schema was already modified. Co-Authored-By: Claude Opus 4.6 * fix: also catch TimeoutExpired in P3018 rollback handler _roll_back_migration uses subprocess.run with timeout=60, so it can raise subprocess.TimeoutExpired in addition to CalledProcessError. Without catching this, a slow database during rollback would escape the handler and bypass _resolve_specific_migration — the same class of bug. Co-Authored-By: Claude Opus 4.6 * fix: make all 85 migration SQL files idempotent, remove test cutoff Fixed all existing migration files to use IF [NOT] EXISTS for DDL statements (ADD COLUMN, DROP COLUMN, DROP INDEX, CREATE INDEX). Removed the date cutoff from the idempotency tests so they now validate all migrations, not just recent ones. Co-Authored-By: Claude Opus 4.6 * fix: make migration failure non-fatal by default, add --require_db_migration flag By default the proxy now warns and continues when database migration fails. Pass --require_db_migration (or set REQUIRE_DB_MIGRATION=true) to restore the previous behavior of exiting with an error. Co-Authored-By: Claude Opus 4.6 * fix: wrap _resolve_specific_migration in try/except, guard RENAME COLUMN and ADD CONSTRAINT Three fixes: 1. _resolve_specific_migration in the P3018 handler was not wrapped in try/except, so failures there would bypass the return True and propagate unexpectedly — partially defeating the rollback fix. 2. Bare RENAME COLUMN in 20260303000000_update_tool_table_policies was non-idempotent. Wrapped in DO $$ IF EXISTS block. Also wrapped all 28 bare ADD CONSTRAINT statements across 9 migration files in DO $$ IF NOT EXISTS (pg_constraint) blocks. 3. Added test_rename_column_is_guarded and test_add_constraint_is_guarded to TestMigrationSQLIdempotency for full DDL coverage. Co-Authored-By: Claude Opus 4.6 * fix: retry after resolving idempotent migration, guard DROP CONSTRAINT Three fixes: 1. Both P3009 and P3018 idempotent handlers returned True after resolving a single migration, exiting before remaining pending migrations were applied. Now they continue the retry loop so prisma migrate deploy runs again for any remaining migrations. 2. Two migration files had bare DROP CONSTRAINT without a DO $$ IF EXISTS guard, which fails if the constraint was already dropped. Wrapped both in idempotent DO $$ blocks. 3. Added test_drop_constraint_is_guarded to catch unguarded DROP CONSTRAINT in future migrations. Co-Authored-By: Claude Opus 4.6 * fix: P3009 try/except, CREATE TABLE IF NOT EXISTS, restore fail-fast default Four fixes: 1. P3009 idempotent handler now has the same try/except around _roll_back_migration and _resolve_specific_migration as the P3018 handler. Previously a rollback or resolve failure in the P3009 path would propagate and leave the migration unresolved. 2. Added IF NOT EXISTS to all 57 bare CREATE TABLE statements across 34 migration files. Added test_create_table_uses_if_not_exists to catch this pattern. 3. Reverted the backwards-incompatible default behavior change: the proxy now fails fast on migration failure (original behavior). Added --skip_db_migration_check / SKIP_DB_MIGRATION_CHECK to opt into warn-and-continue instead. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Haiku 4.5 --- .../20250326162113_baseline/migration.sql | 144 +++++++++++---- .../migration.sql | 12 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 4 +- .../migration.sql | 6 +- .../migration.sql | 2 +- .../migration.sql | 4 +- .../migration.sql | 12 +- .../migration.sql | 18 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 4 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 6 +- .../migration.sql | 38 +++- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 4 +- .../migration.sql | 16 +- .../migration.sql | 2 +- .../migration.sql | 14 +- .../migration.sql | 8 +- .../migration.sql | 14 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 26 +-- .../migration.sql | 10 +- .../migration.sql | 2 +- .../20250802162330_prompt_table/migration.sql | 4 +- .../migration.sql | 2 +- .../migration.sql | 10 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 9 +- .../migration.sql | 4 +- .../migration.sql | 4 +- .../migration.sql | 4 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 32 +++- .../migration.sql | 6 +- .../migration.sql | 2 +- .../migration.sql | 14 +- .../20251114182247_agents_table/migration.sql | 4 +- .../migration.sql | 6 +- .../migration.sql | 2 +- .../migration.sql | 14 +- .../migration.sql | 2 +- .../migration.sql | 16 +- .../migration.sql | 28 +-- .../migration.sql | 2 +- .../migration.sql | 6 +- .../migration.sql | 2 +- .../migration.sql | 48 ++--- .../migration.sql | 4 +- .../migration.sql | 2 +- .../migration.sql | 4 +- .../migration.sql | 16 +- .../migration.sql | 8 +- .../migration.sql | 4 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 6 +- .../migration.sql | 2 +- .../migration.sql | 12 +- .../migration.sql | 6 +- .../migration.sql | 4 +- .../20260214094754_schema_sync/migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 9 +- .../migration.sql | 4 +- .../migration.sql | 2 +- .../migration.sql | 18 +- .../migration.sql | 2 +- .../migration.sql | 24 +-- .../migration.sql | 18 +- .../migration.sql | 2 +- .../migration.sql | 2 +- .../migration.sql | 21 ++- .../migration.sql | 2 +- .../migration.sql | 4 +- .../migration.sql | 2 +- .../migration.sql | 4 +- .../migration.sql | 2 +- .../migration.sql | 8 +- .../migration.sql | 19 +- .../migration.sql | 2 +- .../migration.sql | 8 +- .../migration.sql | 2 +- .../20260306233848_schema_sync/migration.sql | 29 +-- .../20260311180521_schema_sync/migration.sql | 14 +- .../20260312124619_schema_sync/migration.sql | 2 +- .../litellm_proxy_extras/utils.py | 72 +++++--- litellm/proxy/proxy_cli.py | 24 ++- .../test_litellm_proxy_extras_utils.py | 169 ++++++++++++++++++ 99 files changed, 751 insertions(+), 389 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql index fb8a44814f0..7b33d58899e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_BudgetTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetTable" ( "budget_id" TEXT NOT NULL, "max_budget" DOUBLE PRECISION, "soft_budget" DOUBLE PRECISION, @@ -18,7 +18,7 @@ CREATE TABLE "LiteLLM_BudgetTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_CredentialsTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_CredentialsTable" ( "credential_id" TEXT NOT NULL, "credential_name" TEXT NOT NULL, "credential_values" JSONB NOT NULL, @@ -32,7 +32,7 @@ CREATE TABLE "LiteLLM_CredentialsTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_ProxyModelTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ProxyModelTable" ( "model_id" TEXT NOT NULL, "model_name" TEXT NOT NULL, "litellm_params" JSONB NOT NULL, @@ -46,7 +46,7 @@ CREATE TABLE "LiteLLM_ProxyModelTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_OrganizationTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_OrganizationTable" ( "organization_id" TEXT NOT NULL, "organization_alias" TEXT NOT NULL, "budget_id" TEXT NOT NULL, @@ -63,7 +63,7 @@ CREATE TABLE "LiteLLM_OrganizationTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_ModelTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ModelTable" ( "id" SERIAL NOT NULL, "aliases" JSONB, "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, @@ -75,7 +75,7 @@ CREATE TABLE "LiteLLM_ModelTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_TeamTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_TeamTable" ( "team_id" TEXT NOT NULL, "team_alias" TEXT, "organization_id" TEXT, @@ -102,7 +102,7 @@ CREATE TABLE "LiteLLM_TeamTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_UserTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_UserTable" ( "user_id" TEXT NOT NULL, "user_alias" TEXT, "team_id" TEXT, @@ -131,7 +131,7 @@ CREATE TABLE "LiteLLM_UserTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_VerificationToken" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_VerificationToken" ( "token" TEXT NOT NULL, "key_name" TEXT, "key_alias" TEXT, @@ -166,7 +166,7 @@ CREATE TABLE "LiteLLM_VerificationToken" ( ); -- CreateTable -CREATE TABLE "LiteLLM_EndUserTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_EndUserTable" ( "user_id" TEXT NOT NULL, "alias" TEXT, "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, @@ -179,7 +179,7 @@ CREATE TABLE "LiteLLM_EndUserTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_Config" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_Config" ( "param_name" TEXT NOT NULL, "param_value" JSONB, @@ -187,7 +187,7 @@ CREATE TABLE "LiteLLM_Config" ( ); -- CreateTable -CREATE TABLE "LiteLLM_SpendLogs" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs" ( "request_id" TEXT NOT NULL, "call_type" TEXT NOT NULL, "api_key" TEXT NOT NULL DEFAULT '', @@ -218,7 +218,7 @@ CREATE TABLE "LiteLLM_SpendLogs" ( ); -- CreateTable -CREATE TABLE "LiteLLM_ErrorLogs" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ErrorLogs" ( "request_id" TEXT NOT NULL, "startTime" TIMESTAMP(3) NOT NULL, "endTime" TIMESTAMP(3) NOT NULL, @@ -235,7 +235,7 @@ CREATE TABLE "LiteLLM_ErrorLogs" ( ); -- CreateTable -CREATE TABLE "LiteLLM_UserNotifications" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_UserNotifications" ( "request_id" TEXT NOT NULL, "user_id" TEXT NOT NULL, "models" TEXT[], @@ -246,7 +246,7 @@ CREATE TABLE "LiteLLM_UserNotifications" ( ); -- CreateTable -CREATE TABLE "LiteLLM_TeamMembership" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_TeamMembership" ( "user_id" TEXT NOT NULL, "team_id" TEXT NOT NULL, "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, @@ -256,7 +256,7 @@ CREATE TABLE "LiteLLM_TeamMembership" ( ); -- CreateTable -CREATE TABLE "LiteLLM_OrganizationMembership" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_OrganizationMembership" ( "user_id" TEXT NOT NULL, "organization_id" TEXT NOT NULL, "user_role" TEXT, @@ -269,7 +269,7 @@ CREATE TABLE "LiteLLM_OrganizationMembership" ( ); -- CreateTable -CREATE TABLE "LiteLLM_InvitationLink" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_InvitationLink" ( "id" TEXT NOT NULL, "user_id" TEXT NOT NULL, "is_accepted" BOOLEAN NOT NULL DEFAULT false, @@ -284,7 +284,7 @@ CREATE TABLE "LiteLLM_InvitationLink" ( ); -- CreateTable -CREATE TABLE "LiteLLM_AuditLog" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_AuditLog" ( "id" TEXT NOT NULL, "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, "changed_by" TEXT NOT NULL DEFAULT '', @@ -299,62 +299,132 @@ CREATE TABLE "LiteLLM_AuditLog" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_CredentialsTable_credential_name_key" ON "LiteLLM_CredentialsTable"("credential_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_CredentialsTable_credential_name_key" ON "LiteLLM_CredentialsTable"("credential_name"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_TeamTable_model_id_key" ON "LiteLLM_TeamTable"("model_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_TeamTable_model_id_key" ON "LiteLLM_TeamTable"("model_id"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_UserTable_sso_user_id_key" ON "LiteLLM_UserTable"("sso_user_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_UserTable_sso_user_id_key" ON "LiteLLM_UserTable"("sso_user_id"); -- CreateIndex -CREATE INDEX "LiteLLM_SpendLogs_startTime_idx" ON "LiteLLM_SpendLogs"("startTime"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx" ON "LiteLLM_SpendLogs"("startTime"); -- CreateIndex -CREATE INDEX "LiteLLM_SpendLogs_end_user_idx" ON "LiteLLM_SpendLogs"("end_user"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx" ON "LiteLLM_SpendLogs"("end_user"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_OrganizationMembership_user_id_organization_id_key" ON "LiteLLM_OrganizationMembership"("user_id", "organization_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_OrganizationMembership_user_id_organization_id_key" ON "LiteLLM_OrganizationMembership"("user_id", "organization_id"); -- AddForeignKey -ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE RESTRICT ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationTable_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE RESTRICT ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamTable_organization_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_model_id_fkey" FOREIGN KEY ("model_id") REFERENCES "LiteLLM_ModelTable"("id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamTable_model_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_model_id_fkey" FOREIGN KEY ("model_id") REFERENCES "LiteLLM_ModelTable"("id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_UserTable_organization_id_fkey') THEN + ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerificationToken_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerificationToken_organization_id_fkey') THEN + ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_EndUserTable_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationMembership_user_id_fkey') THEN + ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE RESTRICT ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationMembership_organization_id_fkey') THEN + ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE RESTRICT ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationMembership_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_InvitationLink_user_id_fkey') THEN + ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_created_by_fkey" FOREIGN KEY ("created_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_InvitationLink_created_by_fkey') THEN + ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_created_by_fkey" FOREIGN KEY ("created_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_updated_by_fkey" FOREIGN KEY ("updated_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_InvitationLink_updated_by_fkey') THEN + ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_updated_by_fkey" FOREIGN KEY ("updated_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql index 3379d8e9fda..52f20ee0f28 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_DailyUserSpend" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyUserSpend" ( "id" TEXT NOT NULL, "user_id" TEXT NOT NULL, "date" TEXT NOT NULL, @@ -17,17 +17,17 @@ CREATE TABLE "LiteLLM_DailyUserSpend" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_DailyUserSpend_date_idx" ON "LiteLLM_DailyUserSpend"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_date_idx" ON "LiteLLM_DailyUserSpend"("date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyUserSpend_user_id_idx" ON "LiteLLM_DailyUserSpend"("user_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_idx" ON "LiteLLM_DailyUserSpend"("user_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyUserSpend_api_key_idx" ON "LiteLLM_DailyUserSpend"("api_key"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_api_key_idx" ON "LiteLLM_DailyUserSpend"("api_key"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyUserSpend_model_idx" ON "LiteLLM_DailyUserSpend"("model"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_model_idx" ON "LiteLLM_DailyUserSpend"("model"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql index e7c5ab566a9..3865194ce88 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "api_requests" INTEGER NOT NULL DEFAULT 0; +ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "api_requests" INTEGER NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql index e7ea2e9015a..ba3000f7512 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql @@ -2,7 +2,7 @@ CREATE TYPE "JobStatus" AS ENUM ('ACTIVE', 'INACTIVE'); -- CreateTable -CREATE TABLE "LiteLLM_CronJob" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_CronJob" ( "cronjob_id" TEXT NOT NULL, "pod_id" TEXT NOT NULL, "status" "JobStatus" NOT NULL DEFAULT 'INACTIVE', diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql index 9f1693500d0..47962aec772 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql @@ -1,4 +1,4 @@ -- AlterTable -ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "failed_requests" INTEGER NOT NULL DEFAULT 0, -ADD COLUMN "successful_requests" INTEGER NOT NULL DEFAULT 0; +ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "failed_requests" INTEGER NOT NULL DEFAULT 0, +ADD COLUMN IF NOT EXISTS "successful_requests" INTEGER NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql index d14a6294581..602cb5ada37 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_ManagedFileTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedFileTable" ( "id" TEXT NOT NULL, "unified_file_id" TEXT NOT NULL, "file_object" JSONB NOT NULL, @@ -11,8 +11,8 @@ CREATE TABLE "LiteLLM_ManagedFileTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_ManagedFileTable_unified_file_id_key" ON "LiteLLM_ManagedFileTable"("unified_file_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedFileTable_unified_file_id_key" ON "LiteLLM_ManagedFileTable"("unified_file_id"); -- CreateIndex -CREATE INDEX "LiteLLM_ManagedFileTable_unified_file_id_idx" ON "LiteLLM_ManagedFileTable"("unified_file_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedFileTable_unified_file_id_idx" ON "LiteLLM_ManagedFileTable"("unified_file_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql index c07df813796..334b421ca31 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "team_member_permissions" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "team_member_permissions" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql index f47e1c2e91b..1a39ee3579d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql @@ -1,4 +1,4 @@ -- AlterTable -ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0, -ADD COLUMN "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0; +ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0, +ADD COLUMN IF NOT EXISTS "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql index a6eb461bc2f..c8ef4eff7d3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_DailyTeamSpend" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyTeamSpend" ( "id" TEXT NOT NULL, "team_id" TEXT NOT NULL, "date" TEXT NOT NULL, @@ -20,17 +20,17 @@ CREATE TABLE "LiteLLM_DailyTeamSpend" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTeamSpend_date_idx" ON "LiteLLM_DailyTeamSpend"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_date_idx" ON "LiteLLM_DailyTeamSpend"("date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTeamSpend_team_id_idx" ON "LiteLLM_DailyTeamSpend"("team_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_idx" ON "LiteLLM_DailyTeamSpend"("team_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTeamSpend_api_key_idx" ON "LiteLLM_DailyTeamSpend"("api_key"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_api_key_idx" ON "LiteLLM_DailyTeamSpend"("api_key"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTeamSpend_model_idx" ON "LiteLLM_DailyTeamSpend"("model"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_model_idx" ON "LiteLLM_DailyTeamSpend"("model"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql index 8c3cea70937..0a528348144 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql @@ -1,9 +1,9 @@ -- AlterTable -ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0, -ADD COLUMN "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0; +ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0, +ADD COLUMN IF NOT EXISTS "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0; -- CreateTable -CREATE TABLE "LiteLLM_DailyTagSpend" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyTagSpend" ( "id" TEXT NOT NULL, "tag" TEXT NOT NULL, "date" TEXT NOT NULL, @@ -26,20 +26,20 @@ CREATE TABLE "LiteLLM_DailyTagSpend" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_key" ON "LiteLLM_DailyTagSpend"("tag"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_key" ON "LiteLLM_DailyTagSpend"("tag"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTagSpend_date_idx" ON "LiteLLM_DailyTagSpend"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_date_idx" ON "LiteLLM_DailyTagSpend"("date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTagSpend_tag_idx" ON "LiteLLM_DailyTagSpend"("tag"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_idx" ON "LiteLLM_DailyTagSpend"("tag"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTagSpend_api_key_idx" ON "LiteLLM_DailyTagSpend"("api_key"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_api_key_idx" ON "LiteLLM_DailyTagSpend"("api_key"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTagSpend_model_idx" ON "LiteLLM_DailyTagSpend"("model"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_model_idx" ON "LiteLLM_DailyTagSpend"("model"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql index 5c27b84efbf..a5793331eac 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql @@ -1,3 +1,3 @@ -- DropIndex -DROP INDEX "LiteLLM_DailyTagSpend_tag_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_key"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql index 2ee7838dcfa..fd8997813f4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "allowed_routes" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "allowed_routes" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql index 751c75e5f24..f0ad3886eef 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql @@ -1,4 +1,4 @@ -- AlterTable -ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "proxy_server_request" JSONB DEFAULT '{}', -ADD COLUMN "session_id" TEXT; +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "proxy_server_request" JSONB DEFAULT '{}', +ADD COLUMN IF NOT EXISTS "session_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql index 39e7f2f3b20..ae73e40cd59 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_ManagedVectorStoresTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedVectorStoresTable" ( "vector_store_id" TEXT NOT NULL, "custom_llm_provider" TEXT NOT NULL, "vector_store_name" TEXT, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql index 6b8adc6e7e8..95f26d5cb55 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_MCPServerTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_MCPServerTable" ( "server_id" TEXT NOT NULL, "alias" TEXT, "description" TEXT, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161527_add_health_check_fields_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161527_add_health_check_fields_to_mcp_servers/migration.sql index d5c206d1929..71f595173b7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161527_add_health_check_fields_to_mcp_servers/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161527_add_health_check_fields_to_mcp_servers/migration.sql @@ -1,4 +1,4 @@ -- Add health check fields to MCP server table -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "status" TEXT DEFAULT 'unknown'; -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "last_health_check" TIMESTAMP(3); -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "health_check_error" TEXT; \ No newline at end of file +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "status" TEXT DEFAULT 'unknown'; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "last_health_check" TIMESTAMP(3); +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "health_check_error" TEXT; \ No newline at end of file diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql index dcfce07a487..9937857875e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql @@ -1,17 +1,17 @@ -- AlterTable -ALTER TABLE "LiteLLM_OrganizationTable" ADD COLUMN "object_permission_id" TEXT; +ALTER TABLE "LiteLLM_OrganizationTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "object_permission_id" TEXT; +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_UserTable" ADD COLUMN "object_permission_id" TEXT; +ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "object_permission_id" TEXT; +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT; -- CreateTable -CREATE TABLE "LiteLLM_ObjectPermissionTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ObjectPermissionTable" ( "object_permission_id" TEXT NOT NULL, "mcp_servers" TEXT[] DEFAULT ARRAY[]::TEXT[], @@ -19,14 +19,34 @@ CREATE TABLE "LiteLLM_ObjectPermissionTable" ( ); -- AddForeignKey -ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationTable_object_permission_id_fkey') THEN + ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamTable_object_permission_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_UserTable_object_permission_id_fkey') THEN + ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerificationToken_object_permission_id_fkey') THEN + ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql index 8f6c68aa67e..a001977be55 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "status" TEXT; +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "status" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql index eda055d6e56..a647f386d9e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql @@ -1,3 +1,3 @@ -- CreateIndex -CREATE INDEX "LiteLLM_SpendLogs_session_id_idx" ON "LiteLLM_SpendLogs"("session_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx" ON "LiteLLM_SpendLogs"("session_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250514142245_add_guardrails_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250514142245_add_guardrails_table/migration.sql index fa99e3be637..41d8c188947 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250514142245_add_guardrails_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250514142245_add_guardrails_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_GuardrailsTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_GuardrailsTable" ( "guardrail_id" TEXT NOT NULL, "guardrail_name" TEXT NOT NULL, "litellm_params" JSONB NOT NULL, @@ -11,5 +11,5 @@ CREATE TABLE "LiteLLM_GuardrailsTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_GuardrailsTable_guardrail_name_key" ON "LiteLLM_GuardrailsTable"("guardrail_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_GuardrailsTable_guardrail_name_key" ON "LiteLLM_GuardrailsTable"("guardrail_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250522223020_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250522223020_managed_object_table/migration.sql index 95fb8372458..564811655f3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250522223020_managed_object_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250522223020_managed_object_table/migration.sql @@ -1,10 +1,10 @@ -- AlterTable -ALTER TABLE "LiteLLM_ManagedFileTable" ADD COLUMN "created_by" TEXT, -ADD COLUMN "flat_model_file_ids" TEXT[] DEFAULT ARRAY[]::TEXT[], -ADD COLUMN "updated_by" TEXT; +ALTER TABLE "LiteLLM_ManagedFileTable" ADD COLUMN IF NOT EXISTS "created_by" TEXT, +ADD COLUMN IF NOT EXISTS "flat_model_file_ids" TEXT[] DEFAULT ARRAY[]::TEXT[], +ADD COLUMN IF NOT EXISTS "updated_by" TEXT; -- CreateTable -CREATE TABLE "LiteLLM_ManagedObjectTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedObjectTable" ( "id" TEXT NOT NULL, "unified_object_id" TEXT NOT NULL, "model_object_id" TEXT NOT NULL, @@ -19,14 +19,14 @@ CREATE TABLE "LiteLLM_ManagedObjectTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_ManagedObjectTable_unified_object_id_key" ON "LiteLLM_ManagedObjectTable"("unified_object_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedObjectTable_unified_object_id_key" ON "LiteLLM_ManagedObjectTable"("unified_object_id"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_ManagedObjectTable_model_object_id_key" ON "LiteLLM_ManagedObjectTable"("model_object_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedObjectTable_model_object_id_key" ON "LiteLLM_ManagedObjectTable"("model_object_id"); -- CreateIndex -CREATE INDEX "LiteLLM_ManagedObjectTable_unified_object_id_idx" ON "LiteLLM_ManagedObjectTable"("unified_object_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedObjectTable_unified_object_id_idx" ON "LiteLLM_ManagedObjectTable"("unified_object_id"); -- CreateIndex -CREATE INDEX "LiteLLM_ManagedObjectTable_model_object_id_idx" ON "LiteLLM_ManagedObjectTable"("model_object_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedObjectTable_model_object_id_idx" ON "LiteLLM_ManagedObjectTable"("model_object_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250528185438_add_vector_stores_to_object_permissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250528185438_add_vector_stores_to_object_permissions/migration.sql index 39db701056e..f1f76153496 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250528185438_add_vector_stores_to_object_permissions/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250528185438_add_vector_stores_to_object_permissions/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "vector_stores" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "vector_stores" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql index 3d36e42577c..ef53fbf3f4a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql @@ -1,6 +1,16 @@ -- DropForeignKey -ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey"; +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey"; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE CASCADE ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE CASCADE ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250618225828_add_health_check_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250618225828_add_health_check_table/migration.sql index da6f4c23c81..2b4dd1d17d3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250618225828_add_health_check_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250618225828_add_health_check_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_HealthCheckTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_HealthCheckTable" ( "health_check_id" TEXT NOT NULL, "model_name" TEXT NOT NULL, "model_id" TEXT, @@ -18,11 +18,11 @@ CREATE TABLE "LiteLLM_HealthCheckTable" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_HealthCheckTable_model_name_idx" ON "LiteLLM_HealthCheckTable"("model_name"); +CREATE INDEX IF NOT EXISTS "LiteLLM_HealthCheckTable_model_name_idx" ON "LiteLLM_HealthCheckTable"("model_name"); -- CreateIndex -CREATE INDEX "LiteLLM_HealthCheckTable_checked_at_idx" ON "LiteLLM_HealthCheckTable"("checked_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_HealthCheckTable_checked_at_idx" ON "LiteLLM_HealthCheckTable"("checked_at"); -- CreateIndex -CREATE INDEX "LiteLLM_HealthCheckTable_status_idx" ON "LiteLLM_HealthCheckTable"("status"); +CREATE INDEX IF NOT EXISTS "LiteLLM_HealthCheckTable_status_idx" ON "LiteLLM_HealthCheckTable"("status"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql index 51461b82058..75d7e0e74b1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql @@ -1,9 +1,19 @@ -- DropForeignKey -ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey"; +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey"; + END IF; +END $$; -- AlterTable ALTER TABLE "LiteLLM_ManagedFileTable" ALTER COLUMN "file_object" DROP NOT NULL; -- AddForeignKey -ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625213625_add_status_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625213625_add_status_to_managed_object_table/migration.sql index 7ca7b2c3705..a4f0ca07944 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625213625_add_status_to_managed_object_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625213625_add_status_to_managed_object_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN "status" TEXT; +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "status" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707212517_add_mcp_info_column_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707212517_add_mcp_info_column_mcp_servers/migration.sql index efe68ff4792..e942383ded7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707212517_add_mcp_info_column_mcp_servers/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707212517_add_mcp_info_column_mcp_servers/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "mcp_info" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "mcp_info" JSONB DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707230009_add_mcp_namespaced_tool_name/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707230009_add_mcp_namespaced_tool_name/migration.sql index 3130619a773..963b29e1875 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707230009_add_mcp_namespaced_tool_name/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707230009_add_mcp_namespaced_tool_name/migration.sql @@ -1,42 +1,42 @@ -- DropIndex -DROP INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; -- DropIndex -DROP INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key"; -- DropIndex -DROP INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key"; -- AlterTable -ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN "mcp_namespaced_tool_name" TEXT, +ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "mcp_namespaced_tool_name" TEXT, ALTER COLUMN "model" DROP NOT NULL; -- AlterTable -ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN "mcp_namespaced_tool_name" TEXT, +ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "mcp_namespaced_tool_name" TEXT, ALTER COLUMN "model" DROP NOT NULL; -- AlterTable -ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "mcp_namespaced_tool_name" TEXT, +ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "mcp_namespaced_tool_name" TEXT, ALTER COLUMN "model" DROP NOT NULL; -- AlterTable -ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "mcp_namespaced_tool_name" TEXT; +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "mcp_namespaced_tool_name" TEXT; -- CreateIndex -CREATE INDEX "LiteLLM_DailyTagSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyTagSpend"("mcp_namespaced_tool_name"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyTagSpend"("mcp_namespaced_tool_name"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTeamSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyTeamSpend"("mcp_namespaced_tool_name"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyTeamSpend"("mcp_namespaced_tool_name"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyUserSpend"("mcp_namespaced_tool_name"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyUserSpend"("mcp_namespaced_tool_name"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250711220620_add_stdio_mcp/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250711220620_add_stdio_mcp/migration.sql index ebe7a6adb58..685dc6c1d39 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250711220620_add_stdio_mcp/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250711220620_add_stdio_mcp/migration.sql @@ -1,10 +1,10 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "args" TEXT[] DEFAULT ARRAY[]::TEXT[], -ADD COLUMN "command" TEXT, -ADD COLUMN "env" JSONB DEFAULT '{}', -ADD COLUMN "mcp_access_groups" TEXT[], +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "args" TEXT[] DEFAULT ARRAY[]::TEXT[], +ADD COLUMN IF NOT EXISTS "command" TEXT, +ADD COLUMN IF NOT EXISTS "env" JSONB DEFAULT '{}', +ADD COLUMN IF NOT EXISTS "mcp_access_groups" TEXT[], ALTER COLUMN "url" DROP NOT NULL; -- AlterTable -ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "mcp_access_groups" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "mcp_access_groups" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250718125714_add_litellm_params_to_vector_stores/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250718125714_add_litellm_params_to_vector_stores/migration.sql index ef9956ddd5f..5f502a374e2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250718125714_add_litellm_params_to_vector_stores/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250718125714_add_litellm_params_to_vector_stores/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_ManagedVectorStoresTable" ADD COLUMN "litellm_params" JSONB; +ALTER TABLE "LiteLLM_ManagedVectorStoresTable" ADD COLUMN IF NOT EXISTS "litellm_params" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250802162330_prompt_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250802162330_prompt_table/migration.sql index e5c00ef4adb..81b3574499b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250802162330_prompt_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250802162330_prompt_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_PromptTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_PromptTable" ( "id" TEXT NOT NULL, "prompt_id" TEXT NOT NULL, "litellm_params" JSONB NOT NULL, @@ -11,5 +11,5 @@ CREATE TABLE "LiteLLM_PromptTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_PromptTable_prompt_id_key" ON "LiteLLM_PromptTable"("prompt_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_PromptTable_prompt_id_key" ON "LiteLLM_PromptTable"("prompt_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql index 472e2ea1e0c..5b4db0e500a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql @@ -5,4 +5,4 @@ */ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "spec_version"; +ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN IF EXISTS "spec_version"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250926194702_unnamed_migration/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250926194702_unnamed_migration/migration.sql index ea28db19662..a40c88cc85d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250926194702_unnamed_migration/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250926194702_unnamed_migration/migration.sql @@ -1,7 +1,7 @@ -- AlterTable -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "auto_rotate" BOOLEAN DEFAULT false, -ADD COLUMN "key_rotation_at" TIMESTAMP(3), -ADD COLUMN "last_rotation_at" TIMESTAMP(3), -ADD COLUMN "rotation_count" INTEGER DEFAULT 0, -ADD COLUMN "rotation_interval" TEXT; +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "auto_rotate" BOOLEAN DEFAULT false, +ADD COLUMN IF NOT EXISTS "key_rotation_at" TIMESTAMP(3), +ADD COLUMN IF NOT EXISTS "last_rotation_at" TIMESTAMP(3), +ADD COLUMN IF NOT EXISTS "rotation_count" INTEGER DEFAULT 0, +ADD COLUMN IF NOT EXISTS "rotation_interval" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003165142_add_allowed_tools_to_mcp/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003165142_add_allowed_tools_to_mcp/migration.sql index bdac1e42bc2..d5206e8d3a5 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003165142_add_allowed_tools_to_mcp/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003165142_add_allowed_tools_to_mcp/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "allowed_tools" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "allowed_tools" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003190954_extra_headers_to_mcp_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003190954_extra_headers_to_mcp_table/migration.sql index 1cfcf062eb1..b6275c6421c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003190954_extra_headers_to_mcp_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003190954_extra_headers_to_mcp_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251006143948_add_mcp_tool_permissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251006143948_add_mcp_tool_permissions/migration.sql index 51f3be87582..70584a96286 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251006143948_add_mcp_tool_permissions/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251006143948_add_mcp_tool_permissions/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "mcp_tool_permissions" JSONB; +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251011084309_add_tag_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251011084309_add_tag_table/migration.sql index 541c70c7e48..c1005a75a4a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251011084309_add_tag_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251011084309_add_tag_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_TagTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_TagTable" ( "tag_name" TEXT NOT NULL, "description" TEXT, "models" TEXT[], @@ -14,5 +14,10 @@ CREATE TABLE "LiteLLM_TagTable" ( ); -- AddForeignKey -ALTER TABLE "LiteLLM_TagTable" ADD CONSTRAINT "LiteLLM_TagTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TagTable_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_TagTable" ADD CONSTRAINT "LiteLLM_TagTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql index 4cbe4a7184f..3cd2a4c4899 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_SearchToolsTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_SearchToolsTable" ( "search_tool_id" TEXT NOT NULL, "search_tool_name" TEXT NOT NULL, "litellm_params" JSONB NOT NULL, @@ -11,5 +11,5 @@ CREATE TABLE "LiteLLM_SearchToolsTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_SearchToolsTable_search_tool_name_key" ON "LiteLLM_SearchToolsTable"("search_tool_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_SearchToolsTable_search_tool_name_key" ON "LiteLLM_SearchToolsTable"("search_tool_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251031181430_add_cache_config_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251031181430_add_cache_config_table/migration.sql index 705a6fd4d9b..bb475b03d25 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251031181430_add_cache_config_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251031181430_add_cache_config_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_SSOConfig" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_SSOConfig" ( "id" TEXT NOT NULL DEFAULT 'sso_config', "sso_settings" JSONB NOT NULL, "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, @@ -9,7 +9,7 @@ CREATE TABLE "LiteLLM_SSOConfig" ( ); -- CreateTable -CREATE TABLE "LiteLLM_CacheConfig" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_CacheConfig" ( "id" TEXT NOT NULL DEFAULT 'cache_config', "cache_settings" JSONB NOT NULL, "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251101131415_add_managed_vector_store_index_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251101131415_add_managed_vector_store_index_table/migration.sql index af13500d1c7..0997e569171 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251101131415_add_managed_vector_store_index_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251101131415_add_managed_vector_store_index_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_ManagedVectorStoreIndexTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedVectorStoreIndexTable" ( "id" TEXT NOT NULL, "index_name" TEXT NOT NULL, "litellm_params" JSONB NOT NULL, @@ -13,5 +13,5 @@ CREATE TABLE "LiteLLM_ManagedVectorStoreIndexTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_ManagedVectorStoreIndexTable_index_name_key" ON "LiteLLM_ManagedVectorStoreIndexTable"("index_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoreIndexTable_index_name_key" ON "LiteLLM_ManagedVectorStoreIndexTable"("index_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251103072422_add_static_headers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251103072422_add_static_headers/migration.sql index 0bedac76313..452a0b73a51 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251103072422_add_static_headers/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251103072422_add_static_headers/migration.sql @@ -1,2 +1,2 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "static_headers" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "static_headers" JSONB DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251104220043_add_credentials_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251104220043_add_credentials_to_mcp_servers/migration.sql index 800c96f18b7..c86c3f1f4b7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251104220043_add_credentials_to_mcp_servers/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251104220043_add_credentials_to_mcp_servers/migration.sql @@ -1,2 +1,2 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "credentials" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "credentials" JSONB DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000000_add_project_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000000_add_project_table/migration.sql index f1d3129bb36..9769699b33f 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000000_add_project_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000000_add_project_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_ProjectTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ProjectTable" ( "project_id" TEXT NOT NULL, "project_alias" TEXT, "team_id" TEXT, @@ -19,17 +19,37 @@ CREATE TABLE "LiteLLM_ProjectTable" ( ); -- AddForeignKey -ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_team_id_fkey" FOREIGN KEY ("team_id") REFERENCES "LiteLLM_TeamTable"("team_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ProjectTable_team_id_fkey') THEN + ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_team_id_fkey" FOREIGN KEY ("team_id") REFERENCES "LiteLLM_TeamTable"("team_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ProjectTable_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AddForeignKey -ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ProjectTable_object_permission_id_fkey') THEN + ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; -- AlterTable: Add project_id to LiteLLM_VerificationToken -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "project_id" TEXT; +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "project_id" TEXT; -- AddForeignKey -ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_project_id_fkey" FOREIGN KEY ("project_id") REFERENCES "LiteLLM_ProjectTable"("project_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerificationToken_project_id_fkey') THEN + ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_project_id_fkey" FOREIGN KEY ("project_id") REFERENCES "LiteLLM_ProjectTable"("project_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000001_add_project_fields/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000001_add_project_fields/migration.sql index 48328b4d6a2..69dbb6790c7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000001_add_project_fields/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000001_add_project_fields/migration.sql @@ -1,5 +1,5 @@ -- AlterTable: Add new fields to LiteLLM_ProjectTable -ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN "description" TEXT; -ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN "model_rpm_limit" JSONB NOT NULL DEFAULT '{}'; -ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN "model_tpm_limit" JSONB NOT NULL DEFAULT '{}'; +ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN IF NOT EXISTS "description" TEXT; +ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN IF NOT EXISTS "model_rpm_limit" JSONB NOT NULL DEFAULT '{}'; +ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN IF NOT EXISTS "model_tpm_limit" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114173537_add_request_id_to_daily_tag_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114173537_add_request_id_to_daily_tag_spend/migration.sql index 6871e27a28a..6d66ef36ae2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114173537_add_request_id_to_daily_tag_spend/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114173537_add_request_id_to_daily_tag_spend/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN "request_id" TEXT; +ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "request_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql index 74e0eea3134..387b40461a8 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_DailyOrganizationSpend" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyOrganizationSpend" ( "id" TEXT NOT NULL, "organization_id" TEXT, "date" TEXT NOT NULL, @@ -23,20 +23,20 @@ CREATE TABLE "LiteLLM_DailyOrganizationSpend" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_DailyOrganizationSpend_date_idx" ON "LiteLLM_DailyOrganizationSpend"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_date_idx" ON "LiteLLM_DailyOrganizationSpend"("date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_idx" ON "LiteLLM_DailyOrganizationSpend"("organization_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_idx" ON "LiteLLM_DailyOrganizationSpend"("organization_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyOrganizationSpend_api_key_idx" ON "LiteLLM_DailyOrganizationSpend"("api_key"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_api_key_idx" ON "LiteLLM_DailyOrganizationSpend"("api_key"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyOrganizationSpend_model_idx" ON "LiteLLM_DailyOrganizationSpend"("model"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_model_idx" ON "LiteLLM_DailyOrganizationSpend"("model"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyOrganizationSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyOrganizationSpend"("mcp_namespaced_tool_name"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyOrganizationSpend"("mcp_namespaced_tool_name"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114182247_agents_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114182247_agents_table/migration.sql index 28760dcfe48..35595c7eedf 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114182247_agents_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114182247_agents_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_AgentsTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_AgentsTable" ( "agent_id" TEXT NOT NULL, "agent_name" TEXT NOT NULL, "litellm_params" JSONB, @@ -13,5 +13,5 @@ CREATE TABLE "LiteLLM_AgentsTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_AgentsTable_agent_name_key" ON "LiteLLM_AgentsTable"("agent_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentsTable_agent_name_key" ON "LiteLLM_AgentsTable"("agent_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql index 43eb2401422..f604dcadbd9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql @@ -3,10 +3,10 @@ DROP INDEX IF EXISTS "LiteLLM_PromptTable_prompt_id_key"; -- AlterTable ALTER TABLE "LiteLLM_PromptTable" -ADD COLUMN "version" INTEGER NOT NULL DEFAULT 1; +ADD COLUMN IF NOT EXISTS "version" INTEGER NOT NULL DEFAULT 1; -- CreateIndex -CREATE INDEX "LiteLLM_PromptTable_prompt_id_idx" ON "LiteLLM_PromptTable" ("prompt_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_PromptTable_prompt_id_idx" ON "LiteLLM_PromptTable" ("prompt_id"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_PromptTable_prompt_id_version_key" ON "LiteLLM_PromptTable" ("prompt_id", "version"); \ No newline at end of file +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_PromptTable_prompt_id_version_key" ON "LiteLLM_PromptTable" ("prompt_id", "version"); \ No newline at end of file diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql index 4ea082f2750..bc483aace76 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "organization_id" TEXT; +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "organization_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql index c4234785c54..3544768cc16 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_DailyEndUserSpend" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyEndUserSpend" ( "id" TEXT NOT NULL, "end_user_id" TEXT, "date" TEXT NOT NULL, @@ -23,20 +23,20 @@ CREATE TABLE "LiteLLM_DailyEndUserSpend" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_DailyEndUserSpend_date_idx" ON "LiteLLM_DailyEndUserSpend"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_date_idx" ON "LiteLLM_DailyEndUserSpend"("date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_idx" ON "LiteLLM_DailyEndUserSpend"("end_user_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_idx" ON "LiteLLM_DailyEndUserSpend"("end_user_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyEndUserSpend_api_key_idx" ON "LiteLLM_DailyEndUserSpend"("api_key"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_api_key_idx" ON "LiteLLM_DailyEndUserSpend"("api_key"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyEndUserSpend_model_idx" ON "LiteLLM_DailyEndUserSpend"("model"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_model_idx" ON "LiteLLM_DailyEndUserSpend"("model"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyEndUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyEndUserSpend"("mcp_namespaced_tool_name"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyEndUserSpend"("mcp_namespaced_tool_name"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251209112246_add_ui_settings_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251209112246_add_ui_settings_table/migration.sql index 1719ce646d4..01e2c9fa761 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251209112246_add_ui_settings_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251209112246_add_ui_settings_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_UISettings" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_UISettings" ( "id" TEXT NOT NULL DEFAULT 'ui_settings', "ui_settings" JSONB NOT NULL, "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql index 964904c14c1..4cc894f9b22 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql @@ -1,8 +1,8 @@ -- AlterTable -ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "agent_id" TEXT; +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "agent_id" TEXT; -- CreateTable -CREATE TABLE "LiteLLM_DailyAgentSpend" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyAgentSpend" ( "id" TEXT NOT NULL, "agent_id" TEXT, "date" TEXT NOT NULL, @@ -26,20 +26,20 @@ CREATE TABLE "LiteLLM_DailyAgentSpend" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_DailyAgentSpend_date_idx" ON "LiteLLM_DailyAgentSpend"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_date_idx" ON "LiteLLM_DailyAgentSpend"("date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyAgentSpend_agent_id_idx" ON "LiteLLM_DailyAgentSpend"("agent_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_idx" ON "LiteLLM_DailyAgentSpend"("agent_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyAgentSpend_api_key_idx" ON "LiteLLM_DailyAgentSpend"("api_key"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_api_key_idx" ON "LiteLLM_DailyAgentSpend"("api_key"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyAgentSpend_model_idx" ON "LiteLLM_DailyAgentSpend"("model"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_model_idx" ON "LiteLLM_DailyAgentSpend"("model"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyAgentSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyAgentSpend"("mcp_namespaced_tool_name"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyAgentSpend"("mcp_namespaced_tool_name"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql index 6ca66ddaad2..8acff490bb7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_DeletedTeamTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DeletedTeamTable" ( "id" TEXT NOT NULL, "team_id" TEXT NOT NULL, "team_alias" TEXT, @@ -33,7 +33,7 @@ CREATE TABLE "LiteLLM_DeletedTeamTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_DeletedVerificationToken" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DeletedVerificationToken" ( "id" TEXT NOT NULL, "token" TEXT NOT NULL, "key_name" TEXT, @@ -80,38 +80,38 @@ CREATE TABLE "LiteLLM_DeletedVerificationToken" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_team_id_idx" ON "LiteLLM_DeletedTeamTable"("team_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_team_id_idx" ON "LiteLLM_DeletedTeamTable"("team_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_deleted_at_idx" ON "LiteLLM_DeletedTeamTable"("deleted_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_deleted_at_idx" ON "LiteLLM_DeletedTeamTable"("deleted_at"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_organization_id_idx" ON "LiteLLM_DeletedTeamTable"("organization_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_organization_id_idx" ON "LiteLLM_DeletedTeamTable"("organization_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_team_alias_idx" ON "LiteLLM_DeletedTeamTable"("team_alias"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_team_alias_idx" ON "LiteLLM_DeletedTeamTable"("team_alias"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_created_at_idx" ON "LiteLLM_DeletedTeamTable"("created_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_created_at_idx" ON "LiteLLM_DeletedTeamTable"("created_at"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_token_idx" ON "LiteLLM_DeletedVerificationToken"("token"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_token_idx" ON "LiteLLM_DeletedVerificationToken"("token"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_deleted_at_idx" ON "LiteLLM_DeletedVerificationToken"("deleted_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_deleted_at_idx" ON "LiteLLM_DeletedVerificationToken"("deleted_at"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_user_id_idx" ON "LiteLLM_DeletedVerificationToken"("user_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_user_id_idx" ON "LiteLLM_DeletedVerificationToken"("user_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_team_id_idx" ON "LiteLLM_DeletedVerificationToken"("team_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_team_id_idx" ON "LiteLLM_DeletedVerificationToken"("team_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_organization_id_idx" ON "LiteLLM_DeletedVerificationToken"("organization_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_organization_id_idx" ON "LiteLLM_DeletedVerificationToken"("organization_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_key_alias_idx" ON "LiteLLM_DeletedVerificationToken"("key_alias"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_key_alias_idx" ON "LiteLLM_DeletedVerificationToken"("key_alias"); -- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_created_at_idx" ON "LiteLLM_DeletedVerificationToken"("created_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_created_at_idx" ON "LiteLLM_DeletedVerificationToken"("created_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251220144550_schema_update/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251220144550_schema_update/migration.sql index b40defec309..c79ce17b9ee 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251220144550_schema_update/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251220144550_schema_update/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_SkillsTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_SkillsTable" ( "skill_id" TEXT NOT NULL, "display_title" TEXT, "description" TEXT, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql index 8eebb797e2c..a854693d57a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql @@ -1,5 +1,5 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "authorization_url" TEXT, -ADD COLUMN "registration_url" TEXT, -ADD COLUMN "token_url" TEXT; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "authorization_url" TEXT, +ADD COLUMN IF NOT EXISTS "registration_url" TEXT, +ADD COLUMN IF NOT EXISTS "token_url" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql index 8d3e02bd051..26319c5ed84 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "allow_all_keys" BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "allow_all_keys" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260106155622_add_endpoint_to_daily_activity_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260106155622_add_endpoint_to_daily_activity_tables/migration.sql index 4ed7feb9ca0..3e94f736ea3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260106155622_add_endpoint_to_daily_activity_tables/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260106155622_add_endpoint_to_daily_activity_tables/migration.sql @@ -1,72 +1,72 @@ -- DropIndex -DROP INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key"; +DROP INDEX IF EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key"; -- DropIndex -DROP INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key"; -- DropIndex -DROP INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key"; -- DropIndex -DROP INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; -- DropIndex -DROP INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key"; -- DropIndex -DROP INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key"; +DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key"; -- AlterTable -ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN "endpoint" TEXT; +ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN "endpoint" TEXT; +ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN "endpoint" TEXT; +ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN "endpoint" TEXT; +ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN "endpoint" TEXT; +ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "endpoint" TEXT; +ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT; -- CreateIndex -CREATE INDEX "LiteLLM_DailyAgentSpend_endpoint_idx" ON "LiteLLM_DailyAgentSpend"("endpoint"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_endpoint_idx" ON "LiteLLM_DailyAgentSpend"("endpoint"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyEndUserSpend_endpoint_idx" ON "LiteLLM_DailyEndUserSpend"("endpoint"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_endpoint_idx" ON "LiteLLM_DailyEndUserSpend"("endpoint"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyOrganizationSpend_endpoint_idx" ON "LiteLLM_DailyOrganizationSpend"("endpoint"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_endpoint_idx" ON "LiteLLM_DailyOrganizationSpend"("endpoint"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTagSpend_endpoint_idx" ON "LiteLLM_DailyTagSpend"("endpoint"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_endpoint_idx" ON "LiteLLM_DailyTagSpend"("endpoint"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTeamSpend_endpoint_idx" ON "LiteLLM_DailyTeamSpend"("endpoint"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_endpoint_idx" ON "LiteLLM_DailyTeamSpend"("endpoint"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyUserSpend_endpoint_idx" ON "LiteLLM_DailyUserSpend"("endpoint"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_endpoint_idx" ON "LiteLLM_DailyUserSpend"("endpoint"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260107111013_add_router_settings_to_keys_teams/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260107111013_add_router_settings_to_keys_teams/migration.sql index 95566950118..07cf2f95f90 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260107111013_add_router_settings_to_keys_teams/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260107111013_add_router_settings_to_keys_teams/migration.sql @@ -1,6 +1,6 @@ -- AlterTable -ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "router_settings" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "router_settings" JSONB DEFAULT '{}'; -- AlterTable -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "router_settings" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "router_settings" JSONB DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql index add80b39e7f..561c36530d9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql @@ -6,4 +6,4 @@ -- With this index, PostgreSQL can use an Index Scan for O(log n) performance. -- -- Related: GitHub Issue #18411 -CREATE INDEX "LiteLLM_UserTable_user_email_lower_idx" ON "LiteLLM_UserTable"(LOWER("user_email")); +CREATE INDEX IF NOT EXISTS "LiteLLM_UserTable_user_email_lower_idx" ON "LiteLLM_UserTable"(LOWER("user_email")); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260116142756_update_deleted_keys_teams_table_routing_settings/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260116142756_update_deleted_keys_teams_table_routing_settings/migration.sql index 9426bed0da2..b948ffcad66 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260116142756_update_deleted_keys_teams_table_routing_settings/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260116142756_update_deleted_keys_teams_table_routing_settings/migration.sql @@ -1,6 +1,6 @@ -- AlterTable -ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "router_settings" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "router_settings" JSONB DEFAULT '{}'; -- AlterTable -ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "router_settings" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "router_settings" JSONB DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260123131407_add_policy_tables_and_policies_field/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260123131407_add_policy_tables_and_policies_field/migration.sql index 595d8f4a0c5..8b6ac7a7f3e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260123131407_add_policy_tables_and_policies_field/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260123131407_add_policy_tables_and_policies_field/migration.sql @@ -1,20 +1,20 @@ -- AlterTable -ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- AlterTable -ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- AlterTable -ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- AlterTable -ALTER TABLE "LiteLLM_UserTable" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- AlterTable -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- CreateTable -CREATE TABLE "LiteLLM_PolicyTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_PolicyTable" ( "policy_id" TEXT NOT NULL, "policy_name" TEXT NOT NULL, "inherit" TEXT, @@ -31,7 +31,7 @@ CREATE TABLE "LiteLLM_PolicyTable" ( ); -- CreateTable -CREATE TABLE "LiteLLM_PolicyAttachmentTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_PolicyAttachmentTable" ( "attachment_id" TEXT NOT NULL, "policy_name" TEXT NOT NULL, "scope" TEXT, @@ -47,5 +47,5 @@ CREATE TABLE "LiteLLM_PolicyAttachmentTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_PolicyTable_policy_name_key" ON "LiteLLM_PolicyTable"("policy_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_PolicyTable_policy_name_key" ON "LiteLLM_PolicyTable"("policy_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql index 51d88444191..58b2d4048c2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_DeprecatedVerificationToken" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DeprecatedVerificationToken" ( "id" TEXT NOT NULL, "token" TEXT NOT NULL, "active_token_id" TEXT NOT NULL, @@ -10,10 +10,10 @@ CREATE TABLE "LiteLLM_DeprecatedVerificationToken" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_DeprecatedVerificationToken_token_key" ON "LiteLLM_DeprecatedVerificationToken"("token"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DeprecatedVerificationToken_token_key" ON "LiteLLM_DeprecatedVerificationToken"("token"); -- CreateIndex -CREATE INDEX "LiteLLM_DeprecatedVerificationToken_token_revoke_at_idx" ON "LiteLLM_DeprecatedVerificationToken"("token", "revoke_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeprecatedVerificationToken_token_revoke_at_idx" ON "LiteLLM_DeprecatedVerificationToken"("token", "revoke_at"); -- CreateIndex -CREATE INDEX "LiteLLM_DeprecatedVerificationToken_revoke_at_idx" ON "LiteLLM_DeprecatedVerificationToken"("revoke_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DeprecatedVerificationToken_revoke_at_idx" ON "LiteLLM_DeprecatedVerificationToken"("revoke_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205091235_allow_team_guardrail_config/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205091235_allow_team_guardrail_config/migration.sql index 000b96b3b87..53ef243f7d0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205091235_allow_team_guardrail_config/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205091235_allow_team_guardrail_config/migration.sql @@ -1,6 +1,6 @@ -- AlterTable -ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false; -- AlterTable -ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql index a64f1de342f..6a3c3aa66db 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "soft_budget" DOUBLE PRECISION; +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "soft_budget" DOUBLE PRECISION; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207093506_add_available_on_public_internet_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207093506_add_available_on_public_internet_to_mcp_servers/migration.sql index 1efde3dbe0f..9b969c1e3b7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207093506_add_available_on_public_internet_to_mcp_servers/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207093506_add_available_on_public_internet_to_mcp_servers/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "available_on_public_internet" BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "available_on_public_internet" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207110613_add_soft_budget_to_deleted_teams_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207110613_add_soft_budget_to_deleted_teams_table/migration.sql index abfb153061b..229032e2b13 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207110613_add_soft_budget_to_deleted_teams_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207110613_add_soft_budget_to_deleted_teams_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "soft_budget" DOUBLE PRECISION; +ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "soft_budget" DOUBLE PRECISION; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql index 572eea9b529..92baf0d25e4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql @@ -1,8 +1,8 @@ -- CreateIndex -CREATE INDEX "LiteLLM_VerificationToken_user_id_team_id_idx" ON "LiteLLM_VerificationToken"("user_id", "team_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_VerificationToken_user_id_team_id_idx" ON "LiteLLM_VerificationToken"("user_id", "team_id"); -- CreateIndex -CREATE INDEX "LiteLLM_VerificationToken_team_id_idx" ON "LiteLLM_VerificationToken"("team_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_VerificationToken_team_id_idx" ON "LiteLLM_VerificationToken"("team_id"); -- CreateIndex -CREATE INDEX "LiteLLM_VerificationToken_budget_reset_at_expires_idx" ON "LiteLLM_VerificationToken"("budget_reset_at", "expires"); +CREATE INDEX IF NOT EXISTS "LiteLLM_VerificationToken_budget_reset_at_expires_idx" ON "LiteLLM_VerificationToken"("budget_reset_at", "expires"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212103349_adjust_tags_policy_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212103349_adjust_tags_policy_table/migration.sql index f3a0821d37f..de03c010bf4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212103349_adjust_tags_policy_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212103349_adjust_tags_policy_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN "tags" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "tags" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql index 67e75e84c4a..b3ae05e202b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql @@ -1,17 +1,17 @@ -- AlterTable -ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- AlterTable -ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- AlterTable -ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- AlterTable -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; -- CreateTable -CREATE TABLE "LiteLLM_AccessGroupTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_AccessGroupTable" ( "access_group_id" TEXT NOT NULL, "access_group_name" TEXT NOT NULL, "description" TEXT, @@ -29,5 +29,5 @@ CREATE TABLE "LiteLLM_AccessGroupTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_AccessGroupTable_access_group_name_key" ON "LiteLLM_AccessGroupTable"("access_group_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AccessGroupTable_access_group_name_key" ON "LiteLLM_AccessGroupTable"("access_group_name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213105436_add_managed_vector_store_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213105436_add_managed_vector_store_table/migration.sql index 0835875220f..eb3dd90602d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213105436_add_managed_vector_store_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213105436_add_managed_vector_store_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_ManagedVectorStoreTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedVectorStoreTable" ( "id" TEXT NOT NULL, "unified_resource_id" TEXT NOT NULL, "resource_object" JSONB, @@ -16,7 +16,7 @@ CREATE TABLE "LiteLLM_ManagedVectorStoreTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_ManagedVectorStoreTable_unified_resource_id_key" ON "LiteLLM_ManagedVectorStoreTable"("unified_resource_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoreTable_unified_resource_id_key" ON "LiteLLM_ManagedVectorStoreTable"("unified_resource_id"); -- CreateIndex -CREATE INDEX "LiteLLM_ManagedVectorStoreTable_unified_resource_id_idx" ON "LiteLLM_ManagedVectorStoreTable"("unified_resource_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoreTable_unified_resource_id_idx" ON "LiteLLM_ManagedVectorStoreTable"("unified_resource_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213170952_access_group_change_to_model_name/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213170952_access_group_change_to_model_name/migration.sql index c940d3aca8b..921bee97752 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213170952_access_group_change_to_model_name/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213170952_access_group_change_to_model_name/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_AccessGroupTable" DROP COLUMN "access_model_ids", -ADD COLUMN "access_model_names" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_AccessGroupTable" DROP COLUMN IF EXISTS "access_model_ids", +ADD COLUMN IF NOT EXISTS "access_model_names" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214094754_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214094754_schema_sync/migration.sql index b5d5b978580..31e4320bb17 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214094754_schema_sync/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214094754_schema_sync/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN "team_id" TEXT; +ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN IF NOT EXISTS "team_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214163027_add_pipeline_to_policy_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214163027_add_pipeline_to_policy_table/migration.sql index e57b9ef29c5..bb121466c9a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214163027_add_pipeline_to_policy_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214163027_add_pipeline_to_policy_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_PolicyTable" ADD COLUMN "pipeline" JSONB; +ALTER TABLE "LiteLLM_PolicyTable" ADD COLUMN IF NOT EXISTS "pipeline" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql index 5c5dc6fd6f1..4b909471e02 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql @@ -1,6 +1,11 @@ -- AlterTable -ALTER TABLE "LiteLLM_EndUserTable" ADD COLUMN "object_permission_id" TEXT; +ALTER TABLE "LiteLLM_EndUserTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT; -- AddForeignKey -ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_EndUserTable_object_permission_id_fkey') THEN + ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260218231534_add_last_active_to_key_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260218231534_add_last_active_to_key_table/migration.sql index ded1856059b..bfd50fe33be 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260218231534_add_last_active_to_key_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260218231534_add_last_active_to_key_table/migration.sql @@ -1,6 +1,6 @@ -- AlterTable -ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "last_active" TIMESTAMP(3); +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "last_active" TIMESTAMP(3); -- AlterTable -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "last_active" TIMESTAMP(3); +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "last_active" TIMESTAMP(3); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219105005_add_project_id_to_deleted_keys/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219105005_add_project_id_to_deleted_keys/migration.sql index 59bdc86adbb..ef1ae258562 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219105005_add_project_id_to_deleted_keys/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219105005_add_project_id_to_deleted_keys/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "project_id" TEXT; +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "project_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219181415_baseline_diff/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219181415_baseline_diff/migration.sql index dd95d9d84a3..13be156db8e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219181415_baseline_diff/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219181415_baseline_diff/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_DailyGuardrailMetrics" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGuardrailMetrics" ( "guardrail_id" TEXT NOT NULL, "date" TEXT NOT NULL, "requests_evaluated" BIGINT NOT NULL DEFAULT 0, @@ -15,7 +15,7 @@ CREATE TABLE "LiteLLM_DailyGuardrailMetrics" ( ); -- CreateTable -CREATE TABLE "LiteLLM_DailyPolicyMetrics" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyPolicyMetrics" ( "policy_id" TEXT NOT NULL, "date" TEXT NOT NULL, "requests_evaluated" BIGINT NOT NULL DEFAULT 0, @@ -31,7 +31,7 @@ CREATE TABLE "LiteLLM_DailyPolicyMetrics" ( ); -- CreateTable -CREATE TABLE "LiteLLM_SpendLogGuardrailIndex" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogGuardrailIndex" ( "request_id" TEXT NOT NULL, "guardrail_id" TEXT NOT NULL, "policy_id" TEXT, @@ -41,20 +41,20 @@ CREATE TABLE "LiteLLM_SpendLogGuardrailIndex" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_DailyGuardrailMetrics_date_idx" ON "LiteLLM_DailyGuardrailMetrics"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGuardrailMetrics_date_idx" ON "LiteLLM_DailyGuardrailMetrics"("date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyGuardrailMetrics_guardrail_id_idx" ON "LiteLLM_DailyGuardrailMetrics"("guardrail_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGuardrailMetrics_guardrail_id_idx" ON "LiteLLM_DailyGuardrailMetrics"("guardrail_id"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyPolicyMetrics_date_idx" ON "LiteLLM_DailyPolicyMetrics"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyPolicyMetrics_date_idx" ON "LiteLLM_DailyPolicyMetrics"("date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyPolicyMetrics_policy_id_idx" ON "LiteLLM_DailyPolicyMetrics"("policy_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyPolicyMetrics_policy_id_idx" ON "LiteLLM_DailyPolicyMetrics"("policy_id"); -- CreateIndex -CREATE INDEX "LiteLLM_SpendLogGuardrailIndex_guardrail_id_start_time_idx" ON "LiteLLM_SpendLogGuardrailIndex"("guardrail_id", "start_time"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogGuardrailIndex_guardrail_id_start_time_idx" ON "LiteLLM_SpendLogGuardrailIndex"("guardrail_id", "start_time"); -- CreateIndex -CREATE INDEX "LiteLLM_SpendLogGuardrailIndex_policy_id_start_time_idx" ON "LiteLLM_SpendLogGuardrailIndex"("policy_id", "start_time"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogGuardrailIndex_policy_id_start_time_idx" ON "LiteLLM_SpendLogGuardrailIndex"("policy_id", "start_time"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220124742_add_spec_path_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220124742_add_spec_path_to_mcp_servers/migration.sql index 4f4e72a8798..55c12d9476f 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220124742_add_spec_path_to_mcp_servers/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220124742_add_spec_path_to_mcp_servers/migration.sql @@ -1,2 +1,2 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "spec_path" TEXT; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "spec_path" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220153844_add_composite_index_aggregate_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220153844_add_composite_index_aggregate_tables/migration.sql index a10f123b02e..a95a3c5c9c5 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220153844_add_composite_index_aggregate_tables/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220153844_add_composite_index_aggregate_tables/migration.sql @@ -1,36 +1,36 @@ -- DropIndex -DROP INDEX "LiteLLM_DailyAgentSpend_agent_id_idx"; +DROP INDEX IF EXISTS "LiteLLM_DailyAgentSpend_agent_id_idx"; -- DropIndex -DROP INDEX "LiteLLM_DailyEndUserSpend_end_user_id_idx"; +DROP INDEX IF EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_idx"; -- DropIndex -DROP INDEX "LiteLLM_DailyOrganizationSpend_organization_id_idx"; +DROP INDEX IF EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_idx"; -- DropIndex -DROP INDEX "LiteLLM_DailyTagSpend_tag_idx"; +DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_idx"; -- DropIndex -DROP INDEX "LiteLLM_DailyTeamSpend_team_id_idx"; +DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_idx"; -- DropIndex -DROP INDEX "LiteLLM_DailyUserSpend_user_id_idx"; +DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_idx"; -- CreateIndex -CREATE INDEX "LiteLLM_DailyAgentSpend_agent_id_date_idx" ON "LiteLLM_DailyAgentSpend"("agent_id", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_idx" ON "LiteLLM_DailyAgentSpend"("agent_id", "date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_idx" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_idx" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_idx" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_idx" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTagSpend_tag_date_idx" ON "LiteLLM_DailyTagSpend"("tag", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_idx" ON "LiteLLM_DailyTagSpend"("tag", "date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyTeamSpend_team_id_date_idx" ON "LiteLLM_DailyTeamSpend"("team_id", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_idx" ON "LiteLLM_DailyTeamSpend"("team_id", "date"); -- CreateIndex -CREATE INDEX "LiteLLM_DailyUserSpend_user_id_date_idx" ON "LiteLLM_DailyUserSpend"("user_id", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_idx" ON "LiteLLM_DailyUserSpend"("user_id", "date"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221183800_add_policy_versioning/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221183800_add_policy_versioning/migration.sql index 087c5ecc01a..67647da66b0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221183800_add_policy_versioning/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221183800_add_policy_versioning/migration.sql @@ -1,17 +1,17 @@ -- DropIndex -DROP INDEX "LiteLLM_PolicyTable_policy_name_key"; +DROP INDEX IF EXISTS "LiteLLM_PolicyTable_policy_name_key"; -- AlterTable -ALTER TABLE "LiteLLM_PolicyTable" ADD COLUMN "is_latest" BOOLEAN NOT NULL DEFAULT true, -ADD COLUMN "parent_version_id" TEXT, -ADD COLUMN "production_at" TIMESTAMP(3), -ADD COLUMN "published_at" TIMESTAMP(3), -ADD COLUMN "version_number" INTEGER NOT NULL DEFAULT 1, -ADD COLUMN "version_status" TEXT NOT NULL DEFAULT 'production'; +ALTER TABLE "LiteLLM_PolicyTable" ADD COLUMN IF NOT EXISTS "is_latest" BOOLEAN NOT NULL DEFAULT true, +ADD COLUMN IF NOT EXISTS "parent_version_id" TEXT, +ADD COLUMN IF NOT EXISTS "production_at" TIMESTAMP(3), +ADD COLUMN IF NOT EXISTS "published_at" TIMESTAMP(3), +ADD COLUMN IF NOT EXISTS "version_number" INTEGER NOT NULL DEFAULT 1, +ADD COLUMN IF NOT EXISTS "version_status" TEXT NOT NULL DEFAULT 'production'; -- CreateIndex -CREATE INDEX "LiteLLM_PolicyTable_policy_name_version_status_idx" ON "LiteLLM_PolicyTable"("policy_name", "version_status"); +CREATE INDEX IF NOT EXISTS "LiteLLM_PolicyTable_policy_name_version_status_idx" ON "LiteLLM_PolicyTable"("policy_name", "version_status"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_PolicyTable_policy_name_version_number_key" ON "LiteLLM_PolicyTable"("policy_name", "version_number"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_PolicyTable_policy_name_version_number_key" ON "LiteLLM_PolicyTable"("policy_name", "version_number"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql index ac390d164d3..361a4705e62 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql @@ -1,3 +1,3 @@ -- Add batch_processed column to LiteLLM_ManagedObjectTable -- Set to true by CheckBatchCost after cost has been computed for a completed batch -ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN "batch_processed" BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "batch_processed" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224201417_spend_logs_request_duration/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224201417_spend_logs_request_duration/migration.sql index 892aa59e9f8..aa0237eae1f 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224201417_spend_logs_request_duration/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224201417_spend_logs_request_duration/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "request_duration_ms" INTEGER; +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "request_duration_ms" INTEGER; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql index 78e364d5478..ecfe47a44d3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql @@ -1,14 +1,14 @@ -- AlterTable -ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "object_permission_id" TEXT; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT; -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "spec_path"; +ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN IF EXISTS "spec_path"; -- AlterTable -ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "agent_id" TEXT; +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "agent_id" TEXT; -- CreateTable -CREATE TABLE "LiteLLM_ToolTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ToolTable" ( "tool_id" TEXT NOT NULL, "tool_name" TEXT NOT NULL, "origin" TEXT, @@ -27,14 +27,19 @@ CREATE TABLE "LiteLLM_ToolTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_ToolTable_tool_name_key" ON "LiteLLM_ToolTable"("tool_name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ToolTable_tool_name_key" ON "LiteLLM_ToolTable"("tool_name"); -- CreateIndex -CREATE INDEX "LiteLLM_ToolTable_call_policy_idx" ON "LiteLLM_ToolTable"("call_policy"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ToolTable_call_policy_idx" ON "LiteLLM_ToolTable"("call_policy"); -- CreateIndex -CREATE INDEX "LiteLLM_ToolTable_team_id_idx" ON "LiteLLM_ToolTable"("team_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ToolTable_team_id_idx" ON "LiteLLM_ToolTable"("team_id"); -- AddForeignKey -ALTER TABLE "LiteLLM_AgentsTable" ADD CONSTRAINT "LiteLLM_AgentsTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_AgentsTable_object_permission_id_fkey') THEN + ALTER TABLE "LiteLLM_AgentsTable" ADD CONSTRAINT "LiteLLM_AgentsTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql index cba06684193..24724cb18e0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql @@ -1,2 +1,2 @@ -- AlterTable -ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "blocked_tools" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "blocked_tools" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql index e3199679ce2..d9c234696c8 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_SpendLogToolIndex" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogToolIndex" ( "request_id" TEXT NOT NULL, "tool_name" TEXT NOT NULL, "start_time" TIMESTAMP(3) NOT NULL, @@ -8,4 +8,4 @@ CREATE TABLE "LiteLLM_SpendLogToolIndex" ( ); -- CreateIndex -CREATE INDEX "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx" ON "LiteLLM_SpendLogToolIndex"("tool_name", "start_time"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx" ON "LiteLLM_SpendLogToolIndex"("tool_name", "start_time"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql index 594ab9ac1a2..a41160dac6a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "agent_id" TEXT; +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "agent_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql index e2a3694e8ef..7fa354d9d80 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql @@ -1,5 +1,5 @@ -- CreateTable -CREATE TABLE "LiteLLM_ClaudeCodePluginTable" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginTable" ( "id" TEXT NOT NULL, "name" TEXT NOT NULL, "version" TEXT, @@ -15,4 +15,4 @@ CREATE TABLE "LiteLLM_ClaudeCodePluginTable" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_ClaudeCodePluginTable_name_key" ON "LiteLLM_ClaudeCodePluginTable"("name"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginTable_name_key" ON "LiteLLM_ClaudeCodePluginTable"("name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql index b347a8d5895..66d792fe386 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql @@ -1,2 +1,2 @@ -- CreateIndex -CREATE INDEX "LiteLLM_SpendLogs_startTime_request_id_idx" ON "LiteLLM_SpendLogs"("startTime", "request_id"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" ON "LiteLLM_SpendLogs"("startTime", "request_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228170127_support_team_based_guardrails/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228170127_support_team_based_guardrails/migration.sql index 8af167950ec..44d079ad194 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228170127_support_team_based_guardrails/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228170127_support_team_based_guardrails/migration.sql @@ -1,8 +1,8 @@ -- AlterTable -ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN "reviewed_at" TIMESTAMP(3), -ADD COLUMN "status" TEXT NOT NULL DEFAULT 'active', -ADD COLUMN "submitted_at" TIMESTAMP(3); +ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN IF NOT EXISTS "reviewed_at" TIMESTAMP(3), +ADD COLUMN IF NOT EXISTS "status" TEXT NOT NULL DEFAULT 'active', +ADD COLUMN IF NOT EXISTS "submitted_at" TIMESTAMP(3); -- CreateIndex -CREATE INDEX "LiteLLM_GuardrailsTable_status_idx" ON "LiteLLM_GuardrailsTable"("status"); +CREATE INDEX IF NOT EXISTS "LiteLLM_GuardrailsTable_status_idx" ON "LiteLLM_GuardrailsTable"("status"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260303000000_update_tool_table_policies/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260303000000_update_tool_table_policies/migration.sql index 2e2d722ed4c..7aa329c6bb1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260303000000_update_tool_table_policies/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260303000000_update_tool_table_policies/migration.sql @@ -1,20 +1,25 @@ --- Rename call_policy to input_policy -ALTER TABLE "LiteLLM_ToolTable" RENAME COLUMN "call_policy" TO "input_policy"; +-- Rename call_policy to input_policy (only if the old name still exists) +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'LiteLLM_ToolTable' AND column_name = 'call_policy') THEN + ALTER TABLE "LiteLLM_ToolTable" RENAME COLUMN "call_policy" TO "input_policy"; + END IF; +END $$; -- Add output_policy column -ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN "output_policy" TEXT NOT NULL DEFAULT 'untrusted'; +ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN IF NOT EXISTS "output_policy" TEXT NOT NULL DEFAULT 'untrusted'; -- Add user_agent column -ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN "user_agent" TEXT; +ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN IF NOT EXISTS "user_agent" TEXT; -- Add last_used_at column -ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN "last_used_at" TIMESTAMP(3); +ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN IF NOT EXISTS "last_used_at" TIMESTAMP(3); -- Drop old index on call_policy DROP INDEX IF EXISTS "LiteLLM_ToolTable_call_policy_idx"; -- CreateIndex -CREATE INDEX "LiteLLM_ToolTable_input_policy_idx" ON "LiteLLM_ToolTable"("input_policy"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ToolTable_input_policy_idx" ON "LiteLLM_ToolTable"("input_policy"); -- CreateIndex -CREATE INDEX "LiteLLM_ToolTable_output_policy_idx" ON "LiteLLM_ToolTable"("output_policy"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ToolTable_output_policy_idx" ON "LiteLLM_ToolTable"("output_policy"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql index 01f3936a6fc..a045b7d1d66 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql index 3cd8ca638a4..c1556822c1a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql @@ -1,5 +1,5 @@ -- AlterTable -ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "tpm_limit" INTEGER; -ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "rpm_limit" INTEGER; -ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "session_tpm_limit" INTEGER; -ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "session_rpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "tpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "rpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "session_tpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "session_rpm_limit" INTEGER; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306175056_add_configs_override_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306175056_add_configs_override_table/migration.sql index aad5e2b3889..616463a9e2e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306175056_add_configs_override_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306175056_add_configs_override_table/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "spec_path" TEXT; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "spec_path" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306233848_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306233848_schema_sync/migration.sql index 6395d5b1f8b..4f3c0b7485b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306233848_schema_sync/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306233848_schema_sync/migration.sql @@ -1,12 +1,12 @@ -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "byok_api_key_help_url" TEXT, -ADD COLUMN "byok_description" TEXT[] DEFAULT ARRAY[]::TEXT[], -ADD COLUMN "is_byok" BOOLEAN NOT NULL DEFAULT false, -ADD COLUMN "tool_name_to_description" JSONB DEFAULT '{}', -ADD COLUMN "tool_name_to_display_name" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "byok_api_key_help_url" TEXT, +ADD COLUMN IF NOT EXISTS "byok_description" TEXT[] DEFAULT ARRAY[]::TEXT[], +ADD COLUMN IF NOT EXISTS "is_byok" BOOLEAN NOT NULL DEFAULT false, +ADD COLUMN IF NOT EXISTS "tool_name_to_description" JSONB DEFAULT '{}', +ADD COLUMN IF NOT EXISTS "tool_name_to_display_name" JSONB DEFAULT '{}'; -- CreateTable -CREATE TABLE "LiteLLM_MCPUserCredentials" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_MCPUserCredentials" ( "id" TEXT NOT NULL, "user_id" TEXT NOT NULL, "server_id" TEXT NOT NULL, @@ -18,7 +18,7 @@ CREATE TABLE "LiteLLM_MCPUserCredentials" ( ); -- CreateTable -CREATE TABLE "LiteLLM_JWTKeyMapping" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_JWTKeyMapping" ( "id" TEXT NOT NULL, "jwt_claim_name" TEXT NOT NULL, "jwt_claim_value" TEXT NOT NULL, @@ -34,7 +34,7 @@ CREATE TABLE "LiteLLM_JWTKeyMapping" ( ); -- CreateTable -CREATE TABLE "LiteLLM_ConfigOverrides" ( +CREATE TABLE IF NOT EXISTS "LiteLLM_ConfigOverrides" ( "config_type" TEXT NOT NULL, "config_value" JSONB NOT NULL, "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, @@ -44,14 +44,19 @@ CREATE TABLE "LiteLLM_ConfigOverrides" ( ); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_MCPUserCredentials_user_id_server_id_key" ON "LiteLLM_MCPUserCredentials"("user_id", "server_id"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MCPUserCredentials_user_id_server_id_key" ON "LiteLLM_MCPUserCredentials"("user_id", "server_id"); -- CreateIndex -CREATE INDEX "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value", "is_active"); +CREATE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value", "is_active"); -- CreateIndex -CREATE UNIQUE INDEX "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value"); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value"); -- AddForeignKey -ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE RESTRICT ON UPDATE CASCADE; +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN + ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE RESTRICT ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql index 5ab834695b8..84eb70ce097 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql @@ -1,11 +1,11 @@ -- DropIndex -DROP INDEX "LiteLLM_MCPServerTable_approval_status_idx"; +DROP INDEX IF EXISTS "LiteLLM_MCPServerTable_approval_status_idx"; -- AlterTable -ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "approval_status", -DROP COLUMN "review_notes", -DROP COLUMN "reviewed_at", -DROP COLUMN "source_url", -DROP COLUMN "submitted_at", -DROP COLUMN "submitted_by"; +ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN IF EXISTS "approval_status", +DROP COLUMN IF EXISTS "review_notes", +DROP COLUMN IF EXISTS "reviewed_at", +DROP COLUMN IF EXISTS "source_url", +DROP COLUMN IF EXISTS "submitted_at", +DROP COLUMN IF EXISTS "submitted_by"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql index 8854fd1e205..cc48a742f20 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql @@ -1,3 +1,3 @@ -- AlterTable -ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "models" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "models" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f3155722187..7eff0c00f75 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -415,16 +415,26 @@ class ProxyExtrasDBManager: logger.info( f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied" ) - ProxyExtrasDBManager._roll_back_migration( - failed_migration - ) - ProxyExtrasDBManager._resolve_specific_migration( - failed_migration - ) - logger.info( - f"✅ Migration {failed_migration} resolved." - ) - return True + try: + ProxyExtrasDBManager._roll_back_migration( + failed_migration + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as rollback_err: + logger.warning( + f"Failed to roll back migration {failed_migration}: {rollback_err}. " + f"It may already be in a rolled-back state." + ) + try: + ProxyExtrasDBManager._resolve_specific_migration( + failed_migration + ) + logger.info( + f"✅ Migration {failed_migration} resolved, retrying to apply remaining migrations" + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as resolve_err: + logger.warning( + f"Failed to resolve migration {failed_migration}: {resolve_err}" + ) else: logger.info( f"Found failed migration: {failed_migration}, marking as rolled back" @@ -514,20 +524,34 @@ class ProxyExtrasDBManager: ) if migration_match: migration_name = migration_match.group(1) - logger.info( - f"Rolling back migration {migration_name}" - ) - ProxyExtrasDBManager._roll_back_migration( - migration_name - ) - logger.info( - f"Resolving migration {migration_name} that failed " - f"due to existing schema objects" - ) - ProxyExtrasDBManager._resolve_specific_migration( - migration_name - ) - logger.info("✅ Migration resolved.") + try: + logger.info( + f"Rolling back migration {migration_name}" + ) + ProxyExtrasDBManager._roll_back_migration( + migration_name + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as rollback_err: + logger.warning( + f"Failed to roll back migration {migration_name}: {rollback_err}. " + f"It may already be in a rolled-back state." + ) + try: + logger.info( + f"Resolving migration {migration_name} that failed " + f"due to existing schema objects" + ) + ProxyExtrasDBManager._resolve_specific_migration( + migration_name + ) + logger.info( + f"✅ Migration {migration_name} resolved, " + f"retrying to apply remaining migrations" + ) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as resolve_err: + logger.warning( + f"Failed to resolve migration {migration_name}: {resolve_err}" + ) else: # Unknown P3018 error - log and re-raise for safety logger.warning( diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 153c3c2dba0..0f6cc592ff5 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -556,6 +556,13 @@ class ProxyInitializationHelpers: help="Restart worker after this many requests (uvicorn: limit_max_requests, gunicorn: max_requests)", envvar="MAX_REQUESTS_BEFORE_RESTART", ) +@click.option( + "--skip_db_migration_check", + is_flag=True, + default=False, + help="Warn and continue instead of exiting when database migration fails.", + envvar="SKIP_DB_MIGRATION_CHECK", +) def run_server( # noqa: PLR0915 host, port, @@ -595,6 +602,7 @@ def run_server( # noqa: PLR0915 skip_server_startup, keepalive_timeout, max_requests_before_restart, + skip_db_migration_check: bool, ): args = locals() if local: @@ -857,11 +865,17 @@ def run_server( # noqa: PLR0915 if not PrismaManager.setup_database( use_migrate=not use_prisma_db_push ): - print( # noqa - "\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. " - "The proxy cannot start safely. Please check your database connection and migration status.\033[0m" - ) - sys.exit(1) + if skip_db_migration_check: + print( # noqa + "\033[1;33mLiteLLM Proxy: Database migration failed but continuing startup. " + "Pass --skip_db_migration_check to allow this.\033[0m" + ) + else: + print( # noqa + "\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. " + "The proxy cannot start safely. Please check your database connection and migration status.\033[0m" + ) + sys.exit(1) else: print( # noqa f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 597c0845d43..c1e3fd5072f 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -1,6 +1,10 @@ +import glob import os +import re import sys +import pytest + sys.path.insert( 0, os.path.abspath( @@ -10,6 +14,14 @@ sys.path.insert( from litellm_proxy_extras.utils import ProxyExtrasDBManager +# Path to the migrations directory +_MIGRATIONS_DIR = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "../../litellm-proxy-extras/litellm_proxy_extras/migrations", + ) +) + def test_custom_prisma_dir(monkeypatch): import tempfile @@ -138,3 +150,160 @@ class TestErrorClassificationPriority: error_message = "connection timeout" assert ProxyExtrasDBManager._is_permission_error(error_message) is False assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False + + +def _get_all_migrations(): + """Return (migration_name, sql_content) pairs for all migrations.""" + migration_files = sorted(glob.glob(os.path.join(_MIGRATIONS_DIR, "*/migration.sql"))) + results = [] + for path in migration_files: + migration_name = os.path.basename(os.path.dirname(path)) + with open(path) as f: + results.append((migration_name, f.read())) + return results + + +class TestMigrationSQLIdempotency: + """Ensure all migration SQL files use idempotent DDL (IF [NOT] EXISTS). + + Migrations on pre-existing instances can fail when DDL statements assume + the target object doesn't already exist (or still exists for drops). + These tests enforce that all migrations use safe, re-runnable SQL patterns. + """ + + @pytest.fixture(scope="class") + def all_migrations(self): + migrations = _get_all_migrations() + assert len(migrations) > 0, ( + f"No migrations found. " + f"Check that _MIGRATIONS_DIR ({_MIGRATIONS_DIR}) is correct." + ) + return migrations + + def test_create_table_uses_if_not_exists(self, all_migrations): + """CREATE TABLE statements must use IF NOT EXISTS""" + violations = [] + for migration_name, sql in all_migrations: + for line_num, line in enumerate(sql.splitlines(), 1): + if re.search(r"CREATE\s+TABLE\s+", line, re.IGNORECASE) and not re.search( + r"CREATE\s+TABLE\s+IF\s+NOT\s+EXISTS", line, re.IGNORECASE + ): + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "CREATE TABLE without IF NOT EXISTS found in migrations:\n" + + "\n".join(violations) + ) + + def test_add_column_uses_if_not_exists(self, all_migrations): + """ADD COLUMN statements must use IF NOT EXISTS""" + violations = [] + for migration_name, sql in all_migrations: + for line_num, line in enumerate(sql.splitlines(), 1): + if re.search(r"ADD\s+COLUMN\s+", line, re.IGNORECASE) and not re.search( + r"ADD\s+COLUMN\s+IF\s+NOT\s+EXISTS", line, re.IGNORECASE + ): + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "ADD COLUMN without IF NOT EXISTS found in recent migrations:\n" + + "\n".join(violations) + ) + + def test_drop_column_uses_if_exists(self, all_migrations): + """DROP COLUMN statements must use IF EXISTS""" + violations = [] + for migration_name, sql in all_migrations: + for line_num, line in enumerate(sql.splitlines(), 1): + if re.search(r"DROP\s+COLUMN\s+", line, re.IGNORECASE) and not re.search( + r"DROP\s+COLUMN\s+IF\s+EXISTS", line, re.IGNORECASE + ): + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "DROP COLUMN without IF EXISTS found in recent migrations:\n" + + "\n".join(violations) + ) + + def test_drop_index_uses_if_exists(self, all_migrations): + """DROP INDEX statements must use IF EXISTS""" + violations = [] + for migration_name, sql in all_migrations: + for line_num, line in enumerate(sql.splitlines(), 1): + if re.search(r"DROP\s+INDEX\s+", line, re.IGNORECASE) and not re.search( + r"DROP\s+INDEX\s+IF\s+EXISTS", line, re.IGNORECASE + ): + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "DROP INDEX without IF EXISTS found in recent migrations:\n" + + "\n".join(violations) + ) + + def test_create_index_uses_if_not_exists(self, all_migrations): + """CREATE INDEX statements must use IF NOT EXISTS""" + violations = [] + for migration_name, sql in all_migrations: + for line_num, line in enumerate(sql.splitlines(), 1): + if re.search( + r"CREATE\s+(?:UNIQUE\s+)?INDEX\s+", line, re.IGNORECASE + ) and not re.search( + r"CREATE\s+(?:UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?IF\s+NOT\s+EXISTS", + line, + re.IGNORECASE, + ): + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "CREATE INDEX without IF NOT EXISTS found in recent migrations:\n" + + "\n".join(violations) + ) + + def test_rename_column_is_guarded(self, all_migrations): + """RENAME COLUMN must be inside a DO $$ IF EXISTS block""" + violations = [] + for migration_name, sql in all_migrations: + lines = sql.splitlines() + in_do_block = False + for line_num, line in enumerate(lines, 1): + if re.search(r"DO\s+\$\$", line, re.IGNORECASE): + in_do_block = True + if re.search(r"END\s+\$\$", line, re.IGNORECASE): + in_do_block = False + if re.search(r"RENAME\s+COLUMN\s+", line, re.IGNORECASE) and not in_do_block: + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "RENAME COLUMN without DO $$ IF EXISTS guard found in migrations:\n" + + "\n".join(violations) + ) + + def test_add_constraint_is_guarded(self, all_migrations): + """ADD CONSTRAINT must be inside a DO $$ IF NOT EXISTS block""" + violations = [] + for migration_name, sql in all_migrations: + lines = sql.splitlines() + in_do_block = False + for line_num, line in enumerate(lines, 1): + if re.search(r"DO\s+\$\$", line, re.IGNORECASE): + in_do_block = True + if re.search(r"END\s+\$\$", line, re.IGNORECASE): + in_do_block = False + if re.search(r"ADD\s+CONSTRAINT\s+", line, re.IGNORECASE) and not in_do_block: + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "ADD CONSTRAINT without DO $$ IF NOT EXISTS guard found in migrations:\n" + + "\n".join(violations) + ) + + def test_drop_constraint_is_guarded(self, all_migrations): + """DROP CONSTRAINT must be inside a DO $$ IF EXISTS block""" + violations = [] + for migration_name, sql in all_migrations: + lines = sql.splitlines() + in_do_block = False + for line_num, line in enumerate(lines, 1): + if re.search(r"DO\s+\$\$", line, re.IGNORECASE): + in_do_block = True + if re.search(r"END\s+\$\$", line, re.IGNORECASE): + in_do_block = False + if re.search(r"DROP\s+CONSTRAINT\s+", line, re.IGNORECASE) and not in_do_block: + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "DROP CONSTRAINT without DO $$ IF EXISTS guard found in migrations:\n" + + "\n".join(violations) + )