scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

[Feat] UI - show vector store permissions for Key, Team, Org (#11277)

Browse source 
* fix LiteLLM_ObjectPermissionTable

* fix include object_permission for list key

* fix key list to inclue obj permissions

* fix object permissions for vector stores on key info

* add key edit view with vector stores

* allow editing vector stores permissions

* fixes obj permissions

* feat: add obj permission on UI

* fix: add object_permission:true

* ui show org vector stores on org info

* fix: show object permissions on /org/info

* feat: allow updating obj permissions for keys

* fixes: key object permissions

* fixes: team object permissions

* fixes: org object permissions

* fix vector store selector for Orgs
This commit is contained in:
Ishaan Jaff 2025-05-30 17:23:50 -07:00 • committed by GitHub
parent b13d1b3ff1
commit ea841eeb9b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 1160 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_types.py
View file

@ -1157,6 +1157,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
tags: Optional[list] = None
model_aliases: Optional[dict] = None
guardrails: Optional[List[str]] = None
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase):
@ -1254,6 +1255,7 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase):
object_permission_id: str
mcp_servers: List[str]
vector_stores: List[str]
class LiteLLM_TeamTable(TeamBase):
@ -1636,6 +1638,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
created_by: Optional[str] = None
updated_at: Optional[datetime] = None
updated_by: Optional[str] = None
object_permission_id: Optional[str] = None
object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
model_config = ConfigDict(protected_namespaces=())
@ -1757,6 +1760,7 @@ class LiteLLM_OrganizationTableUpdate(LiteLLMPydanticObjectBase):
metadata: Optional[dict] = None
models: Optional[List[str]] = None
updated_by: Optional[str] = None
object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
class LiteLLM_UserTable(LiteLLMPydanticObjectBase):

