diff --git a/.circleci/config.yml b/.circleci/config.yml index 1af15b03a5..e5e015f519 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -72,6 +72,7 @@ jobs: pip install "jsonschema==4.22.0" pip install "pytest-xdist==3.6.1" pip install "websockets==10.4" + pip uninstall posthog -y - save_cache: paths: - ./venv diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7b2435e67c..75494bcf78 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1042,6 +1042,9 @@ class LiteLLM_TeamTable(TeamBase): "model_aliases", ] + if isinstance(values, BaseModel): + values = values.model_dump() + if ( isinstance(values.get("members_with_roles"), dict) and not values["members_with_roles"] @@ -1489,6 +1492,7 @@ class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase): class LiteLLM_OrganizationTableWithMembers(LiteLLM_OrganizationTable): members: List[LiteLLM_OrganizationMembershipTable] + teams: List[LiteLLM_TeamTable] class NewOrganizationResponse(LiteLLM_OrganizationTable): @@ -2424,3 +2428,7 @@ class PrismaCompatibleUpdateDBModel(TypedDict, total=False): model_info: str updated_at: str updated_by: str + + +class SpecialManagementEndpointEnums(enum.Enum): + DEFAULT_ORGANIZATION = "default_organization" diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e1efa23df6..f0b8c84683 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -168,6 +168,8 @@ def _team_key_generation_check( user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest, ): + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return True if ( litellm.key_generation_settings is not None and "team_key_generation" in litellm.key_generation_settings diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 34cee0cc7e..c6bd51c0e7 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -227,7 +227,7 @@ async def list_organization( # if proxy admin - get all orgs if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: response = await prisma_client.db.litellm_organizationtable.find_many( - include={"members": True} + include={"members": True, "teams": True} ) # if internal user - get orgs they are a member of else: @@ -242,7 +242,7 @@ async def list_organization( "in": [membership.organization_id for membership in org_memberships] } }, - include={"members": True}, + include={"members": True, "teams": True}, ) response = org_objects diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b57411d556..35fbfe433e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -39,6 +39,7 @@ from litellm.proxy._types import ( NewTeamRequest, ProxyErrorTypes, ProxyException, + SpecialManagementEndpointEnums, TeamAddMemberResponse, TeamInfoResponseObject, TeamListResponseObject, @@ -1482,6 +1483,7 @@ async def list_team( user_id: Optional[str] = fastapi.Query( default=None, description="Only return teams which this 'user_id' belongs to" ), + organization_id: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1492,6 +1494,7 @@ async def list_team( Parameters: - user_id: str - Optional. If passed will only return teams that the user_id is a member of. + - organization_id: str - Optional. If passed will only return teams that belong to the organization_id. Pass 'default_organization' to get all teams without organization_id. """ from litellm.proxy.proxy_server import prisma_client @@ -1565,6 +1568,19 @@ async def list_team( continue # Sort the responses by team_alias returned_responses.sort(key=lambda x: (getattr(x, "team_alias", "") or "")) + + if organization_id is not None: + if organization_id == SpecialManagementEndpointEnums.DEFAULT_ORGANIZATION.value: + returned_responses = [ + team for team in returned_responses if team.organization_id is None + ] + else: + returned_responses = [ + team + for team in returned_responses + if team.organization_id == organization_id + ] + return returned_responses diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index df337c94d0..a09ef97009 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -56,6 +56,7 @@ model LiteLLM_OrganizationTable { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) teams LiteLLM_TeamTable[] users LiteLLM_UserTable[] + keys LiteLLM_VerificationToken[] members LiteLLM_OrganizationMembership[] @relation("OrganizationToMembership") } @@ -158,9 +159,11 @@ model LiteLLM_VerificationToken { model_spend Json @default("{}") model_max_budget Json @default("{}") budget_id String? + organization_id String? created_at DateTime? @default(now()) @map("created_at") updated_at DateTime? @default(now()) @updatedAt @map("updated_at") litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) } model LiteLLM_EndUserTable { diff --git a/schema.prisma b/schema.prisma index df337c94d0..a09ef97009 100644 --- a/schema.prisma +++ b/schema.prisma @@ -56,6 +56,7 @@ model LiteLLM_OrganizationTable { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) teams LiteLLM_TeamTable[] users LiteLLM_UserTable[] + keys LiteLLM_VerificationToken[] members LiteLLM_OrganizationMembership[] @relation("OrganizationToMembership") } @@ -158,9 +159,11 @@ model LiteLLM_VerificationToken { model_spend Json @default("{}") model_max_budget Json @default("{}") budget_id String? + organization_id String? created_at DateTime? @default(now()) @map("created_at") updated_at DateTime? @default(now()) @updatedAt @map("updated_at") litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) } model LiteLLM_EndUserTable { diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index a08ae1c005..5894df2c71 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -183,6 +183,7 @@ export default function CreateKeyPage() { } } } + setTeams(null); } return ( @@ -201,6 +202,7 @@ export default function CreateKeyPage() { setTeams={setTeams} setKeys={setKeys} setOrganizations={setOrganizations} + currentOrg={currentOrg} /> ) : (
@@ -237,6 +239,7 @@ export default function CreateKeyPage() { setTeams={setTeams} setKeys={setKeys} setOrganizations={setOrganizations} + currentOrg={currentOrg} /> ) : page == "models" ? ( void) => { + let givenTeams; + if (userRole != "Admin" && userRole != "Admin Viewer") { + givenTeams = await teamListCall(accessToken, currentOrg?.organization_id || DEFAULT_ORGANIZATION, userID) + } else { + givenTeams = await teamListCall(accessToken, currentOrg?.organization_id || DEFAULT_ORGANIZATION) + } + + console.log(`givenTeams: ${givenTeams}`) + + setTeams(givenTeams) + } \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/fetch_available_models_team_key.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/fetch_available_models_team_key.tsx index 6a34ff0cb4..d42399179c 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/fetch_available_models_team_key.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/fetch_available_models_team_key.tsx @@ -44,7 +44,6 @@ export const fetchAvailableModelsForTeamOrKey = async ( }; export const getModelDisplayName = (model: string) => { - console.log("getModelDisplayName", model); if (model.endsWith('/*')) { const provider = model.replace('/*', ''); return `All ${provider} models`; diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index a6ce2d812b..e304f931a3 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -59,46 +59,45 @@ const Navbar: React.FC = ({ ]; const orgMenuItems: MenuProps["items"] = [ - ...(userRole === "Admin" - ? [{ - key: 'global', - label: ( -
- Global View -
- ), - onClick: () => onOrgChange({ organization_id: "global", organization_alias: "Global View" } as Organization) - }] - : [ - { - key: 'header', - label: 'Organizations', - type: 'group' as const, - style: { - color: '#6B7280', - fontSize: '0.875rem' - } - }, - ...organizations.map(org => ({ - key: org.organization_id, - label: ( -
- {org.organization_alias} -
- ), - onClick: () => onOrgChange(org) - })), - { - key: "note", - label: ( -
- Switching between organizations on the UI is currently in beta. -
- ), - disabled: true - } - ] - ) + { + key: 'header', + label: 'Organizations', + type: 'group' as const, + style: { + color: '#6B7280', + fontSize: '0.875rem' + } + }, + { + key: "default", + label: ( +
+ Default Organization +
+ ), + onClick: () => onOrgChange({ + organization_id: null, + organization_alias: "Default Organization" + } as Organization) + }, + ...organizations.filter(org => org.organization_id !== null).map(org => ({ + key: org.organization_id ?? "default", + label: ( +
+ {org.organization_alias} +
+ ), + onClick: () => onOrgChange(org) + })), + { + key: "note", + label: ( +
+ Switching between organizations on the UI is currently in beta. +
+ ), + disabled: true + } ]; return ( diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 58bbd9885c..1df3f228f7 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -9,6 +9,8 @@ if (isLocal != true) { console.log = function() {}; } +export const DEFAULT_ORGANIZATION = "default_organization"; + export interface Model { model_name: string; litellm_params: Object; @@ -16,7 +18,7 @@ export interface Model { } export interface Organization { - organization_id: string; + organization_id: string | null; organization_alias: string; budget_id: string; metadata: Record; @@ -724,7 +726,8 @@ export const teamInfoCall = async ( export const teamListCall = async ( accessToken: String, - userID: String | null = null + organizationID: string | null, + userID: String | null = null, ) => { /** * Get all available teams on proxy @@ -732,9 +735,21 @@ export const teamListCall = async ( try { let url = proxyBaseUrl ? `${proxyBaseUrl}/team/list` : `/team/list`; console.log("in teamInfoCall"); + const queryParams = new URLSearchParams(); + if (userID) { - url += `?user_id=${userID}`; + queryParams.append('user_id', userID.toString()); } + + if (organizationID) { + queryParams.append('organization_id', organizationID.toString()); + } + + const queryString = queryParams.toString(); + if (queryString) { + url += `?${queryString}`; + } + const response = await fetch(url, { method: "GET", headers: { diff --git a/ui/litellm-dashboard/src/components/teams.tsx b/ui/litellm-dashboard/src/components/teams.tsx index 433987eb1a..7cb7d440a4 100644 --- a/ui/litellm-dashboard/src/components/teams.tsx +++ b/ui/litellm-dashboard/src/components/teams.tsx @@ -1,8 +1,9 @@ import React, { useState, useEffect } from "react"; import Link from "next/link"; import { Typography } from "antd"; -import { teamDeleteCall, teamUpdateCall, teamInfoCall, Organization } from "./networking"; +import { teamDeleteCall, teamUpdateCall, teamInfoCall, Organization, DEFAULT_ORGANIZATION } from "./networking"; import TeamMemberModal from "@/components/team/edit_membership"; +import { fetchTeams } from "./common_components/fetch_teams"; import { InformationCircleIcon, PencilAltIcon, @@ -95,23 +96,12 @@ const Team: React.FC = ({ }) => { const [lastRefreshed, setLastRefreshed] = useState(""); - const fetchTeams = async (accessToken: string, userID: string | null, userRole: string | null) => { - let givenTeams; - if (userRole != "Admin" && userRole != "Admin Viewer") { - givenTeams = await teamListCall(accessToken, userID) - } else { - givenTeams = await teamListCall(accessToken) - } - - console.log(`givenTeams: ${givenTeams}`) - setTeams(givenTeams) - } useEffect(() => { console.log(`inside useeffect - ${teams}`) if (teams === null && accessToken) { // Call your function here - fetchTeams(accessToken, userID, userRole) + fetchTeams(accessToken, userID, userRole, currentOrg, setTeams) } }, [teams]); @@ -119,7 +109,7 @@ const Team: React.FC = ({ console.log(`inside useeffect - ${lastRefreshed}`) if (accessToken) { // Call your function here - fetchTeams(accessToken, userID, userRole) + fetchTeams(accessToken, userID, userRole, currentOrg, setTeams) } handleRefreshClick() }, [lastRefreshed]); @@ -257,9 +247,9 @@ const Team: React.FC = ({ let _team_id_to_info: Record = {}; let teamList; if (userRole != "Admin" && userRole != "Admin Viewer") { - teamList = await teamListCall(accessToken, userID) + teamList = await teamListCall(accessToken, currentOrg?.organization_id || DEFAULT_ORGANIZATION, userID) } else { - teamList = await teamListCall(accessToken) + teamList = await teamListCall(accessToken, currentOrg?.organization_id || DEFAULT_ORGANIZATION) } for (let i = 0; i < teamList.length; i++) { @@ -405,6 +395,10 @@ const Team: React.FC = ({ {teams && teams.length > 0 ? teams + .filter((team) => { + const targetOrgId = currentOrg ? currentOrg.organization_id : null; + return team.organization_id === targetOrgId; + }) .sort((a, b) => new Date(b.created_at).getTime() - new Date(a.created_at).getTime()) .map((team: any) => ( diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index a6f60e22a2..5468fab9cd 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -7,8 +7,10 @@ import { getProxyUISettings, teamListCall, Organization, - organizationListCall + organizationListCall, + DEFAULT_ORGANIZATION } from "./networking"; +import { fetchTeams } from "./common_components/fetch_teams"; import { Grid, Col, Card, Text, Title } from "@tremor/react"; import CreateKey from "./create_key_button"; import ViewKeyTable from "./view_key_table"; @@ -62,6 +64,7 @@ interface UserDashboardProps { setKeys: React.Dispatch>; setOrganizations: React.Dispatch>; premiumUser: boolean; + currentOrg: Organization | null; } type TeamInterface = { @@ -82,15 +85,15 @@ const UserDashboard: React.FC = ({ setKeys, setOrganizations, premiumUser, + currentOrg }) => { + console.log(`currentOrg in user dashboard: ${JSON.stringify(currentOrg)}`) const [userSpendData, setUserSpendData] = useState( null ); // Assuming useSearchParams() hook exists and works in your setup const searchParams = useSearchParams()!; - const viewSpend = searchParams.get("viewSpend"); - const router = useRouter(); const token = getCookie('token'); @@ -173,22 +176,11 @@ const UserDashboard: React.FC = ({ } } if (userID && accessToken && userRole && !keys && !userSpendData) { - // const cachedUserModels = sessionStorage.getItem("userModels" + userID); - // if (cachedUserModels) { - // setUserModels(JSON.parse(cachedUserModels)); - // } else { - const fetchTeams = async () => { - let givenTeams; - if (userRole != "Admin" && userRole != "Admin Viewer") { - givenTeams = await teamListCall(accessToken, userID) - } else { - givenTeams = await teamListCall(accessToken) - } - - console.log(`givenTeams: ${givenTeams}`) - - setTeams(givenTeams) - } + const cachedUserModels = sessionStorage.getItem("userModels" + userID); + if (cachedUserModels) { + setUserModels(JSON.parse(cachedUserModels)); + } else { + console.log(`currentOrg: ${JSON.stringify(currentOrg)}`) const fetchData = async () => { try { const proxy_settings: ProxySettings = await getProxyUISettings(accessToken); @@ -205,12 +197,10 @@ const UserDashboard: React.FC = ({ setUserSpendData(response["user_info"]); console.log(`userSpendData: ${JSON.stringify(userSpendData)}`) - - console.log(`response["user_info"]["organization_memberships"]: ${JSON.stringify(response["user_info"]["organization_memberships"])}`) // set keys for admin and users - if (!response?.teams?.[0]?.keys) { + if (!response?.teams[0].keys) { setKeys(response["keys"]); } else { setKeys( @@ -270,12 +260,20 @@ const UserDashboard: React.FC = ({ setOrganizations(organizations); } fetchData(); - fetchTeams(); + fetchTeams(accessToken, userID, userRole, currentOrg, setTeams); fetchOrganizations(); } - // } + } }, [userID, token, accessToken, keys, userRole]); + useEffect(() => { + console.log(`currentOrg: ${JSON.stringify(currentOrg)}, accessToken: ${accessToken}, userID: ${userID}, userRole: ${userRole}`) + if (accessToken) { + console.log(`fetching teams`) + fetchTeams(accessToken, userID, userRole, currentOrg, setTeams); + } + }, [currentOrg]); + useEffect(() => { // This code will run every time selectedTeam changes if ( @@ -375,6 +373,7 @@ const UserDashboard: React.FC = ({ setData={setKeys} premiumUser={premiumUser} teams={teams} + currentOrg={currentOrg} /> >; teams: any[] | null; premiumUser: boolean; + currentOrg: Organization | null; } interface ItemData { @@ -132,6 +134,7 @@ const ViewKeyTable: React.FC = ({ setData, teams, premiumUser, + currentOrg }) => { const [isButtonClicked, setIsButtonClicked] = useState(false); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); @@ -167,55 +170,64 @@ const ViewKeyTable: React.FC = ({ // Combine all keys that user should have access to const all_keys_to_display = React.useMemo(() => { + if (!data) return []; + + // Helper function for default team org check + const matchesDefaultTeamOrg = (key: any) => { + console.log(`Checking if key matches default team org: ${JSON.stringify(key)}, currentOrg: ${JSON.stringify(currentOrg)}`) + if (!currentOrg || currentOrg.organization_id === null) { + return !('organization_id' in key) || key.organization_id === null; + } + return key.organization_id === currentOrg.organization_id; + }; + let allKeys: any[] = []; - // If no teams, return personal keys - if (!teams || teams.length === 0) { - return data; - } + // Handle no team selected or Default Team case + if (!selectedTeam || selectedTeam.team_alias === "Default Team") { - teams.forEach((team) => { - // For default team or when user is not admin, use personal keys (data) - if (team.team_id === "default-team" || !isUserTeamAdmin(team)) { - if (selectedTeam && selectedTeam.team_id === team.team_id && data) { - allKeys = [ - ...allKeys, - ...data.filter((key) => key.team_id === team.team_id), - ]; - } - } - // For teams where user is admin, use team keys - else if (isUserTeamAdmin(team)) { - if (selectedTeam && selectedTeam.team_id === team.team_id) { - allKeys = [ - ...allKeys, - ...(data?.filter((key) => key?.team_id === team?.team_id) || []), - ]; - } - } - }); - - // If no team is selected, show all accessible keys - if ((!selectedTeam || selectedTeam.team_alias === "Default Team") && data) { - const personalKeys = data.filter( - (key) => !key.team_id || key.team_id === "default-team" + console.log(`inside personal keys`) + // Get personal keys (with org check) + const personalKeys = data.filter(key => + key.team_id == null && + matchesDefaultTeamOrg(key) ); - const adminTeamKeys = teams - .filter((team) => isUserTeamAdmin(team)) - .flatMap((team) => team.keys || []); + + console.log(`personalKeys: ${JSON.stringify(personalKeys)}`) + + // Get admin team keys (no org check) + const adminTeamKeys = data.filter(key => { + const keyTeam = teams?.find(team => team.team_id === key.team_id); + return keyTeam && isUserTeamAdmin(keyTeam) && key.team_id !== "default-team"; + }); + + console.log(`adminTeamKeys: ${JSON.stringify(adminTeamKeys)}`) + allKeys = [...personalKeys, ...adminTeamKeys]; } + // Handle specific team selected + else { + const selectedTeamData = teams?.find(t => t.team_id === selectedTeam.team_id); + if (selectedTeamData) { + const teamKeys = data.filter(key => { + if (selectedTeamData.team_id === "default-team") { + return key.team_id == null && matchesDefaultTeamOrg(key); + } + return key.team_id === selectedTeamData.team_id; + }); + allKeys = teamKeys; + } + } - // Filter out litellm-dashboard keys - allKeys = allKeys.filter((key) => key.team_id !== "litellm-dashboard"); - - // Remove duplicates based on token - const uniqueKeys = Array.from( - new Map(allKeys.map((key) => [key.token, key])).values() + // Final filtering and deduplication + return Array.from( + new Map( + allKeys + .filter(key => key.team_id !== "litellm-dashboard") + .map(key => [key.token, key]) + ).values() ); - - return uniqueKeys; - }, [data, teams, selectedTeam, userID]); + }, [data, teams, selectedTeam, currentOrg]); useEffect(() => { const calculateNewExpiryTime = (duration: string | undefined) => {