From 2e6fe80f909cdfd55e22abbeef9938e31c073682 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 10 Jan 2026 14:44:28 -0800 Subject: [PATCH] optional query params for org list --- .../organization_endpoints.py | 89 +++++++++++-- .../test_organization_endpoints.py | 125 ++++++++++++++++++ 2 files changed, 200 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index a088f46d67..94cfa82c0c 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -11,8 +11,9 @@ Endpoints for /organization operations #### ORGANIZATION MANAGEMENT #### -from typing import List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple +import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, status from litellm._logging import verbose_proxy_logger @@ -595,11 +596,33 @@ async def delete_organization( response_model=List[LiteLLM_OrganizationTableWithMembers], ) async def list_organization( + org_id: Optional[str] = fastapi.Query( + default=None, description="Filter organizations by exact organization_id match" + ), + org_alias: Optional[str] = fastapi.Query( + default=None, + description="Filter organizations by partial organization_alias match. Supports case-insensitive search.", + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ + Get a list of organizations with optional filtering. + + Parameters: + org_id: Optional[str] + Filter organizations by exact organization_id match + org_alias: Optional[str] + Filter organizations by partial organization_alias match (case-insensitive) + + Example: ``` - curl --location --request GET 'http://0.0.0.0:4000/organization/list' \ + curl --location --request GET 'http://0.0.0.0:4000/organization/list?org_alias=my-org' \ + --header 'Authorization: Bearer sk-1234' + ``` + + Example with org_id: + ``` + curl --location --request GET 'http://0.0.0.0:4000/organization/list?org_id=123e4567-e89b-12d3-a456-426614174000' \ --header 'Authorization: Bearer sk-1234' ``` """ @@ -614,28 +637,66 @@ async def list_organization( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - # if proxy admin - get all orgs + # Build where conditions based on provided filters + where_conditions: Dict[str, Any] = {} + + if org_id: + where_conditions["organization_id"] = org_id + + if org_alias: + where_conditions["organization_alias"] = { + "contains": org_alias, + "mode": "insensitive", # Case-insensitive search + } + + # if proxy admin - get all orgs (with optional filters) if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: response = await prisma_client.db.litellm_organizationtable.find_many( - include={"litellm_budget_table": True, "members": True, "teams": True} + where=where_conditions if where_conditions else None, + include={"litellm_budget_table": True, "members": True, "teams": True}, ) - # if internal user - get orgs they are a member of + # if internal user - get orgs they are a member of (with optional filters) else: org_memberships = ( await prisma_client.db.litellm_organizationmembership.find_many( where={"user_id": user_api_key_dict.user_id} ) ) - org_objects = await prisma_client.db.litellm_organizationtable.find_many( - where={ - "organization_id": { - "in": [membership.organization_id for membership in org_memberships] - } - }, - include={"litellm_budget_table": True, "members": True, "teams": True}, - ) + membership_org_ids = [ + membership.organization_id for membership in org_memberships + ] - response = org_objects + # Combine membership filter with provided filters + if membership_org_ids: + if org_id: + # If org_id is provided, ensure user is a member of that org + if org_id not in membership_org_ids: + # User is not a member of the requested org, return empty list + response = [] + else: + where_conditions["organization_id"] = org_id + response = await prisma_client.db.litellm_organizationtable.find_many( + where=where_conditions, + include={ + "litellm_budget_table": True, + "members": True, + "teams": True, + }, + ) + else: + # Filter by membership and any additional filters + where_conditions["organization_id"] = {"in": membership_org_ids} + response = await prisma_client.db.litellm_organizationtable.find_many( + where=where_conditions, + include={ + "litellm_budget_table": True, + "members": True, + "teams": True, + }, + ) + else: + # User is not a member of any orgs + response = [] return response diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index c02db727bf..2a2c37d03c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -410,3 +410,128 @@ async def test_organization_update_object_permissions_missing_permission_record( # Verify upsert was called to create new record mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once() + + +@pytest.mark.asyncio +async def test_list_organization_filter_by_org_id(monkeypatch): + """ + Test filtering organizations by org_id query parameter. + + This test verifies that when org_id is provided, only the organization + with that exact organization_id is returned. + """ + from types import SimpleNamespace + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import ( + list_organization, + ) + + # Mock prisma client + mock_prisma_client = AsyncMock() + + # Mock organization data + mock_org1 = SimpleNamespace( + organization_id="org-123", + organization_alias="Test Org 1", + model_dump=lambda: { + "organization_id": "org-123", + "organization_alias": "Test Org 1", + }, + ) + + # Mock find_many to return filtered results + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( + return_value=[mock_org1] + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Test as proxy admin + auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ) + + result = await list_organization(org_id="org-123", org_alias=None, user_api_key_dict=auth) + + # Verify the correct organization was returned + assert len(result) == 1 + assert result[0].organization_id == "org-123" + assert result[0].organization_alias == "Test Org 1" + + # Verify find_many was called with correct where conditions + mock_prisma_client.db.litellm_organizationtable.find_many.assert_called_once() + call_args = mock_prisma_client.db.litellm_organizationtable.find_many.call_args + assert call_args.kwargs["where"] == {"organization_id": "org-123"} + assert call_args.kwargs["include"] == { + "litellm_budget_table": True, + "members": True, + "teams": True, + } + + +@pytest.mark.asyncio +async def test_list_organization_filter_by_org_alias(monkeypatch): + """ + Test filtering organizations by org_alias query parameter with case-insensitive partial matching. + + This test verifies that when org_alias is provided, organizations with matching + organization_alias (case-insensitive partial match) are returned. + """ + from types import SimpleNamespace + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import ( + list_organization, + ) + + # Mock prisma client + mock_prisma_client = AsyncMock() + + # Mock organization data + mock_org1 = SimpleNamespace( + organization_id="org-123", + organization_alias="My Test Organization", + model_dump=lambda: { + "organization_id": "org-123", + "organization_alias": "My Test Organization", + }, + ) + mock_org2 = SimpleNamespace( + organization_id="org-456", + organization_alias="Another Test Org", + model_dump=lambda: { + "organization_id": "org-456", + "organization_alias": "Another Test Org", + }, + ) + + # Mock find_many to return filtered results + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( + return_value=[mock_org1, mock_org2] + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Test as proxy admin with org_alias filter + auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ) + + result = await list_organization(org_id=None, org_alias="test", user_api_key_dict=auth) + + # Verify organizations with "test" in alias were returned + assert len(result) == 2 + assert all("test" in org.organization_alias.lower() for org in result) + + # Verify find_many was called with correct where conditions (case-insensitive contains) + mock_prisma_client.db.litellm_organizationtable.find_many.assert_called_once() + call_args = mock_prisma_client.db.litellm_organizationtable.find_many.call_args + assert call_args.kwargs["where"] == { + "organization_alias": {"contains": "test", "mode": "insensitive"} + } + assert call_args.kwargs["include"] == { + "litellm_budget_table": True, + "members": True, + "teams": True, + }