70
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -673,8 +673,9 @@ async def _set_object_permission(
return data_json
def prepare_key_update_data(
data: Union[UpdateKeyRequest, RegenerateKeyRequest], existing_key_row
async def prepare_key_update_data(
data: Union[UpdateKeyRequest, RegenerateKeyRequest],
existing_key_row: LiteLLM_VerificationToken,
):
data_json: dict = data.model_dump(exclude_unset=True)
data_json.pop("key", None)
@ -704,6 +705,12 @@ def prepare_key_update_data(
non_default_values["budget_reset_at"] = key_reset_at
non_default_values["budget_duration"] = budget_duration
if "object_permission" in non_default_values:
non_default_values = await _handle_update_object_permission(
data_json=non_default_values,
existing_key_row=existing_key_row,
)
_metadata = existing_key_row.metadata or {}
# validate model_max_budget
@ -717,6 +724,60 @@ def prepare_key_update_data(
return non_default_values
async def _handle_update_object_permission(
data_json: dict,
existing_key_row: LiteLLM_VerificationToken,
) -> dict:
"""
Handle the update of object permission.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise ValueError("Prisma client not found")
#########################################################
# Ensure `object_permission` is not added to the data_json
# We need to update the entity at the object_permission_id level in the LiteLLM_ObjectPermissionTable
#########################################################
new_object_permission = data_json.pop("object_permission")
if new_object_permission is None:
return data_json
# lookup existing object permission ID and update that entry
existing_object_permission_id = existing_key_row.object_permission_id
existing_object_permission = (
await prisma_client.db.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": existing_object_permission_id},
)
)
existing_object_permissions_dict: dict = {}
if existing_object_permission is not None:
# update the object permission
existing_object_permissions_dict = existing_object_permission.model_dump(
exclude_unset=True, exclude_none=True
)
existing_object_permissions_dict.update(dict(new_object_permission))
#########################################################
# Commit the update to the LiteLLM_ObjectPermissionTable
#########################################################
new_object_permission_row = (
await prisma_client.db.litellm_objectpermissiontable.upsert(
where={"object_permission_id": existing_object_permission_id},
data={
"create": existing_object_permissions_dict,
"update": existing_object_permissions_dict,
},
)
)
verbose_proxy_logger.debug(
f"new_object_permission_row: {new_object_permission_row}"
)
data_json["object_permission_id"] = new_object_permission_row.object_permission_id
return data_json
def is_different_team(
data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken
) -> bool:
@ -853,7 +914,7 @@ async def update_key_fn(
change_initiated_by=user_api_key_dict,
llm_router=llm_router,
)
non_default_values = prepare_key_update_data(
non_default_values = await prepare_key_update_data(
data=data, existing_key_row=existing_key_row
)
@ -1945,7 +2006,7 @@ async def regenerate_key_fn(
non_default_values = {}
if data is not None:
# Update with any provided parameters from GenerateKeyRequest
non_default_values = prepare_key_update_data(
non_default_values = await prepare_key_update_data(
data=data, existing_key_row=_key_in_db
)
verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
@ -2372,6 +2433,7 @@ async def _list_key_helper(
{"created_at": "desc"},
{"token": "desc"}, # fallback sort
],
include={"object_permission": True},
)
verbose_proxy_logger.debug(f"Fetched {len(keys)} keys")

91
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -273,6 +273,22 @@ async def update_organization(
updated_organization_row = prisma_client.jsonify_object(
data.model_dump(exclude_none=True)
)
existing_organization_row = (
await prisma_client.db.litellm_organizationtable.find_unique(
where={"organization_id": data.organization_id},
)
)
if existing_organization_row is None:
raise ValueError(
f"Organization not found for organization_id={data.organization_id}"
)
if data.object_permission is not None:
updated_organization_row = await handle_update_object_permission(
data_json=updated_organization_row,
existing_organization_row=existing_organization_row,
)
response = await prisma_client.db.litellm_organizationtable.update(
where={"organization_id": data.organization_id},
@ -283,6 +299,74 @@ async def update_organization(
return response
async def handle_update_object_permission(
data_json: dict,
existing_organization_row: LiteLLM_OrganizationTable,
) -> dict:
"""
Handle the update of object permission for an organization.
- Upserts the new object permission into the LiteLLM_ObjectPermissionTable
- Adds object_permission_id to data_json (this gets added in the DB)
- Pops the object_permission from data_json
-
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise ValueError("Prisma client not found")
#########################################################
# Ensure `object_permission` is not added to the data_json
# We need to update the entity at the object_permission_id level in the LiteLLM_ObjectPermissionTable
#########################################################
new_object_permission: Union[dict, str] = data_json.pop("object_permission") or {}
if new_object_permission is None:
return data_json
# lookup existing object permission ID and update that entry
existing_object_permission_id = existing_organization_row.object_permission_id
existing_object_permissions_dict = {}
existing_object_permission = (
await prisma_client.db.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": existing_object_permission_id},
)
)
# update the object permission
if existing_object_permission is not None:
existing_object_permissions_dict = existing_object_permission.model_dump(
exclude_unset=True, exclude_none=True
)
if isinstance(new_object_permission, str):
new_object_permission = json.loads(new_object_permission)
if isinstance(new_object_permission, dict):
existing_object_permissions_dict.update(new_object_permission)
#########################################################
# Commit the update to the LiteLLM_ObjectPermissionTable
#########################################################
created_object_permission_row = (
await prisma_client.db.litellm_objectpermissiontable.upsert(
where={"object_permission_id": existing_object_permission_id},
data={
"create": existing_object_permissions_dict,
"update": existing_object_permissions_dict,
},
)
)
data_json[
"object_permission_id"
] = created_object_permission_row.object_permission_id
verbose_proxy_logger.debug(
f"created_object_permission_row: {created_object_permission_row}"
)
return data_json
@router.delete(
"/organization/delete",
tags=["organization management"],
@ -414,7 +498,12 @@ async def info_organization(organization_id: str):
LiteLLM_OrganizationTableWithMembers
] = await prisma_client.db.litellm_organizationtable.find_unique(
where={"organization_id": organization_id},
include={"litellm_budget_table": True, "members": True, "teams": True},
include={
"litellm_budget_table": True,
"members": True,
"teams": True,
"object_permission": True,
},
)
if response is None:

66
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -663,6 +663,13 @@ async def update_team(
# set the budget_reset_at in DB
updated_kv["budget_reset_at"] = reset_at
# Check object permission
if data.object_permission is not None:
updated_kv = await handle_update_object_permission(
data_json=updated_kv,
existing_team_row=existing_team_row,
)
# update team metadata fields
_team_metadata_fields = LiteLLM_ManagementEndpoint_MetadataFields_Premium
for field in _team_metadata_fields:
@ -734,6 +741,62 @@ async def update_team(
return {"team_id": team_row.team_id, "data": team_row}
async def handle_update_object_permission(
data_json: dict, existing_team_row: LiteLLM_TeamTable
) -> dict:
"""
Handle the update of object permission for a team.
- IF there's no object_permission_id, then create a new entry in LiteLLM_ObjectPermissionTable
- IF there's an object_permission_id, then update the entry in LiteLLM_ObjectPermissionTable
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise ValueError("Prisma client not found")
#########################################################
# Ensure `object_permission` is not added to the data_json
# We need to update the entity at the object_permission_id level in the LiteLLM_ObjectPermissionTable
#########################################################
new_object_permission = data_json.pop("object_permission")
if new_object_permission is None:
return data_json
# lookup existing object permission ID and update that entry
existing_object_permission_id = existing_team_row.object_permission_id
existing_object_permission = (
await prisma_client.db.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": existing_object_permission_id},
)
)
existing_object_permissions_dict: Dict = {}
# update the object permission
if existing_object_permission is not None:
existing_object_permissions_dict = existing_object_permission.model_dump(
exclude_unset=True, exclude_none=True
)
existing_object_permissions_dict.update(dict(new_object_permission))
created_object_permission_row = (
await prisma_client.db.litellm_objectpermissiontable.upsert(
where={"object_permission_id": existing_object_permission_id},
data={
"create": existing_object_permissions_dict,
"update": existing_object_permissions_dict,
},
)
)
data_json[
"object_permission_id"
] = created_object_permission_row.object_permission_id
verbose_proxy_logger.debug(
f"created_object_permission_row: {created_object_permission_row}"
)
return data_json
def _check_team_member_admin_add(
member: Union[Member, List[Member]],
premium_user: bool,
@ -1486,7 +1549,8 @@ async def team_info(
team_info: Optional[
BaseModel
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
where={"team_id": team_id},
include={"object_permission": True},
)
if team_info is None:
raise Exception

9
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -639,7 +639,8 @@ async def test_proxy_config_update_from_db():
}
def test_prepare_key_update_data():
@pytest.mark.asyncio
async def test_prepare_key_update_data():
from litellm.proxy.management_endpoints.key_management_endpoints import (
prepare_key_update_data,
)
@ -647,15 +648,15 @@ def test_prepare_key_update_data():
existing_key_row = MagicMock()
data = UpdateKeyRequest(key="test_key", models=["gpt-4"], duration="120s")
updated_data = prepare_key_update_data(data, existing_key_row)
updated_data = await prepare_key_update_data(data, existing_key_row)
assert "expires" in updated_data
data = UpdateKeyRequest(key="test_key", metadata={})
updated_data = prepare_key_update_data(data, existing_key_row)
updated_data = await prepare_key_update_data(data, existing_key_row)
assert updated_data["metadata"] == {}
data = UpdateKeyRequest(key="test_key", metadata=None)
updated_data = prepare_key_update_data(data, existing_key_row)
updated_data = await prepare_key_update_data(data, existing_key_row)
assert updated_data["metadata"] is None

214
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -251,3 +251,217 @@ async def test_key_generation_with_object_permission(monkeypatch):
]
assert len(key_insert_calls) == 1
assert key_insert_calls[0]["data"].get("object_permission_id") == "objperm123"
@pytest.mark.asyncio
async def test_key_update_object_permissions_existing_permission(monkeypatch):
"""
Test updating object permissions when a key already has an existing object_permission_id.
This test verifies that when updating vector stores for a key that already has an
object_permission_id, the existing LiteLLM_ObjectPermissionTable record is updated
with the new permissions and the object_permission_id remains the same.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_VerificationToken,
)
from litellm.proxy.management_endpoints.key_management_endpoints import (
_handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Mock existing key with object_permission_id
existing_key_row = LiteLLM_VerificationToken(
token="test_token_hash",
object_permission_id="existing_perm_id_123",
user_id="user123",
team_id=None,
)
# Mock existing object permission record
existing_object_permission = MagicMock()
existing_object_permission.model_dump.return_value = {
"object_permission_id": "existing_perm_id_123",
"vector_stores": ["old_store_1", "old_store_2"],
}
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=existing_object_permission
)
# Mock upsert operation
updated_permission = MagicMock()
updated_permission.object_permission_id = "existing_perm_id_123"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=updated_permission
)
# Test data with new object permission
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["new_store_1", "new_store_2", "new_store_3"]
).model_dump(exclude_unset=True, exclude_none=True),
"user_id": "user123",
}
# Call the function
result = await _handle_update_object_permission(
data_json=data_json,
existing_key_row=existing_key_row,
)
# Verify the object_permission was removed from data_json and object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "existing_perm_id_123"
# Verify database operations were called correctly
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": "existing_perm_id_123"}
)
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
@pytest.mark.asyncio
async def test_key_update_object_permissions_no_existing_permission(monkeypatch):
"""
Test creating object permissions when a key has no existing object_permission_id.
This test verifies that when updating object permissions for a key that has
object_permission_id set to None, a new entry is created in the
LiteLLM_ObjectPermissionTable and the key is updated with the new object_permission_id.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_VerificationToken,
)
from litellm.proxy.management_endpoints.key_management_endpoints import (
_handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
existing_key_row_no_perm = LiteLLM_VerificationToken(
token="test_token_hash_2",
object_permission_id=None,
user_id="user456",
team_id=None,
)
# Mock find_unique to return None (no existing permission)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=None
)
# Mock upsert to create new record
new_permission = MagicMock()
new_permission.object_permission_id = "new_perm_id_456"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=new_permission
)
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["brand_new_store"]
).model_dump(exclude_unset=True, exclude_none=True),
"user_id": "user456",
}
result = await _handle_update_object_permission(
data_json=data_json,
existing_key_row=existing_key_row_no_perm,
)
# Verify new object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "new_perm_id_456"
# Verify find_unique was called with None
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": None}
)
# Verify upsert was called to create new record
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
@pytest.mark.asyncio
async def test_key_update_object_permissions_missing_permission_record(monkeypatch):
"""
Test creating object permissions when existing object_permission_id record is not found.
This test verifies that when updating object permissions for a key that has an
object_permission_id but the corresponding record cannot be found in the database,
a new entry is created in the LiteLLM_ObjectPermissionTable with the new permissions.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_VerificationToken,
)
from litellm.proxy.management_endpoints.key_management_endpoints import (
_handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
existing_key_row_missing_perm = LiteLLM_VerificationToken(
token="test_token_hash_3",
object_permission_id="missing_perm_id_789",
user_id="user789",
team_id=None,
)
# Mock find_unique to return None (permission record not found)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=None
)
# Mock upsert to create new record
new_permission = MagicMock()
new_permission.object_permission_id = "recreated_perm_id_789"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=new_permission
)
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["recreated_store"]
).model_dump(exclude_unset=True, exclude_none=True),
"user_id": "user789",
}
result = await _handle_update_object_permission(
data_json=data_json,
existing_key_row=existing_key_row_missing_perm,
)
# Verify new object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "recreated_perm_id_789"
# Verify find_unique was called with the missing permission ID
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": "missing_perm_id_789"}
)
# Verify upsert was called to create new record
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()

