Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions litellm/proxy/db/db_spend_update_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
)
Expand Down
103 changes: 103 additions & 0 deletions tests/test_litellm/proxy/db/test_db_spend_update_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
126 changes: 126 additions & 0 deletions ui/litellm-dashboard/src/components/OldTeams.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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(
<OldTeams
teams={[
{
team_id: "1",
team_alias: "Test Team",
organization_id: "org-123",
models: ["gpt-4"],
max_budget: 100,
budget_duration: "1d",
tpm_limit: 1000,
rpm_limit: 1000,
created_at: new Date().toISOString(),
keys: [],
members_with_roles: [],
},
]}
searchParams={{}}
accessToken="test-token"
setTeams={vi.fn()}
userID="user-123"
userRole="Admin"
organizations={[]}
/>,
);

expect(getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument();
});

it("should show Default Team Settings tab for proxy_admin role", () => {
const { getByRole } = render(
<OldTeams
teams={[
{
team_id: "1",
team_alias: "Test Team",
organization_id: "org-123",
models: ["gpt-4"],
max_budget: 100,
budget_duration: "1d",
tpm_limit: 1000,
rpm_limit: 1000,
created_at: new Date().toISOString(),
keys: [],
members_with_roles: [],
},
]}
searchParams={{}}
accessToken="test-token"
setTeams={vi.fn()}
userID="user-123"
userRole="proxy_admin"
organizations={[]}
/>,
);

expect(getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument();
});

it("should not show Default Team Settings tab for proxy_admin_viewer role", () => {
const { queryByRole } = render(
<OldTeams
teams={[
{
team_id: "1",
team_alias: "Test Team",
organization_id: "org-123",
models: ["gpt-4"],
max_budget: 100,
budget_duration: "1d",
tpm_limit: 1000,
rpm_limit: 1000,
created_at: new Date().toISOString(),
keys: [],
members_with_roles: [],
},
]}
searchParams={{}}
accessToken="test-token"
setTeams={vi.fn()}
userID="user-123"
userRole="proxy_admin_viewer"
organizations={[]}
/>,
);

expect(queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument();
});

it("should not show Default Team Settings tab for Admin Viewer role", () => {
const { queryByRole } = render(
<OldTeams
teams={[
{
team_id: "1",
team_alias: "Test Team",
organization_id: "org-123",
models: ["gpt-4"],
max_budget: 100,
budget_duration: "1d",
tpm_limit: 1000,
rpm_limit: 1000,
created_at: new Date().toISOString(),
keys: [],
members_with_roles: [],
},
]}
searchParams={{}}
accessToken="test-token"
setTeams={vi.fn()}
userID="user-123"
userRole="Admin Viewer"
organizations={[]}
/>,
);

expect(queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument();
});
});
24 changes: 9 additions & 15 deletions ui/litellm-dashboard/src/components/OldTeams.tsx
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down Expand Up @@ -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";
Expand All @@ -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";

Expand Down Expand Up @@ -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[];
}
Expand Down Expand Up @@ -615,7 +609,7 @@ const Teams: React.FC<TeamProps> = ({
<div className="flex">
<Tab>Your Teams</Tab>
<Tab>Available Teams</Tab>
{isAdminRole(userRole || "") && <Tab>Default Team Settings</Tab>}
{isProxyAdminRole(userRole || "") && <Tab>Default Team Settings</Tab>}
</div>
<div className="flex items-center space-x-2">
{lastRefreshed && <Text>Last Refreshed: {lastRefreshed}</Text>}
Expand Down Expand Up @@ -1004,7 +998,7 @@ const Teams: React.FC<TeamProps> = ({
<TabPanel>
<AvailableTeamsPanel accessToken={accessToken} userID={userID} />
</TabPanel>
{isAdminRole(userRole || "") && (
{isProxyAdminRole(userRole || "") && (
<TabPanel>
<TeamSSOSettings accessToken={accessToken} userID={userID || ""} userRole={userRole || ""} />
</TabPanel>
Expand Down
41 changes: 41 additions & 0 deletions ui/litellm-dashboard/src/utils/roles.test.tsx
Original file line number Diff line number Diff line change
@@ -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);
});
});
});
4 changes: 4 additions & 0 deletions ui/litellm-dashboard/src/utils/roles.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
};
Loading