From 26226d475f7a2074b638e9d5dc3f10682e3b58b5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 14 Mar 2025 19:34:06 -0700 Subject: [PATCH 1/5] feat(proxy_server.py): support retrieving models for a team, if user is a member - via `/models?team_id` Allows user to see team models on UI when creating a key --- .../proxy/_experimental/out/onboarding.html | 1 - litellm/proxy/auth/model_checks.py | 8 ++--- .../management_endpoints/team_endpoints.py | 11 ++----- litellm/proxy/proxy_server.py | 32 +++++++++++++++---- 4 files changed, 32 insertions(+), 20 deletions(-) delete mode 100644 litellm/proxy/_experimental/out/onboarding.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index e95db384ad..0000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 041d23bf0b..a48ef6ae87 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -85,7 +85,7 @@ def get_key_models( def get_team_models( - user_api_key_dict: UserAPIKeyAuth, + team_models: List[str], proxy_model_list: List[str], model_access_groups: Dict[str, List[str]], ) -> List[str]: @@ -96,10 +96,10 @@ def get_team_models( - If model_access_groups is provided, only return models that are in the access groups """ all_models = [] - if len(user_api_key_dict.team_models) > 0: - all_models = user_api_key_dict.team_models + if len(team_models) > 0: + all_models = team_models if SpecialModelNames.all_team_models.value in all_models: - all_models = user_api_key_dict.team_models + all_models = team_models if SpecialModelNames.all_proxy_models.value in all_models: all_models = proxy_model_list diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 81f66421a7..52b2f43282 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1352,9 +1352,9 @@ async def team_info( else: _team_info = LiteLLM_TeamTable() - ## UNFURL 'all-proxy-models' into the team_info.models list ## - if llm_router is not None: - _team_info = _unfurl_all_proxy_models(_team_info, llm_router) + # ## UNFURL 'all-proxy-models' into the team_info.models list ## + # if llm_router is not None: + # _team_info = _unfurl_all_proxy_models(_team_info, llm_router) response_object = TeamInfoResponseObject( team_id=team_id, team_info=_team_info, @@ -1615,11 +1615,6 @@ async def list_team( ) try: - # unfurl all-proxy-models - if llm_router is not None: - team = _unfurl_all_proxy_models( - LiteLLM_TeamTable(**team.model_dump()), llm_router - ) returned_responses.append( TeamListResponseObject( **team.model_dump(), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 250551b019..a40522fd22 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -122,7 +122,7 @@ from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router -from litellm.proxy.auth.auth_checks import log_db_metrics +from litellm.proxy.auth.auth_checks import get_team_object, log_db_metrics from litellm.proxy.auth.auth_utils import check_response_size_is_safe from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import LicenseCheck @@ -213,7 +213,10 @@ from litellm.proxy.management_endpoints.team_callback_endpoints import ( router as team_callback_router, ) from litellm.proxy.management_endpoints.team_endpoints import router as team_router -from litellm.proxy.management_endpoints.team_endpoints import update_team +from litellm.proxy.management_endpoints.team_endpoints import ( + update_team, + validate_membership, +) from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) @@ -3380,13 +3383,14 @@ class ProxyStartupEvent: async def model_list( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), return_wildcard_routes: Optional[bool] = False, + team_id: Optional[str] = None, ): """ Use `/model/info` - to get detailed model information, example - pricing, mode, etc. This is just for compatibility with openai projects like aider. """ - global llm_model_list, general_settings, llm_router + global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj all_models = [] model_access_groups: Dict[str, List[str]] = defaultdict(list) ## CHECK IF MODEL RESTRICTIONS ARE SET AT KEY/TEAM LEVEL ## @@ -3401,19 +3405,33 @@ async def model_list( model_access_groups=model_access_groups, ) + team_models: List[str] = user_api_key_dict.team_models + + if team_id: + team_object = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) + team_models = team_object.models + team_models = get_team_models( - user_api_key_dict=user_api_key_dict, + team_models=team_models, proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, ) + all_models = get_complete_model_list( - key_models=key_models, + key_models=key_models if not team_models else [], team_models=team_models, proxy_model_list=proxy_model_list, user_model=user_model, infer_model_from_keys=general_settings.get("infer_model_from_keys", False), return_wildcard_routes=return_wildcard_routes, ) + return dict( data=[ { @@ -6117,7 +6135,7 @@ async def model_info_v1( # noqa: PLR0915 model_access_groups=model_access_groups, ) team_models = get_team_models( - user_api_key_dict=user_api_key_dict, + team_models=user_api_key_dict.team_models, proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, ) @@ -6344,7 +6362,7 @@ async def model_group_info( model_access_groups=model_access_groups, ) team_models = get_team_models( - user_api_key_dict=user_api_key_dict, + team_models=user_api_key_dict.team_models, proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, ) From aaaac25e2dc7953db4dbe423b8ee31e48363c704 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 14 Mar 2025 19:50:49 -0700 Subject: [PATCH 2/5] fix(create_key_button.tsx): retrieve all available team models on team select --- .../src/components/create_key_button.tsx | 54 +++++++++++-------- .../src/components/networking.tsx | 13 ++++- 2 files changed, 42 insertions(+), 25 deletions(-) diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index 205d4226dc..b460fccb67 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/create_key_button.tsx @@ -91,28 +91,31 @@ const getPredefinedTags = (data: any[] | null) => { return uniqueTags; } -export const getTeamModels = (team: Team | null, allAvailableModels: string[]): string[] => { - let tempModelsToPick = []; - - if (team) { - if (team.models.length > 0) { - if (team.models.includes("all-proxy-models")) { - // if the team has all-proxy-models show all available models - tempModelsToPick = allAvailableModels; - } else { - // show team models - tempModelsToPick = team.models; - } - } else { - // show all available models if the team has no models set - tempModelsToPick = allAvailableModels; +export const fetchTeamModels = async (userID: string, userRole: string, accessToken: string, teamID: string): Promise => { + try { + if (userID === null || userRole === null) { + return []; } - } else { - // no team set, show all available models - tempModelsToPick = allAvailableModels; - } - return unfurlWildcardModelsInList(tempModelsToPick, allAvailableModels); + if (accessToken !== null) { + const model_available = await modelAvailableCall( + accessToken, + userID, + userRole, + true, + teamID + ); + let available_model_names = model_available["data"].map( + (element: { id: string }) => element.id + ); + console.log("available_model_names:", available_model_names); + return available_model_names; + } + return []; + } catch (error) { + console.error("Error fetching user models:", error); + return []; + } }; export const fetchUserModels = async (userID: string, userRole: string, accessToken: string, setUserModels: (models: string[]) => void) => { @@ -182,6 +185,7 @@ const CreateKey: React.FC = ({ } }, [accessToken, userID, userRole]); + useEffect(() => { const fetchGuardrails = async () => { try { @@ -277,10 +281,14 @@ const CreateKey: React.FC = ({ }; useEffect(() => { - const models = getTeamModels(selectedCreateKeyTeam, userModels); - setModelsToPick(models); + if (userID && userRole && accessToken && selectedCreateKeyTeam) { + fetchTeamModels(userID, userRole, accessToken, selectedCreateKeyTeam.team_id).then((models) => { + let allModels = Array.from(new Set([...selectedCreateKeyTeam.models, ...models])); + setModelsToPick(allModels); + }); + } form.setFieldValue('models', []); - }, [selectedCreateKeyTeam, userModels]); + }, [selectedCreateKeyTeam]); // Add a callback function to handle user creation const handleUserCreated = (userId: string) => { diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 831ff7d102..c30cf727b0 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1284,6 +1284,7 @@ export const modelInfoV1Call = async (accessToken: String, modelId: String) => { } }; + export const modelHubCall = async (accessToken: String) => { /** * Get all models on proxy @@ -1581,7 +1582,8 @@ export const modelAvailableCall = async ( accessToken: String, userID: String, userRole: String, - return_wildcard_routes: boolean = false + return_wildcard_routes: boolean = false, + teamID: String | null = null ) => { /** * Get all the models user has access to @@ -1589,8 +1591,15 @@ export const modelAvailableCall = async ( console.log("in /models calls, globalLitellmHeaderName", globalLitellmHeaderName) try { let url = proxyBaseUrl ? `${proxyBaseUrl}/models` : `/models`; + const params = new URLSearchParams(); if (return_wildcard_routes === true) { - url += `?return_wildcard_routes=True`; + params.append('return_wildcard_routes', 'True'); + } + if (teamID) { + params.append('team_id', teamID.toString()); + } + if (params.toString()) { + url += `?${params.toString()}`; } //message.info("Requesting model data"); From 6ed995952f66fd022368d4a2a2578133d395e0d5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 14 Mar 2025 20:28:50 -0700 Subject: [PATCH 3/5] fix: fix test --- tests/proxy_unit_tests/test_proxy_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 909647cac2..dccf0d1842 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -939,8 +939,9 @@ def test_get_team_models(): model_access_groups["default"].extend(["gpt-4o-mini"]) model_access_groups["team2"].extend(["gpt-3.5-turbo"]) + team_models = user_api_key_dict.team_models result = get_team_models( - user_api_key_dict=user_api_key_dict, + team_models=team_models, proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, ) From a93ef15b5531deb2a3026daca50498cdcc5d009e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 14 Mar 2025 20:48:30 -0700 Subject: [PATCH 4/5] fix(key_edit_view.tsx): fix showing available models on key edit for non-team key --- .../src/components/key_edit_view.tsx | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/ui/litellm-dashboard/src/components/key_edit_view.tsx b/ui/litellm-dashboard/src/components/key_edit_view.tsx index e6fcadf3fc..bbe75e0d0e 100644 --- a/ui/litellm-dashboard/src/components/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/key_edit_view.tsx @@ -2,7 +2,7 @@ import React, { useState, useEffect } from "react"; import { Form, Input, InputNumber, Select } from "antd"; import { Button, TextInput } from "@tremor/react"; import { KeyResponse } from "./key_team_helpers/key_list"; -import { getTeamModels } from "../components/create_key_button"; +import { fetchTeamModels } from "../components/create_key_button"; import { modelAvailableCall } from "./networking"; interface KeyEditViewProps { @@ -45,31 +45,36 @@ export function KeyEditView({ const [form] = Form.useForm(); const [userModels, setUserModels] = useState([]); const team = teams?.find(team => team.team_id === keyData.team_id); - const availableModels = getTeamModels(team, userModels); - + const [availableModels, setAvailableModels] = useState([]); useEffect(() => { - const fetchUserModels = async () => { + const fetchModels = async () => { + if (!userID || !userRole || !accessToken) return; + try { - if (accessToken && userID && userRole) { + if (keyData.team_id === null) { + // Fetch user models if no team const model_available = await modelAvailableCall( accessToken, - userID, + userID, userRole ); - let available_model_names = model_available["data"].map( + const available_model_names = model_available["data"].map( (element: { id: string }) => element.id ); - console.log("available_model_names:", available_model_names); - setUserModels(available_model_names); + setAvailableModels(available_model_names); + } else if (team?.team_id) { + // Fetch team models if team exists + const models = await fetchTeamModels(userID, userRole, accessToken, team.team_id); + setAvailableModels(Array.from(new Set([...team.models, ...models]))); } } catch (error) { - console.error("Error fetching user models:", error); + console.error("Error fetching models:", error); } }; - fetchUserModels(); - }, []); + fetchModels(); + }, [userID, userRole, accessToken, team, keyData.team_id]); // Convert API budget duration to form format const getBudgetDuration = (duration: string | null) => { From b15c06ee946efce1b38df25a1ff1259f2b7ea8f5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 14 Mar 2025 20:51:21 -0700 Subject: [PATCH 5/5] fix(team_endpoints.py): fix linting error --- litellm/proxy/management_endpoints/team_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 52b2f43282..1994e27ecf 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1279,7 +1279,7 @@ async def team_info( --header 'Authorization: Bearer your_api_key_here' ``` """ - from litellm.proxy.proxy_server import llm_router, prisma_client + from litellm.proxy.proxy_server import prisma_client try: if prisma_client is None: @@ -1556,7 +1556,7 @@ async def list_team( - 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 llm_router, prisma_client + from litellm.proxy.proxy_server import prisma_client if not allowed_route_check_inside_route( user_api_key_dict=user_api_key_dict, requested_user_id=user_id