diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 9009ce8995b..324416cb05d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1052,7 +1052,7 @@ async def fetch_and_validate_organization( organization_row = await prisma_client.db.litellm_organizationtable.find_unique( where={"organization_id": organization_id}, - include={"litellm_budget_table": True, "members": True}, + include={"litellm_budget_table": True, "members": True, "teams": True}, ) if organization_row is None: 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 d096b5515a0..c20d4aa2027 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -3726,4 +3726,139 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): ) # Verify team was updated - assert result["team_id"] == "org-team-update-bypass-123" \ No newline at end of file + assert result["team_id"] == "org-team-update-bypass-123" + + +@pytest.mark.asyncio +async def test_update_team_guardrails_with_org_id(): + """ + Test that updating team guardrails works when team has an organization_id. + The fix ensures 'teams' field is included when fetching organization data. + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_TeamTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create user (org admin) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-guardrails-test", + models=[], + ) + + # Update request to add guardrails to team + update_request = UpdateTeamRequest( + team_id="team-guardrails-123", + guardrails=["aporia-pre-call", "aporia-post-call"], + organization_id="test-org-guardrails", # Changing org triggers fetch_and_validate_organization + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with all required fields including teams (the fix) + from datetime import datetime + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-guardrails" + mock_org.models = ["gpt-4", "gpt-3.5-turbo"] + mock_org.budget_id = "budget-123" + mock_org.created_by = "admin" + mock_org.updated_by = "admin" + mock_org.created_at = datetime(2024, 1, 1) + mock_org.updated_at = datetime(2024, 1, 1) + mock_org.litellm_budget_table = None + mock_org.members = [] + mock_org.teams = [] # Must be a list, not None + mock_org.model_dump.return_value = { + "organization_id": "test-org-guardrails", + "models": ["gpt-4", "gpt-3.5-turbo"], + "budget_id": "budget-123", + "created_by": "admin", + "updated_by": "admin", + "created_at": datetime(2024, 1, 1), + "updated_at": datetime(2024, 1, 1), + "litellm_budget_table": None, + "members": [], + "teams": [], + } + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ): + # Mock existing team - must have compatible models with organization + mock_existing_team = MagicMock() + mock_existing_team.team_id = "team-guardrails-123" + mock_existing_team.organization_id = None + mock_existing_team.metadata = {} + mock_existing_team.model_id = None + mock_existing_team.models = ["gpt-4"] # Subset of org models to pass validation + mock_existing_team.max_budget = None + mock_existing_team.tpm_limit = None + mock_existing_team.rpm_limit = None + mock_existing_team.model_dump.return_value = { + "team_id": "team-guardrails-123", + "organization_id": None, + "metadata": {}, + "models": ["gpt-4"], + "max_budget": None, + "tpm_limit": None, + "rpm_limit": None, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team + ) + mock_cache.async_set_cache = AsyncMock() + + # Mock organization fetch - this is where the bug occurred + # The fix ensures 'teams: True' is in the include clause + mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock( + return_value=mock_org + ) + + # Mock team update + mock_updated_team = MagicMock(spec=LiteLLM_TeamTable) + mock_updated_team.team_id = "team-guardrails-123" + mock_updated_team.organization_id = "test-org-guardrails" + mock_updated_team.metadata = {"guardrails": ["aporia-pre-call", "aporia-post-call"]} + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "team-guardrails-123", + "organization_id": "test-org-guardrails", + "metadata": {"guardrails": ["aporia-pre-call", "aporia-post-call"]}, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) + mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + # Mock llm_router + mock_router = MagicMock() + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + # This should succeed without Pydantic validation error + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify the team was updated successfully with guardrails + assert result is not None + assert result["data"].organization_id == "test-org-guardrails" + assert result["data"].metadata["guardrails"] == ["aporia-pre-call", "aporia-post-call"] + + # Verify that organization fetch was called with proper include clause + # The function is called twice: once by fetch_and_validate_organization (with include) + # and once by get_org_object (without include). We verify the first call has 'teams'. + assert mock_prisma.db.litellm_organizationtable.find_unique.call_count >= 1 + + # Get the first call (from fetch_and_validate_organization) + first_call_kwargs = mock_prisma.db.litellm_organizationtable.find_unique.call_args_list[0].kwargs + + # Verify that 'teams' is included in the fetch + assert "include" in first_call_kwargs + assert "teams" in first_call_kwargs["include"] + assert first_call_kwargs["include"]["teams"] is True