From 2ab00bd12cedeb0f64dda0e4d981baf475e2fc8f Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 26 Apr 2025 18:24:49 -0700 Subject: [PATCH] Litellm UI improvements 04 26 2025 p1 (#10346) * fix(provider_info_helpers.tsx): fix together ai provider name * fix(user_edit_view.tsx): all admin to edit user's personal models * fix(model_dashboard.tsx): fix model filtering by team id check * fix(proxy_server.py): support returning available models for a user on `/model_group/info` Fixes issue where personal key with 'all team models' would see all proxy models (not just the ones user can call) * test(test_models.py): add unit test for model group info check on personal key --- litellm/proxy/_new_secret_config.yaml | 6 +++ litellm/proxy/proxy_server.py | 37 ++++++++++++++++--- tests/test_models.py | 17 +++++++++ .../src/components/model_dashboard.tsx | 2 +- .../src/components/provider_info_helpers.tsx | 2 +- .../src/components/user_edit_view.tsx | 2 +- 6 files changed, 58 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 976895d1a1..0c78332a8f 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,4 +1,10 @@ model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: azure/chatgpt-v-3 + api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + api_version: "2023-05-15" + api_key: os.environ/AZURE_API_KEY - model_name: "gpt-4o-azure" litellm_params: model: azure/gpt-4o diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5252cb9a97..966d79c607 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -217,6 +217,7 @@ from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.management_endpoints.budget_management_endpoints import ( router as budget_management_router, ) +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.management_endpoints.customer_endpoints import ( router as customer_router, ) @@ -6376,6 +6377,7 @@ async def model_group_info( raise HTTPException( status_code=500, detail={"error": "LLM Router is not loaded in"} ) + ## CHECK IF MODEL RESTRICTIONS ARE SET AT KEY/TEAM LEVEL ## model_access_groups: Dict[str, List[str]] = defaultdict(list) if llm_router is None: @@ -6389,11 +6391,36 @@ async def model_group_info( proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, ) - team_models = get_team_models( - team_models=user_api_key_dict.team_models, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - ) + team_models = [] + if ( + not user_api_key_dict.team_id + and user_api_key_dict.user_id is not None + and not _user_has_admin_view(user_api_key_dict) + ): + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + user_object = await prisma_client.db.litellm_usertable.find_first( + where={"user_id": user_api_key_dict.user_id} + ) + user_object_typed = LiteLLM_UserTable(**user_object.model_dump()) + user_models = [] + if user_object is not None: + user_models = get_team_models( + team_models=user_object_typed.models, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + team_models = user_models + else: + team_models = get_team_models( + team_models=user_api_key_dict.team_models, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + all_models_str = get_complete_model_list( key_models=key_models, team_models=team_models, diff --git a/tests/test_models.py b/tests/test_models.py index 89944c07b3..eaad07052c 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -446,6 +446,23 @@ async def test_add_model_run_health(): # cleanup await delete_model(session=session, model_id=model_id) +@pytest.mark.asyncio +async def test_get_personal_models_for_user(): + """ + Test /models endpoint with team + """ + from test_users import new_user + async with aiohttp.ClientSession() as session: + # Creat a user + user_data = await new_user(session=session, i=0, models=["gpt-3.5-turbo"]) + user_id = user_data["user_id"] + user_api_key = user_data["key"] + + model_group_info = await get_model_group_info(session=session, key=user_api_key) + print(model_group_info) + + assert len(model_group_info["data"]) == 1 + assert model_group_info["data"][0]["model_group"] == "gpt-3.5-turbo" @pytest.mark.asyncio async def test_model_group_info_e2e(): diff --git a/ui/litellm-dashboard/src/components/model_dashboard.tsx b/ui/litellm-dashboard/src/components/model_dashboard.tsx index 1db4b965c5..eaa2a324b4 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard.tsx @@ -1211,7 +1211,7 @@ const ModelDashboard: React.FC = ({ data={modelData.data.filter( (model: any) => ( (selectedModelGroup === "all" || model.model_name === selectedModelGroup || !selectedModelGroup) && - (selectedTeamFilter === "all" || model.team_id === selectedTeamFilter || !selectedTeamFilter) + (selectedTeamFilter === "all" || model.model_info["team_id"] === selectedTeamFilter || !selectedTeamFilter) ) )} isLoading={false} diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index c3112db055..6b675aa617 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -51,7 +51,7 @@ export const provider_map: Record = { Cerebras: "cerebras", Sambanova: "sambanova", Perplexity: "perplexity", - TogetherAI: "togetherai", + TogetherAI: "together_ai", Openrouter: "openrouter", FireworksAI: "fireworks_ai" }; diff --git a/ui/litellm-dashboard/src/components/user_edit_view.tsx b/ui/litellm-dashboard/src/components/user_edit_view.tsx index 9c0fcbced9..6511d74d1a 100644 --- a/ui/litellm-dashboard/src/components/user_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/user_edit_view.tsx @@ -120,7 +120,7 @@ export function UserEditView({ mode="multiple" placeholder="Select models" style={{ width: "100%" }} - disabled={!all_admin_roles.includes(userData.user_info?.user_role || "")} + disabled={!all_admin_roles.includes(userRole || "")} >