242
tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py Normal file
View file

@ -0,0 +1,242 @@
import asyncio
import json
import os
import sys
import uuid
from typing import Optional, cast
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
sys.path.insert(
0, os.path.abspath("../../../")
) # Adds the parent directory to the system path
@pytest.mark.asyncio
async def test_organization_update_object_permissions_existing_permission(monkeypatch):
"""
Test updating object permissions when an organization already has an existing object_permission_id.
This test verifies that when updating vector stores for an organization that already has an
object_permission_id, the existing LiteLLM_ObjectPermissionTable record is updated
with the new permissions and the object_permission_id remains the same.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_OrganizationTable,
)
from litellm.proxy.management_endpoints.organization_endpoints import (
handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Mock existing organization with object_permission_id
existing_organization_row = LiteLLM_OrganizationTable(
organization_id="test_org_id",
object_permission_id="existing_perm_id_123",
organization_alias="test_org",
budget_id="test_budget_id",
models=["test_model_1", "test_model_2"],
created_by="test_created_by",
updated_by="test_updated_by",
)
# Mock existing object permission record
existing_object_permission = MagicMock()
existing_object_permission.model_dump.return_value = {
"object_permission_id": "existing_perm_id_123",
"vector_stores": ["old_store_1", "old_store_2"],
}
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=existing_object_permission
)
# Mock upsert operation
updated_permission = MagicMock()
updated_permission.object_permission_id = "existing_perm_id_123"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=updated_permission
)
# Test data with new object permission
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["new_store_1", "new_store_2", "new_store_3"]
).model_dump(exclude_unset=True, exclude_none=True),
"organization_alias": "updated_org",
}
# Call the function
result = await handle_update_object_permission(
data_json=data_json,
existing_organization_row=existing_organization_row,
)
# Verify the object_permission was removed from data_json and object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "existing_perm_id_123"
# Verify database operations were called correctly
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": "existing_perm_id_123"}
)
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
@pytest.mark.asyncio
async def test_organization_update_object_permissions_no_existing_permission(
monkeypatch,
):
"""
Test creating object permissions when an organization has no existing object_permission_id.
This test verifies that when updating object permissions for an organization that has
object_permission_id set to None, a new entry is created in the
LiteLLM_ObjectPermissionTable and the organization is updated with the new object_permission_id.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_OrganizationTable,
)
from litellm.proxy.management_endpoints.organization_endpoints import (
handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
existing_organization_row_no_perm = LiteLLM_OrganizationTable(
organization_id="test_org_id_2",
object_permission_id=None,
organization_alias="test_org_2",
budget_id="test_budget_id_2",
models=["test_model_1", "test_model_2"],
created_by="test_created_by_2",
updated_by="test_updated_by_2",
)
# Mock find_unique to return None (no existing permission)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=None
)
# Mock upsert to create new record
new_permission = MagicMock()
new_permission.object_permission_id = "new_perm_id_456"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=new_permission
)
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["brand_new_store"]
).model_dump(exclude_unset=True, exclude_none=True),
"organization_alias": "updated_org_2",
}
result = await handle_update_object_permission(
data_json=data_json,
existing_organization_row=existing_organization_row_no_perm,
)
# Verify new object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "new_perm_id_456"
# Verify find_unique was called with None
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": None}
)
# Verify upsert was called to create new record
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
@pytest.mark.asyncio
async def test_organization_update_object_permissions_missing_permission_record(
monkeypatch,
):
"""
Test creating object permissions when existing object_permission_id record is not found.
This test verifies that when updating object permissions for an organization that has an
object_permission_id but the corresponding record cannot be found in the database,
a new entry is created in the LiteLLM_ObjectPermissionTable with the new permissions.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_OrganizationTable,
)
from litellm.proxy.management_endpoints.organization_endpoints import (
handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
existing_organization_row_missing_perm = LiteLLM_OrganizationTable(
organization_id="test_org_id_3",
object_permission_id="missing_perm_id_789",
organization_alias="test_org_3",
budget_id="test_budget_id_3",
models=["test_model_1", "test_model_2"],
created_by="test_created_by_3",
updated_by="test_updated_by_3",
)
# Mock find_unique to return None (permission record not found)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=None
)
# Mock upsert to create new record
new_permission = MagicMock()
new_permission.object_permission_id = "recreated_perm_id_789"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=new_permission
)
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["recreated_store"]
).model_dump(exclude_unset=True, exclude_none=True),
"organization_alias": "updated_org_3",
}
result = await handle_update_object_permission(
data_json=data_json,
existing_organization_row=existing_organization_row_missing_perm,
)
# Verify new object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "recreated_perm_id_789"
# Verify find_unique was called with the missing permission ID
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": "missing_perm_id_789"}
)
# Verify upsert was called to create new record
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()

202
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -313,3 +313,205 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth):
assert mock_team_create.call_count == 1
created_team_kwargs = mock_team_create.call_args.kwargs
assert created_team_kwargs["data"].get("object_permission_id") == "objperm123"
@pytest.mark.asyncio
async def test_team_update_object_permissions_existing_permission(monkeypatch):
"""
Test updating object permissions when a team already has an existing object_permission_id.
This test verifies that when updating vector stores for a team that already has an
object_permission_id, the existing LiteLLM_ObjectPermissionTable record is updated
with the new permissions and the object_permission_id remains the same.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_TeamTable
from litellm.proxy.management_endpoints.team_endpoints import (
handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Mock existing team with object_permission_id
existing_team_row = LiteLLM_TeamTable(
team_id="test_team_id",
object_permission_id="existing_perm_id_123",
team_alias="test_team",
)
# Mock existing object permission record
existing_object_permission = MagicMock()
existing_object_permission.model_dump.return_value = {
"object_permission_id": "existing_perm_id_123",
"vector_stores": ["old_store_1", "old_store_2"],
}
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=existing_object_permission
)
# Mock upsert operation
updated_permission = MagicMock()
updated_permission.object_permission_id = "existing_perm_id_123"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=updated_permission
)
# Test data with new object permission
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["new_store_1", "new_store_2", "new_store_3"]
).model_dump(exclude_unset=True, exclude_none=True),
"team_alias": "updated_team",
}
# Call the function
result = await handle_update_object_permission(
data_json=data_json,
existing_team_row=existing_team_row,
)
# Verify the object_permission was removed from data_json and object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "existing_perm_id_123"
# Verify database operations were called correctly
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": "existing_perm_id_123"}
)
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
@pytest.mark.asyncio
async def test_team_update_object_permissions_no_existing_permission(monkeypatch):
"""
Test creating object permissions when a team has no existing object_permission_id.
This test verifies that when updating object permissions for a team that has
object_permission_id set to None, a new entry is created in the
LiteLLM_ObjectPermissionTable and the team is updated with the new object_permission_id.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_TeamTable
from litellm.proxy.management_endpoints.team_endpoints import (
handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
existing_team_row_no_perm = LiteLLM_TeamTable(
team_id="test_team_id_2",
object_permission_id=None,
team_alias="test_team_2",
)
# Mock find_unique to return None (no existing permission)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=None
)
# Mock upsert to create new record
new_permission = MagicMock()
new_permission.object_permission_id = "new_perm_id_456"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=new_permission
)
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["brand_new_store"]
).model_dump(exclude_unset=True, exclude_none=True),
"team_alias": "updated_team_2",
}
result = await handle_update_object_permission(
data_json=data_json,
existing_team_row=existing_team_row_no_perm,
)
# Verify new object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "new_perm_id_456"
# Verify find_unique was called with None
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": None}
)
# Verify upsert was called to create new record
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
@pytest.mark.asyncio
async def test_team_update_object_permissions_missing_permission_record(monkeypatch):
"""
Test creating object permissions when existing object_permission_id record is not found.
This test verifies that when updating object permissions for a team that has an
object_permission_id but the corresponding record cannot be found in the database,
a new entry is created in the LiteLLM_ObjectPermissionTable with the new permissions.
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_TeamTable
from litellm.proxy.management_endpoints.team_endpoints import (
handle_update_object_permission,
)
# Mock prisma client
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
existing_team_row_missing_perm = LiteLLM_TeamTable(
team_id="test_team_id_3",
object_permission_id="missing_perm_id_789",
team_alias="test_team_3",
)
# Mock find_unique to return None (permission record not found)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
return_value=None
)
# Mock upsert to create new record
new_permission = MagicMock()
new_permission.object_permission_id = "recreated_perm_id_789"
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
return_value=new_permission
)
data_json = {
"object_permission": LiteLLM_ObjectPermissionBase(
vector_stores=["recreated_store"]
).model_dump(exclude_unset=True, exclude_none=True),
"team_alias": "updated_team_3",
}
result = await handle_update_object_permission(
data_json=data_json,
existing_team_row=existing_team_row_missing_perm,
)
# Verify new object_permission_id was set
assert "object_permission" not in result
assert result["object_permission_id"] == "recreated_perm_id_789"
# Verify find_unique was called with the missing permission ID
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
where={"object_permission_id": "missing_perm_id_789"}
)
# Verify upsert was called to create new record
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()

