diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 71f10f94731..8088edf29ed 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/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[] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7aed02465aa..47c6c0f754f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7170430c6f2..d8c117fe281 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -659,6 +659,11 @@ class ModelParams(LiteLLMPydanticObjectBase): 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""" diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8e0d4774a22..46aebee11d6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/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: diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 50a75b18474..03691d5e976 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/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"], diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 56f6a0042d0..63b0bb6d03f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/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: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 71f10f94731..dc16a0ceaf0 100644 --- a/litellm/proxy/schema.prisma +++ b/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[] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d1e725def55..32332b48c48 100644 --- a/litellm/proxy/utils.py +++ b/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)})"), diff --git a/schema.prisma b/schema.prisma index 71f10f94731..8088edf29ed 100644 --- a/schema.prisma +++ b/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[] diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 46e49e9cc6e..27d2828201a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -60,7 +60,9 @@ async def test_key_token_handling(monkeypatch): """ mock_prisma_client = AsyncMock() mock_insert_data = AsyncMock( - return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None) + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) ) mock_prisma_client.insert_data = mock_insert_data mock_prisma_client.db = MagicMock() @@ -73,7 +75,9 @@ async def test_key_token_handling(monkeypatch): ) mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( - return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None) + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) ) from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles @@ -167,3 +171,83 @@ async def test_budget_reset_at_first_of_month(monkeypatch): response_date.month == expected_month ), f"Expected month {expected_month}, got {response_date.month}" assert response_date.day == 1, f"Expected day 1, got {response_date.day}" + + +@pytest.mark.asyncio +async def test_key_generation_with_object_permission(monkeypatch): + """Ensure /key/generate correctly handles `object_permission` input by + 1. Creating a record in litellm_objectpermissiontable + 2. Passing the returned `object_permission_id` into the key insert payload + """ + # --- Setup mocked prisma client --- + mock_prisma_client = AsyncMock() + + # identity helper for jsonify_object (used inside generate_key_helper_fn) + mock_prisma_client.jsonify_object = lambda data: data # type: ignore + + # Mock the prisma_client.db.litellm_objectpermissiontable.create call + mock_object_permission_create = AsyncMock( + return_value=MagicMock(object_permission_id="objperm123") + ) + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_objectpermissiontable = MagicMock() + mock_prisma_client.db.litellm_objectpermissiontable.create = ( + mock_object_permission_create + ) + + # Mock prisma_client.insert_data for both user and key tables + async def _insert_data_side_effect(*args, **kwargs): # type: ignore + table_name = kwargs.get("table_name") + if table_name == "user": + # minimal attributes accessed later in generate_key_helper_fn + return MagicMock(models=[], spend=0) + elif table_name == "key": + return MagicMock( + token="hashed_token_456", + litellm_budget_table=None, + object_permission=None, + ) + return MagicMock() + + mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect) + + # Attach the mocked prisma client to the proxy_server module + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # --- Import objects after monkeypatching --- + from litellm.proxy._types import ( + GenerateKeyRequest, + LiteLLM_ObjectPermissionTable, + LitellmUserRoles, + ) + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + # --- Call generate_key_fn with object_permission --- + request_data = GenerateKeyRequest( + object_permission=LiteLLM_ObjectPermissionTable(vector_stores=["my-vector"]) + ) + + await generate_key_fn( + data=request_data, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="user-1", + ), + ) + + # --- Assertions --- + # 1. Object permission creation was triggered + mock_object_permission_create.assert_called_once() + + # 2. Key insert received the generated object_permission_id + key_insert_calls = [ + call.kwargs + for call in mock_prisma_client.insert_data.call_args_list + if call.kwargs.get("table_name") == "key" + ] + assert len(key_insert_calls) == 1 + assert key_insert_calls[0]["data"].get("object_permission_id") == "objperm123" 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 3187117145e..19256103052 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -236,3 +236,80 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth): # Clean up dependency override app.dependency_overrides = {} + + +@pytest.mark.asyncio +async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): + """Ensure /team/new correctly handles `object_permission` by + 1. Creating a record in litellm_objectpermissiontable + 2. Passing the returned `object_permission_id` into the team insert payload + """ + # --- Configure mocked prisma client --- + # Helper identity converters used by team logic + mock_db_client.jsonify_team_object = lambda db_data: db_data # type: ignore + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + + # Mock DB structure under prisma_client.db + mock_db_client.db = MagicMock() + + # 1. Mock object permission table creation + mock_object_perm_create = AsyncMock( + return_value=MagicMock(object_permission_id="objperm123") + ) + mock_db_client.db.litellm_objectpermissiontable = MagicMock() + mock_db_client.db.litellm_objectpermissiontable.create = mock_object_perm_create + + # 2. Mock model table creation (may be skipped but provided for safety) + mock_db_client.db.litellm_modeltable = MagicMock() + mock_db_client.db.litellm_modeltable.create = AsyncMock( + return_value=MagicMock(id="model123") + ) + + # 3. Capture team table creation + team_create_result = MagicMock( + team_id="team-456", + object_permission_id="objperm123", + ) + team_create_result.model_dump.return_value = { + "team_id": "team-456", + "object_permission_id": "objperm123", + } + mock_team_create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = mock_team_create + + # 4. Mock user table update behaviour (called for each member) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + # --- Import after mocks applied --- + from fastapi import Request + + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Build request objects + team_request = NewTeamRequest( + team_alias="my-team", + object_permission=LiteLLM_ObjectPermissionTable(vector_stores=["my-vector"]), + ) + + # Pass a dummy FastAPI Request object + dummy_request = MagicMock(spec=Request) + + # Execute the endpoint function + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=mock_admin_auth, + ) + + # --- Assertions --- + # 1. Object permission creation should be called exactly once + mock_object_perm_create.assert_awaited_once() + + # 2. Team creation payload should include the generated object_permission_id + 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" diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index 625a849b25e..851bbc47940 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/create_key_button.tsx @@ -33,6 +33,7 @@ import { getPossibleUserRoles, userFilterUICall, } from "./networking"; +import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import { Team } from "./key_team_helpers/key_list"; import TeamDropdown from "./common_components/team_dropdown"; import { InfoCircleOutlined } from '@ant-design/icons'; @@ -264,6 +265,15 @@ const CreateKey: React.FC = ({ formValues.metadata = JSON.stringify(metadata); } + // Transform allowed_vector_store_ids into object_permission format + if (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) { + formValues.object_permission = { + vector_stores: formValues.allowed_vector_store_ids + }; + // Remove the original field as it's now part of object_permission + delete formValues.allowed_vector_store_ids; + } + const response = await keyCreateCall(accessToken, userID, formValues); console.log("key create Response:", response); @@ -676,6 +686,26 @@ const CreateKey: React.FC = ({ options={guardrailsList.map(name => ({ value: name, label: name }))} /> + + Allowed Vector Stores{' '} + + + + + } + name="allowed_vector_store_ids" + className="mt-4" + help="Select vector stores this key can access. Leave empty for access to all vector stores" + > + form.setFieldValue('allowed_vector_store_ids', values)} + value={form.getFieldValue('allowed_vector_store_ids')} + accessToken={accessToken} + placeholder="Select vector stores (optional)" + /> + = ({ console.log(`values in organizations new create call: ${JSON.stringify(values)}`); + // Transform allowed_vector_store_ids into object_permission + if (values.allowed_vector_store_ids && values.allowed_vector_store_ids.length > 0) { + values.object_permission = { + vector_stores: values.allowed_vector_store_ids + }; + delete values.allowed_vector_store_ids; + } + await organizationCreateCall(accessToken, values); setIsOrgModalVisible(false); form.resetFields(); @@ -214,6 +224,27 @@ const OrganizationsTable: React.FC = ({ + + Allowed Vector Stores{' '} + + + + + } + name="allowed_vector_store_ids" + className="mt-4" + help="Select vector stores this organization can access. Leave empty for access to all vector stores" + > + form.setFieldValue('allowed_vector_store_ids', values)} + value={form.getFieldValue('allowed_vector_store_ids')} + accessToken={accessToken || ''} + placeholder="Select vector stores (optional)" + /> + + @@ -488,6 +519,26 @@ const OrganizationsTable: React.FC = ({ + + Allowed Vector Stores{' '} + + + + + } + name="allowed_vector_store_ids" + className="mt-4" + help="Select vector stores this organization can access. Leave empty for access to all vector stores" + > + form.setFieldValue('allowed_vector_store_ids', values)} + value={form.getFieldValue('allowed_vector_store_ids')} + accessToken={accessToken || ''} + placeholder="Select vector stores (optional)" + /> + diff --git a/ui/litellm-dashboard/src/components/teams.tsx b/ui/litellm-dashboard/src/components/teams.tsx index 8989c033e38..7e8271ef510 100644 --- a/ui/litellm-dashboard/src/components/teams.tsx +++ b/ui/litellm-dashboard/src/components/teams.tsx @@ -57,6 +57,7 @@ import { } from "@tremor/react"; import { CogIcon } from "@heroicons/react/outline"; import AvailableTeamsPanel from "@/components/team/available_teams"; +import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import type { KeyResponse, Team } from "./key_team_helpers/key_list"; const isLocal = process.env.NODE_ENV === "development"; const proxyBaseUrl = isLocal ? "http://localhost:4000" : null; @@ -324,6 +325,13 @@ const Teams: React.FC = ({ } message.info("Creating Team"); + // Transform allowed_vector_store_ids into object_permission + if (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) { + formValues.object_permission = { + vector_stores: formValues.allowed_vector_store_ids + }; + delete formValues.allowed_vector_store_ids; + } const response: any = await teamCreateCall(accessToken, formValues); if (teams !== null) { setTeams([...teams, response]); @@ -1065,6 +1073,26 @@ const Teams: React.FC = ({ options={guardrailsList.map(name => ({ value: name, label: name }))} /> + + Allowed Vector Stores{' '} + + + + + } + name="allowed_vector_store_ids" + className="mt-8" + help="Select vector stores this team can access. Leave empty for access to all vector stores" + > + form.setFieldValue('allowed_vector_store_ids', values)} + value={form.getFieldValue('allowed_vector_store_ids')} + accessToken={accessToken || ''} + placeholder="Select vector stores (optional)" + /> + diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx index 93726e82931..965a0a0599a 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx @@ -8,13 +8,17 @@ interface VectorStoreSelectorProps { value?: string[]; className?: string; accessToken: string; + placeholder?: string; + disabled?: boolean; } const VectorStoreSelector: React.FC = ({ onChange, value, className, - accessToken + accessToken, + placeholder = "Select vector stores", + disabled = false }) => { const [vectorStores, setVectorStores] = useState([]); const [loading, setLoading] = useState(false); @@ -43,7 +47,7 @@ const VectorStoreSelector: React.FC = ({