[Feat] Permission management vector stores on LiteLLM Key, Team, Orgs (#11213)

* fix: init commit for object permissions

* fix: init commit for object permissions

* fix: add vector_store_id to permissions

* fix vector store selector

* feat:add vector store permission mgmt

* feat: ui add allowed vector stores dropdown

* feat: add new vector store object permissions

* testing: key mgmt

* fix: stor vector store permissions on team

* ui select vector store for teams

* ui add vector store settings for orgs

* feat: allow setting org vector store permissions

* test: adding team permissions for vector stores

Author: ishaan-jaff

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -155,6 +155,7 @@ model LiteLLM_UserTable {
 model LiteLLM_ObjectPermissionTable {
   object_permission_id  String         @id @default(uuid())
   mcp_servers           String[]       @default([])
+  vector_stores         String[]       @default([])
   
   teams                 LiteLLM_TeamTable[]
   verification_tokens   LiteLLM_VerificationToken[]

litellm/model_prices_and_context_window_backup.json

@@ -421,6 +421,23 @@
         "supports_system_messages": true,
         "supports_tool_choice": true
     },
+    "gpt-4o-mini-audio-preview": {
+        "max_tokens": 16384,
+        "max_input_tokens": 128000,
+        "max_output_tokens": 16384,
+        "input_cost_per_token": 0.00000015,
+        "input_cost_per_audio_token": 0.00001,
+        "output_cost_per_token": 0.0000006,
+        "output_cost_per_audio_token": 0.00002,
+        "litellm_provider": "openai",
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_audio_input": true,
+        "supports_audio_output": true,
+        "supports_system_messages": true,
+        "supports_tool_choice": true
+    },
     "gpt-4o-mini-audio-preview-2024-12-17": {
         "max_tokens": 16384,
         "max_input_tokens": 128000,

litellm/proxy/_types.py

@@ -659,6 +659,11 @@ def set_model_info(cls, values):
         return values
 
 
+class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase):
+    mcp_servers: Optional[List[str]] = None
+    vector_stores: Optional[List[str]] = None
+
+
 class GenerateRequestBase(LiteLLMPydanticObjectBase):
     """
     Overlapping schema between key and user generate/update requests
@@ -689,6 +694,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
     guardrails: Optional[List[str]] = None
     blocked: Optional[bool] = None
     aliases: Optional[dict] = {}
+    object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
 
 
 class KeyRequestBase(GenerateRequestBase):
@@ -1110,6 +1116,7 @@ class NewTeamRequest(TeamBase):
     model_aliases: Optional[dict] = None
     tags: Optional[list] = None
     guardrails: Optional[List[str]] = None
+    object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -1261,6 +1268,12 @@ class LiteLLM_TeamTable(TeamBase):
     updated_at: Optional[datetime] = None
     created_at: Optional[datetime] = None
 
+    #########################################################
+    # Object Permission - MCP, Vector Stores etc.
+    #########################################################
+    object_permission_id: Optional[str] = None
+    object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
+
     model_config = ConfigDict(protected_namespaces=())
 
     @model_validator(mode="before")
@@ -1338,6 +1351,11 @@ class NewOrganizationRequest(LiteLLM_BudgetTable):
     budget_id: Optional[str] = None
     metadata: Optional[dict] = None
 
+    #########################################################
+    # Object Permission - MCP, Vector Stores etc.
+    #########################################################
+    object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
+
 
 class OrganizationRequest(LiteLLMPydanticObjectBase):
     organizations: List[str]
@@ -1619,6 +1637,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
     created_by: Optional[str] = None
     updated_at: Optional[datetime] = None
     updated_by: Optional[str] = None
+    object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -1761,6 +1780,7 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase):
     metadata: Optional[dict] = None
     created_at: Optional[datetime] = None
     updated_at: Optional[datetime] = None
+    object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
 
     @model_validator(mode="before")
     @classmethod
@@ -1790,6 +1810,12 @@ class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase):
     users: Optional[List[LiteLLM_UserTable]] = None
     litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
 
+    #########################################################
+    # Object Permission - MCP, Vector Stores etc.
+    #########################################################
+    object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
+    object_permission_id: Optional[str] = None
+
 
 class LiteLLM_OrganizationTableWithMembers(LiteLLM_OrganizationTable):
     """Returned by the /organization/info endpoint and /organization/list endpoint"""

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -376,6 +376,7 @@ async def generate_key_fn(  # noqa: PLR0915
     - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
     - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
     - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
+    - object_permission: Optional[LiteLLM_ObjectPermissionTable] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission.
     Examples:
 
     1. Allow users to turn on/off pii masking
@@ -571,6 +572,11 @@ async def generate_key_fn(  # noqa: PLR0915
 
             data_json.pop("tags")
 
+        data_json = await _set_object_permission(
+            data_json=data_json,
+            prisma_client=prisma_client,
+        )
+
         await _enforce_unique_key_alias(
             key_alias=data_json.get("key_alias", None),
             prisma_client=prisma_client,
@@ -641,6 +647,32 @@ def prepare_metadata_fields(
     return non_default_values
 
 
+async def _set_object_permission(
+    data_json: dict,
+    prisma_client: Optional[PrismaClient],
+):
+    """
+    Creates the LiteLLM_ObjectPermissionTable record for the key.
+    - Handles permissions for vector stores and mcp servers.
+    """
+    if prisma_client is None:
+        return data_json
+
+    if "object_permission" in data_json:
+        created_object_permission = (
+            await prisma_client.db.litellm_objectpermissiontable.create(
+                data=data_json["object_permission"],
+            )
+        )
+        data_json[
+            "object_permission_id"
+        ] = created_object_permission.object_permission_id
+
+        # delete the object_permission from the data_json
+        data_json.pop("object_permission")
+    return data_json
+
+
 def prepare_key_update_data(
     data: Union[UpdateKeyRequest, RegenerateKeyRequest], existing_key_row
 ):
@@ -1273,6 +1305,9 @@ async def generate_key_helper_fn(  # noqa: PLR0915
     updated_by: Optional[str] = None,
     allowed_routes: Optional[list] = None,
     sso_user_id: Optional[str] = None,
+    object_permission_id: Optional[
+        str
+    ] = None,  # object_permission_id <-> LiteLLM_ObjectPermissionTable
 ):
     from litellm.proxy.proxy_server import (
         litellm_proxy_budget_name,
@@ -1351,6 +1386,7 @@ async def generate_key_helper_fn(  # noqa: PLR0915
             "budget_reset_at": reset_at,
             "allowed_cache_controls": allowed_cache_controls,
             "sso_user_id": sso_user_id,
+            "object_permission_id": object_permission_id,
         }
         if teams is not None:
             user_data["teams"] = teams
@@ -1379,6 +1415,7 @@ async def generate_key_helper_fn(  # noqa: PLR0915
             "created_by": created_by,
             "updated_by": updated_by,
             "allowed_routes": allowed_routes or [],
+            "object_permission_id": object_permission_id,
         }
 
         if (
@@ -1452,6 +1489,9 @@ async def generate_key_helper_fn(  # noqa: PLR0915
             key_data["litellm_budget_table"] = getattr(
                 create_key_response, "litellm_budget_table", None
             )
+            key_data["object_permission"] = getattr(
+                create_key_response, "object_permission", None
+            )
             key_data["created_at"] = getattr(create_key_response, "created_at", None)
             key_data["updated_at"] = getattr(create_key_response, "updated_at", None)
     except Exception as e:

litellm/proxy/management_endpoints/organization_endpoints.py

@@ -168,6 +168,12 @@ async def new_organization(
 
         data.budget_id = _budget.budget_id
 
+    ## Handle Object Permission - MCP, Vector Stores etc.
+    object_permission_id = await _set_object_permission(
+        data=data,
+        prisma_client=prisma_client,
+    )
+
     """
     Ensure only models that user has access to, are given to org
     """
@@ -189,6 +195,7 @@ async def new_organization(
 
     organization_row = LiteLLM_OrganizationTable(
         **data.json(exclude_none=True),
+        object_permission_id=object_permission_id,
         created_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
         updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
     )
@@ -207,6 +214,30 @@ async def new_organization(
     return response
 
 
+async def _set_object_permission(
+    data: NewOrganizationRequest,
+    prisma_client: Optional[PrismaClient],
+) -> Optional[str]:
+    """
+    Creates the LiteLLM_ObjectPermissionTable record for the organization.
+    - Handles permissions for vector stores and mcp servers.
+
+    Returns the object_permission_id if created, otherwise None.
+    """
+    if prisma_client is None:
+        return None
+
+    if data.object_permission is not None:
+        created_object_permission = (
+            await prisma_client.db.litellm_objectpermissiontable.create(
+                data=data.object_permission.model_dump(exclude_none=True),
+            )
+        )
+        del data.object_permission
+        return created_object_permission.object_permission_id
+    return None
+
+
 @router.patch(
     "/organization/update",
     tags=["organization management"],

litellm/proxy/management_endpoints/team_endpoints.py

@@ -298,10 +298,17 @@ async def new_team(  # noqa: PLR0915
 
             _model_id = model_dict.id
 
+        ## Handle Object Permission - MCP, Vector Stores etc.
+        object_permission_id = await _set_object_permission(
+            data=data,
+            prisma_client=prisma_client,
+        )
+
         ## ADD TO TEAM TABLE
         complete_team_data = LiteLLM_TeamTable(
             **data.json(),
             model_id=_model_id,
+            object_permission_id=object_permission_id,
         )
 
         # Set Management Endpoint Metadata Fields
@@ -411,6 +418,30 @@ async def _update_model_table(
     return _model_id
 
 
+async def _set_object_permission(
+    data: NewTeamRequest,
+    prisma_client: Optional[PrismaClient],
+) -> Optional[str]:
+    """
+    Creates the LiteLLM_ObjectPermissionTable record for the team.
+    - Handles permissions for vector stores and mcp servers.
+
+    Returns the object_permission_id if created, otherwise None.
+    """
+    if prisma_client is None:
+        return None
+
+    if data.object_permission is not None:
+        created_object_permission = (
+            await prisma_client.db.litellm_objectpermissiontable.create(
+                data=data.object_permission.model_dump(exclude_none=True),
+            )
+        )
+        del data.object_permission
+        return created_object_permission.object_permission_id
+    return None
+
+
 def validate_team_org_change(
     team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTable, llm_router: Router
 ) -> bool:

litellm/proxy/schema.prisma

@@ -155,7 +155,7 @@ model LiteLLM_UserTable {
 model LiteLLM_ObjectPermissionTable {
   object_permission_id  String         @id @default(uuid())
   mcp_servers           String[]       @default([])
-  
+  vector_stores         String[]       @default([])
   teams                 LiteLLM_TeamTable[]
   verification_tokens   LiteLLM_VerificationToken[]
   organizations         LiteLLM_OrganizationTable[]

litellm/proxy/utils.py

@@ -2798,6 +2798,8 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException:
     """
     from fastapi import status
 
+    verbose_proxy_logger.exception(f"Exception: {e}")
+
     if isinstance(e, HTTPException):
         return ProxyException(
             message=getattr(e, "detail", f"error({str(e)})"),

schema.prisma

@@ -155,6 +155,7 @@ model LiteLLM_UserTable {
 model LiteLLM_ObjectPermissionTable {
   object_permission_id  String         @id @default(uuid())
   mcp_servers           String[]       @default([])
+  vector_stores         String[]       @default([])
   
   teams                 LiteLLM_TeamTable[]
   verification_tokens   LiteLLM_VerificationToken[]