14
ui/litellm-dashboard/src/components/key_edit_view.tsx
View file

@ -5,6 +5,8 @@ import { KeyResponse } from "./key_team_helpers/key_list";
import { fetchTeamModels } from "../components/create_key_button";
import { modelAvailableCall } from "./networking";
import NumericalInput from "./shared/numerical_input";
import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
interface KeyEditViewProps {
keyData: KeyResponse;
onCancel: () => void;
@ -92,7 +94,8 @@ export function KeyEditView({
...keyData,
budget_duration: getBudgetDuration(keyData.budget_duration),
metadata: keyData.metadata ? JSON.stringify(keyData.metadata, null, 2) : "",
guardrails: keyData.metadata?.guardrails || []
guardrails: keyData.metadata?.guardrails || [],
vector_stores: keyData.object_permission?.vector_stores || []
};
return (
@ -165,6 +168,15 @@ export function KeyEditView({
/>
</Form.Item>
<Form.Item label="Vector Stores" name="vector_stores">
<VectorStoreSelector
onChange={(values) => form.setFieldValue('vector_stores', values)}
value={form.getFieldValue('vector_stores')}
accessToken={accessToken || ""}
placeholder="Select vector stores"
/>
</Form.Item>
<Form.Item label="Metadata" name="metadata">
<Input.TextArea rows={10} />
</Form.Item>

26
ui/litellm-dashboard/src/components/key_info_view.tsx
View file

@ -22,6 +22,7 @@ import { Form, Input, InputNumber, message, Select } from "antd";
import { KeyEditView } from "./key_edit_view";
import { RegenerateKeyModal } from "./regenerate_key_modal";
import { rolesWithWriteAccess } from '../utils/roles';
import ObjectPermissionsView from "./object_permissions_view";
interface KeyInfoViewProps {
keyId: string;
@ -64,6 +65,16 @@ export default function KeyInfoView({ keyId, onClose, keyData, accessToken, user
const currentKey = formValues.token;
formValues.key = currentKey;
// Handle object_permission updates
if (formValues.vector_stores !== undefined) {
formValues.object_permission = {
...keyData.object_permission,
vector_stores: formValues.vector_stores || []
};
// Remove vector_stores from the top level as it should be in object_permission
delete formValues.vector_stores;
}
// Convert metadata back to an object if it exists and is a string
if (formValues.metadata && typeof formValues.metadata === "string") {
try {
@ -249,6 +260,14 @@ export default function KeyInfoView({ keyId, onClose, keyData, accessToken, user
)}
</div>
</Card>
<Card>
<ObjectPermissionsView
objectPermission={keyData.object_permission}
variant="inline"
accessToken={accessToken}
/>
</Card>
</Grid>
</TabPanel>
@ -354,6 +373,13 @@ export default function KeyInfoView({ keyId, onClose, keyData, accessToken, user
{JSON.stringify(keyData.metadata, null, 2)}
</pre>
</div>
<ObjectPermissionsView
objectPermission={keyData.object_permission}
variant="inline"
className="pt-4 border-t border-gray-200"
accessToken={accessToken}
/>
</div>
)}
</Card>

5
ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
View file

@ -75,6 +75,11 @@ export interface KeyResponse {
user_tpm_limit: number;
user_rpm_limit: number;
user_email: string;
object_permission?: {
object_permission_id: string;
mcp_servers: string[];
vector_stores: string[];
};
}
interface KeyListResponse {

5
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -46,6 +46,11 @@ export interface Organization {
teams: any[] | null;
users: any[] | null;
members: any[] | null;
object_permission?: {
object_permission_id: string;
mcp_servers: string[];
vector_stores: string[];
};
}
export interface CredentialItem {

150
ui/litellm-dashboard/src/components/object_permissions_view.tsx Normal file
View file

@ -0,0 +1,150 @@
import React, { useState, useEffect } from "react";
import { Card, Text, Badge } from "@tremor/react";
import { ServerIcon, DatabaseIcon } from "@heroicons/react/outline";
import { vectorStoreListCall } from "./networking";
interface VectorStoreDetails {
vector_store_id: string;
vector_store_name?: string;
}
interface ObjectPermission {
object_permission_id: string;
mcp_servers: string[];
vector_stores: string[];
}
interface ObjectPermissionsViewProps {
objectPermission?: ObjectPermission;
variant?: "card" | "inline";
className?: string;
accessToken?: string | null;
}
export function ObjectPermissionsView({
objectPermission,
variant = "card",
className = "",
accessToken
}: ObjectPermissionsViewProps) {
const vectorStores = objectPermission?.vector_stores || [];
const mcpServers = objectPermission?.mcp_servers || [];
const [vectorStoreDetails, setVectorStoreDetails] = useState<VectorStoreDetails[]>([]);
// Fetch vector store details when component mounts
useEffect(() => {
const fetchVectorStores = async () => {
if (!accessToken || vectorStores.length === 0) return;
try {
const response = await vectorStoreListCall(accessToken);
if (response.data) {
setVectorStoreDetails(response.data.map((store: any) => ({
vector_store_id: store.vector_store_id,
vector_store_name: store.vector_store_name
})));
}
} catch (error) {
console.error("Error fetching vector stores:", error);
}
};
fetchVectorStores();
}, [accessToken, vectorStores.length]);
// Function to get display name for vector store
const getVectorStoreDisplayName = (storeId: string) => {
const storeDetail = vectorStoreDetails.find(store => store.vector_store_id === storeId);
if (storeDetail) {
return `${storeDetail.vector_store_name || storeDetail.vector_store_id} (${storeDetail.vector_store_id})`;
}
return storeId;
};
const content = (
<div className={variant === "card" ? "grid grid-cols-1 md:grid-cols-2 gap-6" : "space-y-4"}>
{/* Vector Stores Section */}
<div className="space-y-3">
<div className="flex items-center gap-2">
<DatabaseIcon className="h-4 w-4 text-blue-600" />
<Text className="font-semibold text-gray-900">Vector Stores</Text>
<Badge color="blue" size="xs">
{vectorStores.length}
</Badge>
</div>
{vectorStores.length > 0 ? (
<div className="flex flex-wrap gap-2">
{vectorStores.map((store, index) => (
<div
key={index}
className="inline-flex items-center px-3 py-1.5 rounded-lg bg-blue-50 border border-blue-200 text-blue-800 text-sm font-medium"
>
{getVectorStoreDisplayName(store)}
</div>
))}
</div>
) : (
<div className="flex items-center gap-2 px-3 py-2 rounded-lg bg-gray-50 border border-gray-200">
<DatabaseIcon className="h-4 w-4 text-gray-400" />
<Text className="text-gray-500 text-sm">No vector stores configured</Text>
</div>
)}
</div>
{/* MCP Servers Section */}
{/* <div className="space-y-3">
<div className="flex items-center gap-2">
<ServerIcon className="h-4 w-4 text-blue-600" />
<Text className="font-semibold text-gray-900">MCP Servers</Text>
<Badge color="blue" size="xs">
{mcpServers.length}
</Badge>
</div>
{mcpServers.length > 0 ? (
<div className="flex flex-wrap gap-2">
{mcpServers.map((server, index) => (
<div
key={index}
className="inline-flex items-center px-3 py-1.5 rounded-lg bg-blue-50 border border-blue-200 text-blue-800 text-sm font-medium"
>
{server}
</div>
))}
</div>
) : (
<div className="flex items-center gap-2 px-3 py-2 rounded-lg bg-gray-50 border border-gray-200">
<ServerIcon className="h-4 w-4 text-gray-400" />
<Text className="text-gray-500 text-sm">No MCP servers configured</Text>
</div>
)}
</div> */}
</div>
);
if (variant === "card") {
return (
<div className={`bg-white border border-gray-200 rounded-lg p-6 ${className}`}>
<div className="flex items-center gap-2 mb-6">
<div>
<Text className="font-semibold text-gray-900">Object Permissions</Text>
<Text className="text-xs text-gray-500">
Access control for Vector Stores
</Text>
</div>
</div>
{content}
</div>
);
}
return (
<div className={`${className}`}>
<Text className="font-medium text-gray-900 mb-3">Object Permissions</Text>
{content}
</div>
);
}
export default ObjectPermissionsView;

35
ui/litellm-dashboard/src/components/organization/organization_view.tsx
View file

@ -28,6 +28,8 @@ import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_
import { Member, Organization, organizationInfoCall, organizationMemberAddCall, organizationMemberUpdateCall, organizationMemberDeleteCall, organizationUpdateCall } from "../networking";
import UserSearchModal from "../common_components/user_search_modal";
import MemberModal from "../team/edit_membership";
import ObjectPermissionsView from "../object_permissions_view";
import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
interface OrganizationInfoProps {
organizationId: string;
@ -139,7 +141,7 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
try {
if (!accessToken) return;
const updateData = {
const updateData: any = {
organization_id: organizationId,
organization_alias: values.organization_alias,
models: values.models,
@ -151,6 +153,14 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
},
metadata: values.metadata ? JSON.parse(values.metadata) : null,
};
// Handle object_permission updates
if (values.vector_stores !== undefined) {
updateData.object_permission = {
...orgData?.object_permission,
vector_stores: values.vector_stores || []
};
}
const response = await organizationUpdateCall(accessToken, updateData);
@ -247,6 +257,12 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
))}
</div>
</Card>
<ObjectPermissionsView
objectPermission={orgData.object_permission}
variant="card"
accessToken={accessToken}
/>
</Grid>
</TabPanel>
@ -346,6 +362,7 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
max_budget: orgData.litellm_budget_table.max_budget,
budget_duration: orgData.litellm_budget_table.budget_duration,
metadata: orgData.metadata ? JSON.stringify(orgData.metadata, null, 2) : "",
vector_stores: orgData.object_permission?.vector_stores || []
}}
layout="vertical"
>
@ -393,6 +410,15 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
<NumericalInput step={1} style={{ width: "100%" }} />
</Form.Item>
<Form.Item label="Vector Stores" name="vector_stores">
<VectorStoreSelector
onChange={(values) => form.setFieldValue('vector_stores', values)}
value={form.getFieldValue('vector_stores')}
accessToken={accessToken || ""}
placeholder="Select vector stores"
/>
</Form.Item>
<Form.Item label="Metadata" name="metadata">
<Input.TextArea rows={4} />
</Form.Item>
@ -440,6 +466,13 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
<div>Max: {orgData.litellm_budget_table.max_budget !== null ? `$${orgData.litellm_budget_table.max_budget}` : 'No Limit'}</div>
<div>Reset: {orgData.litellm_budget_table.budget_duration || 'Never'}</div>
</div>
<ObjectPermissionsView
objectPermission={orgData.object_permission}
variant="inline"
className="pt-4 border-t border-gray-200"
accessToken={accessToken}
/>
</div>
)}
</Card>

40
ui/litellm-dashboard/src/components/team/team_info.tsx
View file

@ -33,6 +33,8 @@ import MemberModal from "./edit_membership";
import UserSearchModal from "@/components/common_components/user_search_modal";
import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
import { isAdminRole } from "@/utils/roles";
import ObjectPermissionsView from "../object_permissions_view";
import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
export interface TeamData {
team_id: string;
@ -58,6 +60,11 @@ export interface TeamData {
model_aliases: Record<string, string>;
} | null;
created_at: string;
object_permission?: {
object_permission_id: string;
mcp_servers: string[];
vector_stores: string[];
};
};
keys: any[];
team_memberships: any[];
@ -209,7 +216,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
return;
}
const updateData = {
const updateData: any = {
team_id: teamId,
team_alias: values.team_alias,
models: values.models,
@ -223,6 +230,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
},
organization_id: values.organization_id,
};
// Handle object_permission updates
if (values.vector_stores !== undefined) {
updateData.object_permission = {
...teamData?.team_info.object_permission,
vector_stores: values.vector_stores || []
};
}
const response = await teamUpdateCall(accessToken, updateData);
@ -306,6 +321,12 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
)}
</div>
</Card>
<ObjectPermissionsView
objectPermission={info.object_permission}
variant="card"
accessToken={accessToken}
/>
</Grid>
</TabPanel>
@ -361,6 +382,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
guardrails: info.metadata?.guardrails || [],
metadata: info.metadata ? JSON.stringify(info.metadata, null, 2) : "",
organization_id: info.organization_id,
vector_stores: info.object_permission?.vector_stores || []
}}
layout="vertical"
>
@ -432,6 +454,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
placeholder="Select or enter guardrails"
/>
</Form.Item>
<Form.Item label="Vector Stores" name="vector_stores">
<VectorStoreSelector
onChange={(values) => form.setFieldValue('vector_stores', values)}
value={form.getFieldValue('vector_stores')}
accessToken={accessToken || ""}
placeholder="Select vector stores"
/>
</Form.Item>
<Form.Item label="Organization ID" name="organization_id">
<Input type=""/>
@ -495,6 +526,13 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
{info.blocked ? 'Blocked' : 'Active'}
</Badge>
</div>
<ObjectPermissionsView
objectPermission={info.object_permission}
variant="inline"
className="pt-4 border-t border-gray-200"
accessToken={accessToken}
/>
</div>
)}
</Card>