diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 187d095a5b73..6c4f3bdd4587 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -994,11 +994,11 @@ async def _update_daily_spend( # If _update_daily_spend ever gets the ability to write to multiple tables at once, the sorting # should sort by the table first. key=lambda x: ( - x[1]["date"], + x[1].get("date") or "", x[1].get(entity_id_field) or "", - x[1]["api_key"], - x[1]["model"], - x[1]["custom_llm_provider"], + x[1].get("api_key") or "", + x[1].get("model") or "", + x[1].get("custom_llm_provider") or "", ), )[:BATCH_SIZE] ) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 6dbbbdd74425..43dbc63689e2 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -221,6 +221,109 @@ async def test_update_daily_spend_sorting(): # Verify that table.upsert was called mock_table.upsert.assert_has_calls(upsert_calls) + + +@pytest.mark.asyncio +async def test_update_daily_spend_with_none_values_in_sorting_fields(): + """ + Test that _update_daily_spend handles None values in sorting fields correctly. + + This test ensures that when fields like date, api_key, model, or custom_llm_provider + are None, the sorting doesn't crash with TypeError: '<' not supported between + instances of 'NoneType' and 'str'. + """ + # Setup + mock_prisma_client = MagicMock() + mock_batcher = MagicMock() + mock_table = MagicMock() + mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher + mock_batcher.litellm_dailyuserspend = mock_table + + # Create transactions with None values in various sorting fields + daily_spend_transactions = { + "key1": { + "user_id": "user1", + "date": None, # None date + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + }, + "key2": { + "user_id": "user2", + "date": "2024-01-01", + "api_key": None, # None api_key + "model": "gpt-4", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + }, + "key3": { + "user_id": "user3", + "date": "2024-01-01", + "api_key": "test-api-key", + "model": None, # None model + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + }, + "key4": { + "user_id": "user4", + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": None, # None custom_llm_provider + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + }, + "key5": { + "user_id": None, # None entity_id + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + }, + } + + # Call the method - this should not raise TypeError + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=1, + prisma_client=mock_prisma_client, + proxy_logging_obj=MagicMock(), + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider", + ) + + # Verify that table.upsert was called (should be called 5 times, once for each transaction) + assert mock_table.upsert.call_count == 5 + + # Tag Spend Tracking Tests diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx index 4ede3fff6d18..261178191f8e 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx @@ -466,3 +466,129 @@ describe("OldTeams - helper functions", () => { }); }); }); + +describe("OldTeams - Default Team Settings tab visibility", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should show Default Team Settings tab for Admin role", () => { + const { getByRole } = render( + , + ); + + expect(getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); + }); + + it("should show Default Team Settings tab for proxy_admin role", () => { + const { getByRole } = render( + , + ); + + expect(getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument(); + }); + + it("should not show Default Team Settings tab for proxy_admin_viewer role", () => { + const { queryByRole } = render( + , + ); + + expect(queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); + }); + + it("should not show Default Team Settings tab for Admin Viewer role", () => { + const { queryByRole } = render( + , + ); + + expect(queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx index 6a8313e56b21..eefb0302a899 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.tsx @@ -1,8 +1,7 @@ import AvailableTeamsPanel from "@/components/team/available_teams"; import TeamInfoView from "@/components/team/team_info"; import TeamSSOSettings from "@/components/TeamSSOSettings"; -import { updateExistingKeys } from "@/utils/dataUtils"; -import { isAdminRole } from "@/utils/roles"; +import { isProxyAdminRole } from "@/utils/roles"; import { InfoCircleOutlined } from "@ant-design/icons"; import { ChevronDownIcon, ChevronRightIcon, PencilAltIcon, RefreshIcon, TrashIcon } from "@heroicons/react/outline"; import { @@ -31,10 +30,10 @@ import { Text, TextInput, } from "@tremor/react"; -import { Button as Button2, Form, Input, Modal, Select as Select2, Switch, Tooltip, Typography } from "antd"; +import { Button as Button2, Form, Input, Modal, Select as Select2, Tooltip, Typography } from "antd"; +import { AlertTriangleIcon, XIcon } from "lucide-react"; import React, { useEffect, useState } from "react"; import { formatNumberWithCommas } from "../utils/dataUtils"; -import DeleteResourceModal from "./common_components/DeleteResourceModal"; import { fetchTeams } from "./common_components/fetch_teams"; import ModelAliasManager from "./common_components/ModelAliasManager"; import PremiumLoggingSettings from "./common_components/PremiumLoggingSettings"; @@ -47,15 +46,7 @@ import type { KeyResponse, Team } from "./key_team_helpers/key_list"; import MCPServerSelector from "./mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "./mcp_server_management/MCPToolPermissions"; import NotificationsManager from "./molecules/notifications_manager"; -import { - Member, - Organization, - fetchMCPAccessGroups, - getGuardrailsList, - teamCreateCall, - teamDeleteCall, - v2TeamListCall, -} from "./networking"; +import { Organization, fetchMCPAccessGroups, getGuardrailsList, teamDeleteCall } from "./networking"; import NumericalInput from "./shared/numerical_input"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; @@ -85,6 +76,9 @@ interface EditTeamModalProps { onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted } +import { updateExistingKeys } from "@/utils/dataUtils"; +import { Member, teamCreateCall, v2TeamListCall } from "./networking"; + interface TeamInfo { members_with_roles: Member[]; } @@ -615,7 +609,7 @@ const Teams: React.FC = ({
Your Teams Available Teams - {isAdminRole(userRole || "") && Default Team Settings} + {isProxyAdminRole(userRole || "") && Default Team Settings}
{lastRefreshed && Last Refreshed: {lastRefreshed}} @@ -1004,7 +998,7 @@ const Teams: React.FC = ({ - {isAdminRole(userRole || "") && ( + {isProxyAdminRole(userRole || "") && ( diff --git a/ui/litellm-dashboard/src/utils/roles.test.tsx b/ui/litellm-dashboard/src/utils/roles.test.tsx new file mode 100644 index 000000000000..871bbaba072a --- /dev/null +++ b/ui/litellm-dashboard/src/utils/roles.test.tsx @@ -0,0 +1,41 @@ +import { describe, it, expect } from "vitest"; +import { isAdminRole, isProxyAdminRole } from "./roles"; + +describe("roles", () => { + describe("isAdminRole", () => { + it("should return true for all admin roles", () => { + expect(isAdminRole("Admin")).toBe(true); + expect(isAdminRole("Admin Viewer")).toBe(true); + expect(isAdminRole("proxy_admin")).toBe(true); + expect(isAdminRole("proxy_admin_viewer")).toBe(true); + expect(isAdminRole("org_admin")).toBe(true); + }); + + it("should return false for non-admin roles", () => { + expect(isAdminRole("Internal User")).toBe(false); + expect(isAdminRole("Internal Viewer")).toBe(false); + expect(isAdminRole("regular_user")).toBe(false); + expect(isAdminRole("")).toBe(false); + }); + }); + + describe("isProxyAdminRole", () => { + it("should return true for proxy_admin and Admin roles", () => { + expect(isProxyAdminRole("proxy_admin")).toBe(true); + expect(isProxyAdminRole("Admin")).toBe(true); + }); + + it("should return false for other admin roles", () => { + expect(isProxyAdminRole("Admin Viewer")).toBe(false); + expect(isProxyAdminRole("proxy_admin_viewer")).toBe(false); + expect(isProxyAdminRole("org_admin")).toBe(false); + }); + + it("should return false for non-admin roles", () => { + expect(isProxyAdminRole("Internal User")).toBe(false); + expect(isProxyAdminRole("Internal Viewer")).toBe(false); + expect(isProxyAdminRole("regular_user")).toBe(false); + expect(isProxyAdminRole("")).toBe(false); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/roles.ts b/ui/litellm-dashboard/src/utils/roles.ts index f542da105f5a..da9b00082e3a 100644 --- a/ui/litellm-dashboard/src/utils/roles.ts +++ b/ui/litellm-dashboard/src/utils/roles.ts @@ -11,3 +11,7 @@ export const rolesWithWriteAccess = ["Internal User", "Admin", "proxy_admin"]; export const isAdminRole = (role: string): boolean => { return all_admin_roles.includes(role); }; + +export const isProxyAdminRole = (role: string): boolean => { + return role === "proxy_admin" || role === "Admin"; +};