optional query params for org list

This commit is contained in:
yuneng-jiang
2026-01-10 14:44:28 -08:00
parent 7e6d0d7691
commit 2e6fe80f90
2 changed files with 200 additions and 14 deletions
@@ -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
@@ -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,
}