diff --git a/litellm/proxy/_experimental/out/_next/static/chunks/webpack-a426aae3231a8df1.js b/litellm/proxy/_experimental/out/_next/static/chunks/webpack-a426aae3231a8df1.js index c82df116cd..2170290813 100644 --- a/litellm/proxy/_experimental/out/_next/static/chunks/webpack-a426aae3231a8df1.js +++ b/litellm/proxy/_experimental/out/_next/static/chunks/webpack-a426aae3231a8df1.js @@ -1 +1 @@ -!function(){"use strict";var e,t,n,r,o,u,i,c,f,a={},l={};function d(e){var t=l[e];if(void 0!==t)return t.exports;var n=l[e]={id:e,loaded:!1,exports:{}},r=!0;try{a[e].call(n.exports,n,n.exports,d),r=!1}finally{r&&delete l[e]}return n.loaded=!0,n.exports}d.m=a,e=[],d.O=function(t,n,r,o){if(n){o=o||0;for(var u=e.length;u>0&&e[u-1][2]>o;u--)e[u]=e[u-1];e[u]=[n,r,o];return}for(var i=1/0,u=0;u=o&&Object.keys(d.O).every(function(e){return d.O[e](n[f])})?n.splice(f--,1):(c=!1,o0&&e[u-1][2]>o;u--)e[u]=e[u-1];e[u]=[n,r,o];return}for(var i=1/0,u=0;u=o&&Object.keys(d.O).every(function(e){return d.O[e](n[f])})?n.splice(f--,1):(c=!1,o - - - - - - - - - - - - - + + + + + + + + + + + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/cerebras.svg b/litellm/proxy/_experimental/out/assets/logos/cerebras.svg index 426f6430c2..1ff347220c 100644 --- a/litellm/proxy/_experimental/out/assets/logos/cerebras.svg +++ b/litellm/proxy/_experimental/out/assets/logos/cerebras.svg @@ -1,89 +1,89 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/deepseek.svg b/litellm/proxy/_experimental/out/assets/logos/deepseek.svg index c4754047da..61760f1319 100644 --- a/litellm/proxy/_experimental/out/assets/logos/deepseek.svg +++ b/litellm/proxy/_experimental/out/assets/logos/deepseek.svg @@ -1,25 +1,25 @@ - - - - - - + + + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/perplexity-ai.svg b/litellm/proxy/_experimental/out/assets/logos/perplexity-ai.svg index e828b6dfbf..e3a32be980 100644 --- a/litellm/proxy/_experimental/out/assets/logos/perplexity-ai.svg +++ b/litellm/proxy/_experimental/out/assets/logos/perplexity-ai.svg @@ -1,16 +1,16 @@ - - - - - - - - - - + + + + + + + + + + \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 8ecf14a6ea..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/_types.py b/litellm/proxy/_types.py index 4546a3d579..6592efca0b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -837,6 +837,7 @@ class SpecialMCPServerName(str, enum.Enum): all_team_servers = "all-team-mcpservers" all_proxy_servers = "all-proxy-mcpservers" + # MCP Proxy Request Types class NewMCPServerRequest(LiteLLMPydanticObjectBase): server_id: Optional[str] = None @@ -1111,6 +1112,9 @@ class NewTeamRequest(TeamBase): tags: Optional[list] = None guardrails: Optional[List[str]] = None object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + team_member_budget: Optional[float] = ( + None # allow user to set a budget for all team members + ) model_config = ConfigDict(protected_namespaces=()) @@ -1152,6 +1156,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): model_aliases: Optional[dict] = None guardrails: Optional[List[str]] = None object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + team_member_budget: Optional[float] = None class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): @@ -2495,6 +2500,7 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): user_id: str team_id: str budget_id: Optional[str] = None + spend: Optional[float] = 0.0 litellm_budget_table: Optional[LiteLLM_BudgetTable] @@ -2646,9 +2652,13 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse): ########################################## +class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): + team_member_budget_table: Optional[LiteLLM_BudgetTable] = None + + class TeamInfoResponseObject(TypedDict): team_id: str - team_info: LiteLLM_TeamTable + team_info: TeamInfoResponseObjectTeamTable keys: List team_memberships: List[LiteLLM_TeamMembership] diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 959ff64f69..e9f7bda681 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -610,23 +610,23 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, ) if _end_user_object is not None: - end_user_params[ - "allowed_model_region" - ] = _end_user_object.allowed_model_region + end_user_params["allowed_model_region"] = ( + _end_user_object.allowed_model_region + ) if _end_user_object.litellm_budget_table is not None: budget_info = _end_user_object.litellm_budget_table if budget_info.tpm_limit is not None: - end_user_params[ - "end_user_tpm_limit" - ] = budget_info.tpm_limit + end_user_params["end_user_tpm_limit"] = ( + budget_info.tpm_limit + ) if budget_info.rpm_limit is not None: - end_user_params[ - "end_user_rpm_limit" - ] = budget_info.rpm_limit + end_user_params["end_user_rpm_limit"] = ( + budget_info.rpm_limit + ) if budget_info.max_budget is not None: - end_user_params[ - "end_user_max_budget" - ] = budget_info.max_budget + end_user_params["end_user_max_budget"] = ( + budget_info.max_budget + ) except Exception as e: if isinstance(e, litellm.BudgetExceededError): raise e @@ -904,6 +904,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Check 3. Check if user is in their team budget if valid_token.team_member_spend is not None: + if prisma_client is not None: _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" @@ -926,6 +927,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 await user_api_key_cache.async_set_cache( key=_cache_key, value=team_member_info, + ttl=5, ) if ( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3c24036ebe..8d7e12a6f8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -45,6 +45,7 @@ from litellm.proxy._types import ( SpecialProxyStrings, TeamAddMemberResponse, TeamInfoResponseObject, + TeamInfoResponseObjectTeamTable, TeamListResponseObject, TeamMemberAddRequest, TeamMemberDeleteRequest, @@ -132,6 +133,93 @@ async def get_all_team_memberships( return returned_tm +async def _create_team_member_budget_table( + data: Union[NewTeamRequest, LiteLLM_TeamTable], + new_team_data_json: dict, + user_api_key_dict: UserAPIKeyAuth, + team_member_budget: float, +) -> dict: + """Allows admin to create 1 budget, that applies to all team members""" + from litellm.proxy._types import BudgetNewRequest + from litellm.proxy.management_endpoints.budget_management_endpoints import ( + new_budget, + ) + + if data.team_alias is not None: + budget_id = ( + f"team-{data.team_alias.replace(' ', '-')}-budget-{uuid.uuid4().hex}" + ) + else: + budget_id = f"team-budget-{uuid.uuid4().hex}" + + team_member_budget_table = await new_budget( + budget_obj=BudgetNewRequest( + max_budget=team_member_budget, + budget_duration=data.budget_duration, + budget_id=budget_id, + ), + user_api_key_dict=user_api_key_dict, + ) + + # Add team_member_budget_id as metadata field to team table + if new_team_data_json.get("metadata") is None: + new_team_data_json["metadata"] = {} + new_team_data_json["metadata"][ + "team_member_budget_id" + ] = team_member_budget_table.budget_id + new_team_data_json.pop( + "team_member_budget", None + ) # remove team_member_budget from new_team_data_json + + return new_team_data_json + + +async def _upsert_team_member_budget_table( + team_table: LiteLLM_TeamTable, + user_api_key_dict: UserAPIKeyAuth, + team_member_budget: float, + updated_kv: dict, +) -> dict: + """ + Add budget if none exists + + If budget exists, update it + """ + from litellm.proxy._types import BudgetNewRequest + from litellm.proxy.management_endpoints.budget_management_endpoints import ( + update_budget, + ) + + if team_table.metadata is None: + team_table.metadata = {} + + team_member_budget_id = team_table.metadata.get("team_member_budget_id") + if team_member_budget_id is not None and isinstance(team_member_budget_id, str): + # Budget exists + budget_row = await update_budget( + budget_obj=BudgetNewRequest( + budget_id=team_member_budget_id, + max_budget=team_member_budget, + ), + user_api_key_dict=user_api_key_dict, + ) + verbose_proxy_logger.info( + f"Updated team member budget table: {budget_row.budget_id}, with team_member_budget={team_member_budget}" + ) + if updated_kv.get("metadata") is None: + updated_kv["metadata"] = {} + updated_kv["metadata"]["team_member_budget_id"] = budget_row.budget_id + updated_kv.pop("team_member_budget", None) + else: # budget does not exist + updated_kv = await _create_team_member_budget_table( + data=team_table, + new_team_data_json=updated_kv, + user_api_key_dict=user_api_key_dict, + team_member_budget=team_member_budget, + ) + return updated_kv + + #### TEAM MANAGEMENT #### @router.post( "/team/new", @@ -311,9 +399,19 @@ async def new_team( # noqa: PLR0915 prisma_client=prisma_client, ) + ## Create Team Member Budget Table + data_json = data.json() + if data.team_member_budget is not None: + data_json = await _create_team_member_budget_table( + data=data, + new_team_data_json=data_json, + user_api_key_dict=user_api_key_dict, + team_member_budget=data.team_member_budget, + ) + ## ADD TO TEAM TABLE complete_team_data = LiteLLM_TeamTable( - **data.json(), + **data_json, model_id=_model_id, object_permission_id=object_permission_id, ) @@ -668,6 +766,14 @@ async def update_team( # set the budget_reset_at in DB updated_kv["budget_reset_at"] = reset_at + if data.team_member_budget is not None: + updated_kv = await _upsert_team_member_budget_table( + team_table=existing_team_row, + updated_kv=updated_kv, + team_member_budget=data.team_member_budget, + user_api_key_dict=user_api_key_dict, + ) + # Check object permission if data.object_permission is not None: updated_kv = await handle_update_object_permission( @@ -849,8 +955,6 @@ async def team_member_add( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - [BETA] - Add new members (either via user_email or user_id) to a team If user doesn't exist, new user row will also be added to User Table @@ -944,6 +1048,11 @@ async def team_member_add( user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, team_id=data.team_id, + default_team_budget_id=( + complete_team_data.metadata.get("team_member_budget_id") + if complete_team_data.metadata is not None + else None + ), ) except Exception as e: raise HTTPException( @@ -969,6 +1078,11 @@ async def team_member_add( user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, team_id=data.team_id, + default_team_budget_id=( + complete_team_data.metadata.get("team_member_budget_id") + if complete_team_data.metadata is not None + else None + ), ) except Exception as e: raise HTTPException( @@ -1471,6 +1585,24 @@ def _unfurl_all_proxy_models( return team_info +async def _add_team_member_budget_table( + team_member_budget_id: str, + prisma_client: PrismaClient, + team_info_response_object: TeamInfoResponseObjectTeamTable, +) -> TeamInfoResponseObjectTeamTable: + try: + team_budget = await prisma_client.db.litellm_budgettable.find_unique( + where={"budget_id": team_member_budget_id} + ) + team_info_response_object.team_member_budget_table = team_budget + except Exception: + verbose_proxy_logger.info( + f"Team member budget table not found, passed team_member_budget_id={team_member_budget_id}" + ) + + return team_info_response_object + + @router.get( "/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) @@ -1493,6 +1625,7 @@ async def team_info( --header 'Authorization: Bearer your_api_key_here' ``` """ + from litellm.proxy._types import TeamInfoResponseObjectTeamTable from litellm.proxy.proxy_server import prisma_client try: @@ -1561,11 +1694,24 @@ async def team_info( ) if isinstance(team_info, dict): - _team_info = LiteLLM_TeamTable(**team_info) + _team_info = TeamInfoResponseObjectTeamTable(**team_info) elif isinstance(team_info, BaseModel): - _team_info = LiteLLM_TeamTable(**team_info.model_dump()) + _team_info = TeamInfoResponseObjectTeamTable(**team_info.model_dump()) else: - _team_info = LiteLLM_TeamTable() + _team_info = TeamInfoResponseObjectTeamTable() + + ## GET TEAM BUDGET (if exists) ## + team_member_budget_id = ( + _team_info.metadata.get("team_member_budget_id") + if _team_info.metadata is not None + else None + ) + if team_member_budget_id is not None: + _team_info = await _add_team_member_budget_table( + team_member_budget_id=team_member_budget_id, + prisma_client=prisma_client, + team_info_response_object=_team_info, + ) # ## UNFURL 'all-proxy-models' into the team_info.models list ## # if llm_router is not None: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 03f740214a..29a078ec06 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -726,7 +726,7 @@ async def auth_callback(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "?login=success" verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}") redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token, secure=True) + redirect_response.set_cookie(key="token", value=jwt_token) return redirect_response @@ -763,9 +763,9 @@ async def insert_sso_user( if user_defined_values.get("max_budget") is None: user_defined_values["max_budget"] = litellm.max_internal_user_budget if user_defined_values.get("budget_duration") is None: - user_defined_values[ - "budget_duration" - ] = litellm.internal_user_budget_duration + user_defined_values["budget_duration"] = ( + litellm.internal_user_budget_duration + ) if user_defined_values["user_role"] is None: user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -783,7 +783,10 @@ async def insert_sso_user( if result_openid: new_user_request.metadata = {"auth_provider": result_openid.provider} - response = await new_user(data=new_user_request, user_api_key_dict=UserAPIKeyAuth()) + response = await new_user( + data=new_user_request, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) return response @@ -960,9 +963,9 @@ class SSOAuthenticationHandler: if state: redirect_params["state"] = state elif "okta" in generic_authorization_endpoint: - redirect_params[ - "state" - ] = uuid.uuid4().hex # set state param for okta - required + redirect_params["state"] = ( + uuid.uuid4().hex + ) # set state param for okta - required return await generic_sso.get_login_redirect(**redirect_params) # type: ignore raise ValueError( "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso" @@ -1209,9 +1212,9 @@ class MicrosoftSSOHandler: # if user is trying to get the raw sso response for debugging, return the raw sso response if return_raw_sso_response: - original_msft_result[ - MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY - ] = user_team_ids + original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = ( + user_team_ids + ) return original_msft_result or {} result = MicrosoftSSOHandler.openid_from_response( @@ -1279,9 +1282,9 @@ class MicrosoftSSOHandler: # Fetch user membership from Microsoft Graph API all_group_ids = [] - next_link: Optional[ - str - ] = MicrosoftSSOHandler.graph_api_user_groups_endpoint + next_link: Optional[str] = ( + MicrosoftSSOHandler.graph_api_user_groups_endpoint + ) auth_headers = {"Authorization": f"Bearer {access_token}"} page_count = 0 diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index ec645c0b15..246fa0b639 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -60,6 +60,7 @@ async def add_new_member( team_id: str, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, + default_team_budget_id: Optional[str] = None, ) -> Tuple[LiteLLM_UserTable, Optional[LiteLLM_TeamMembership]]: """ Add a new member to a team @@ -119,11 +120,8 @@ async def add_new_member( ) # Check if trying to set a budget for team member - if ( - max_budget_in_team is not None - and returned_user is not None - and returned_user.user_id is not None - ): + + if max_budget_in_team is not None: # create a new budget item for this member response = await prisma_client.db.litellm_budgettable.create( data={ @@ -134,6 +132,10 @@ async def add_new_member( ) _budget_id = response.budget_id + else: + _budget_id = default_team_budget_id + + if _budget_id and returned_user is not None and returned_user.user_id is not None: _returned_team_membership = ( await prisma_client.db.litellm_teammembership.create( data={ diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 878a6e1c17..a515f6bbaa 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1037,6 +1037,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): return_value=LiteLLM_TeamTableCachedObj(team_id="1234") ) + print(f"team_member_add_request={team_member_add_request}") await team_member_add( data=team_member_add_request, user_api_key_dict=UserAPIKeyAuth(user_role="proxy_admin"), @@ -2207,7 +2208,9 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch): monkeypatch.setattr(proxy_server, "health_check_interval", 1) monkeypatch.setattr(proxy_server, "health_check_details", None) - monkeypatch.setattr(proxy_server, "llm_model_list", copy.deepcopy(test_model_list_1)) + monkeypatch.setattr( + proxy_server, "llm_model_list", copy.deepcopy(test_model_list_1) + ) monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check) monkeypatch.setattr(proxy_server, "health_check_results", {}) @@ -2221,7 +2224,9 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch): except asyncio.CancelledError: pass - monkeypatch.setattr(proxy_server, "llm_model_list", copy.deepcopy(test_model_list_2)) + monkeypatch.setattr( + proxy_server, "llm_model_list", copy.deepcopy(test_model_list_2) + ) try: await proxy_server._run_background_health_check() diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index be4504fba9..98a50a55b0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -581,3 +581,135 @@ def test_team_member_add_duplication_check_allows_new_member(): except ProxyException: # If a ProxyException is raised, the test should fail pytest.fail("ProxyException should not be raised for a new member") + + +@pytest.mark.asyncio +async def test_add_team_member_budget_table_success(): + """ + Test _add_team_member_budget_table when budget is found successfully + """ + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_member_budget_table, + ) + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock budget record + mock_budget_record = MagicMock() + mock_budget_record.budget_id = "budget-123" + mock_budget_record.max_budget = 1000.0 + + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_budget_record + ) + + # Create team info response object + team_info_response = TeamInfoResponseObjectTeamTable( + team_id="test-team-123", team_alias="Test Team" + ) + + # Call the function + result = await _add_team_member_budget_table( + team_member_budget_id="budget-123", + prisma_client=mock_prisma_client, + team_info_response_object=team_info_response, + ) + + # Verify the result + assert result == team_info_response + assert result.team_member_budget_table == mock_budget_record + + # Verify database call was made correctly + mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( + where={"budget_id": "budget-123"} + ) + + +@pytest.mark.asyncio +async def test_add_team_member_budget_table_exception_handling(): + """ + Test _add_team_member_budget_table when an exception occurs during budget lookup + """ + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_member_budget_table, + ) + + # Mock prisma client to raise an exception + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + side_effect=Exception("Database connection failed") + ) + + # Create team info response object + team_info_response = TeamInfoResponseObjectTeamTable( + team_id="test-team-456", team_alias="Test Team 2" + ) + + # Mock the verbose_proxy_logger to capture log calls + with patch( + "litellm.proxy.management_endpoints.team_endpoints.verbose_proxy_logger" + ) as mock_logger: + # Call the function + result = await _add_team_member_budget_table( + team_member_budget_id="nonexistent-budget-456", + prisma_client=mock_prisma_client, + team_info_response_object=team_info_response, + ) + + # Verify the result is returned even when exception occurs + assert result == team_info_response + + # Verify team_member_budget_table is not set when exception occurs + assert ( + not hasattr(result, "team_member_budget_table") + or result.team_member_budget_table is None + ) + + # Verify the error was logged + mock_logger.info.assert_called_once_with( + "Team member budget table not found, passed team_member_budget_id=nonexistent-budget-456" + ) + + # Verify database call was attempted + mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( + where={"budget_id": "nonexistent-budget-456"} + ) + + +@pytest.mark.asyncio +async def test_add_team_member_budget_table_budget_not_found(): + """ + Test _add_team_member_budget_table when budget record is not found (returns None) + """ + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_member_budget_table, + ) + + # Mock prisma client to return None (budget not found) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + + # Create team info response object + team_info_response = TeamInfoResponseObjectTeamTable( + team_id="test-team-789", team_alias="Test Team 3" + ) + + # Call the function + result = await _add_team_member_budget_table( + team_member_budget_id="nonexistent-budget-789", + prisma_client=mock_prisma_client, + team_info_response_object=team_info_response, + ) + + # Verify the result + assert result == team_info_response + assert result.team_member_budget_table is None + + # Verify database call was made correctly + mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( + where={"budget_id": "nonexistent-budget-789"} + ) diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py new file mode 100644 index 0000000000..2f6ecb776d --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -0,0 +1,293 @@ +import json +import os +import sys +import uuid +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_UserTable, + Member, + UserAPIKeyAuth, +) +from litellm.proxy.management_helpers.utils import add_new_member + + +@pytest.mark.asyncio +async def test_add_new_member_uses_default_team_budget_id(): + """ + Test that add_new_member uses the default_team_budget_id when max_budget_in_team is None. + + This test verifies that: + 1. When max_budget_in_team is None + 2. And default_team_budget_id is provided + 3. The team membership is created with the default_team_budget_id + """ + from litellm.proxy._types import LitellmUserRoles + + # Setup test data + test_user_id = "test_user_123" + test_team_id = "test_team_456" + test_default_budget_id = "default_budget_789" + test_admin_name = "test_admin" + + # Create a Member object with user_id + new_member = Member(user_id=test_user_id, role="user") + + # Create UserAPIKeyAuth object + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock the prisma client + mock_prisma_client = AsyncMock() + + # Mock the user table upsert operation + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": test_user_id, + "user_email": None, + "teams": [test_team_id], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + + # Mock the team membership creation + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": test_team_id, + "user_id": test_user_id, + "budget_id": test_default_budget_id, + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + # Call the function with max_budget_in_team=None and a default_team_budget_id + result_user, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, # This is the key - no max budget specified + prisma_client=mock_prisma_client, + team_id=test_team_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=test_admin_name, + default_team_budget_id=test_default_budget_id, # This should be used + ) + + # Verify that the user was created/updated correctly + assert result_user is not None + assert result_user.user_id == test_user_id + + # Verify that the team membership was created correctly + assert result_team_membership is not None + assert result_team_membership.team_id == test_team_id + assert result_team_membership.user_id == test_user_id + assert result_team_membership.budget_id == test_default_budget_id + + # Verify that the prisma client methods were called correctly + mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() + mock_prisma_client.db.litellm_teammembership.create.assert_called_once() + + # Verify that no budget table creation was called (since max_budget_in_team is None) + assert ( + not hasattr(mock_prisma_client.db, "litellm_budgettable") + or not mock_prisma_client.db.litellm_budgettable.create.called + ) + + # Verify the team membership was created with the correct budget_id + team_membership_call_args = ( + mock_prisma_client.db.litellm_teammembership.create.call_args + ) + assert team_membership_call_args is not None + create_data = team_membership_call_args.kwargs["data"] + assert create_data["budget_id"] == test_default_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_creates_new_budget_when_max_budget_provided(): + """ + Test that add_new_member creates a new budget when max_budget_in_team is provided. + + This test verifies that: + 1. When max_budget_in_team is provided + 2. A new budget is created in the litellm_budgettable + 3. The new budget_id is used for the team membership + """ + from litellm.proxy._types import LitellmUserRoles + + # Setup test data + test_user_id = "test_user_123" + test_team_id = "test_team_456" + test_max_budget = 100.0 + test_new_budget_id = "new_budget_789" + test_admin_name = "test_admin" + + # Create a Member object with user_id + new_member = Member(user_id=test_user_id, role="user") + + # Create UserAPIKeyAuth object + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock the prisma client + mock_prisma_client = AsyncMock() + + # Mock the user table upsert operation + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": test_user_id, + "user_email": None, + "teams": [test_team_id], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + + # Mock the budget table creation + mock_budget_response = MagicMock() + mock_budget_response.budget_id = test_new_budget_id + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_budget_response + ) + + # Mock the team membership creation + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": test_team_id, + "user_id": test_user_id, + "budget_id": test_new_budget_id, + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + # Call the function with max_budget_in_team provided + result_user, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=test_max_budget, # This should trigger budget creation + prisma_client=mock_prisma_client, + team_id=test_team_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=test_admin_name, + default_team_budget_id=None, # Should be ignored since max_budget_in_team is provided + ) + + # Verify that the budget was created + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + budget_call_args = mock_prisma_client.db.litellm_budgettable.create.call_args + budget_data = budget_call_args.kwargs["data"] + assert budget_data["max_budget"] == test_max_budget + assert budget_data["created_by"] == user_api_key_dict.user_id + assert budget_data["updated_by"] == user_api_key_dict.user_id + + # Verify that the team membership was created with the new budget_id + assert result_team_membership is not None + assert result_team_membership.budget_id == test_new_budget_id + + # Verify the team membership was created with the correct budget_id + team_membership_call_args = ( + mock_prisma_client.db.litellm_teammembership.create.call_args + ) + assert team_membership_call_args is not None + create_data = team_membership_call_args.kwargs["data"] + assert create_data["budget_id"] == test_new_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_with_user_email(): + """ + Test add_new_member with user_email instead of user_id and default budget. + + This test verifies that: + 1. When new_member has user_email instead of user_id + 2. And max_budget_in_team is None + 3. The default_team_budget_id is used correctly + """ + from litellm.proxy._types import LitellmUserRoles + + # Setup test data + test_user_email = "test@example.com" + test_team_id = "test_team_456" + test_default_budget_id = "default_budget_789" + test_admin_name = "test_admin" + + # Create a Member object with user_email + new_member = Member(user_email=test_user_email, role="user") + + # Create UserAPIKeyAuth object + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock the prisma client + mock_prisma_client = AsyncMock() + + # Mock get_data to return empty list (no existing user) + mock_prisma_client.get_data = AsyncMock(return_value=[]) + + # Mock insert_data for new user creation + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "generated_user_id", + "user_email": test_user_email, + "teams": [test_team_id], + "user_role": "internal_user", + } + mock_prisma_client.insert_data = AsyncMock(return_value=mock_user_response) + + # Mock the team membership creation + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": test_team_id, + "user_id": "generated_user_id", + "budget_id": test_default_budget_id, + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + # Call the function + result_user, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id=test_team_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=test_admin_name, + default_team_budget_id=test_default_budget_id, + ) + + # Verify that the user was created correctly + assert result_user is not None + assert result_user.user_email == test_user_email + + # Verify that the team membership was created with the default budget_id + assert result_team_membership is not None + assert result_team_membership.budget_id == test_default_budget_id + + # Verify that get_data was called to check for existing user + mock_prisma_client.get_data.assert_called_once_with( + key_val={"user_email": test_user_email}, + table_name="user", + query_type="find_all", + ) + + # Verify that insert_data was called to create new user + mock_prisma_client.insert_data.assert_called_once() + insert_call_args = mock_prisma_client.insert_data.call_args + insert_data = insert_call_args.kwargs["data"] + assert insert_data["user_email"] == test_user_email + assert insert_data["teams"] == [test_team_id] diff --git a/ui/litellm-dashboard/src/components/team/team_info.tsx b/ui/litellm-dashboard/src/components/team/team_info.tsx index 1c6bfd4683..0d79f56262 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.tsx @@ -38,6 +38,23 @@ import VectorStoreSelector from "../vector_store_management/VectorStoreSelector" import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import PremiumVectorStoreSelector from "../common_components/PremiumVectorStoreSelector"; +export interface TeamMembership { + user_id: string; + team_id: string; + budget_id: string; + spend: number; + litellm_budget_table: { + budget_id: string; + soft_budget: number | null; + max_budget: number | null; + max_parallel_requests: number | null; + tpm_limit: number | null; + rpm_limit: number | null; + model_max_budget: Record | null; + budget_duration: string | null; + }; +} + export interface TeamData { team_id: string; team_info: { @@ -67,9 +84,13 @@ export interface TeamData { mcp_servers: string[]; vector_stores: string[]; }; + team_member_budget_table: { + max_budget: number; + budget_duration: string; + } | null; }; keys: any[]; - team_memberships: any[]; + team_memberships: TeamMembership[]; } export interface TeamInfoProps { @@ -235,6 +256,10 @@ const TeamInfoView: React.FC = ({ organization_id: values.organization_id, }; + if (values.team_member_budget !== undefined) { + updateData.team_member_budget = Number(values.team_member_budget); + } + // Handle object_permission updates if (values.vector_stores !== undefined || values.mcp_servers !== undefined) { updateData.object_permission = { @@ -298,6 +323,10 @@ const TeamInfoView: React.FC = ({ {info.budget_duration && ( Reset: {info.budget_duration} )} +
+ {info.team_member_budget_table && ( + Team Member Budget: ${info.team_member_budget_table.max_budget} + )} @@ -420,6 +449,10 @@ const TeamInfoView: React.FC = ({ + + + +