From 68fd17d15ecaadfad67de524842cc12e8eeef51d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 31 May 2025 09:41:05 -0700 Subject: [PATCH] [Fix] QA Fixes - Vector Store Object Permissions (#11291) * fix: QA for key,team,org permissions * fix: add_vector_store_to_registry * fix: refactor bedrock guard * fix: refactor using us east 1 with vector stores * fix: code QA checks * fix: testing for mgmt endpoints --- .../vector_stores/bedrock_vector_store.py | 54 +++++--- litellm/llms/bedrock/base_aws_llm.py | 30 +++++ .../guardrail_hooks/bedrock_guardrails.py | 21 +--- .../key_management_endpoints.py | 53 +++----- .../organization_endpoints.py | 60 ++------- .../management_endpoints/team_endpoints.py | 53 +++----- .../object_permission_utils.py | 93 ++++++++++++++ .../vector_stores/vector_store_registry.py | 19 ++- .../test_key_management_endpoints.py | 6 - .../test_organization_endpoints.py | 5 - .../test_team_endpoints.py | 5 - .../test_vector_store_registry.py | 115 ++++++++++++++++++ 12 files changed, 337 insertions(+), 177 deletions(-) create mode 100644 litellm/proxy/management_helpers/object_permission_utils.py create mode 100644 tests/test_litellm/vector_stores/test_vector_store_registry.py diff --git a/litellm/integrations/vector_stores/bedrock_vector_store.py b/litellm/integrations/vector_stores/bedrock_vector_store.py index 9015757000..0523dac8ed 100644 --- a/litellm/integrations/vector_stores/bedrock_vector_store.py +++ b/litellm/integrations/vector_stores/bedrock_vector_store.py @@ -34,7 +34,6 @@ from litellm.types.vector_stores import ( VectorStoreSearchResponse, VectorStoreSearchResult, ) -from litellm.utils import load_credentials_from_list if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -258,22 +257,49 @@ class BedrockVectorStore(BaseVectorStore, BaseAWSLLM): from fastapi import HTTPException non_default_params = non_default_params or {} - load_credentials_from_list(kwargs=non_default_params) + credentials_dict: Dict[str, Any] = {} + if litellm.vector_store_registry is not None: + credentials_dict = ( + litellm.vector_store_registry.get_credentials_for_vector_store( + knowledge_base_id + ) + ) + credentials = self.get_credentials( - aws_access_key_id=non_default_params.get("aws_access_key_id", None), - aws_secret_access_key=non_default_params.get("aws_secret_access_key", None), - aws_session_token=non_default_params.get("aws_session_token", None), - aws_region_name=non_default_params.get("aws_region_name", None), - aws_session_name=non_default_params.get("aws_session_name", None), - aws_profile_name=non_default_params.get("aws_profile_name", None), - aws_role_name=non_default_params.get("aws_role_name", None), - aws_web_identity_token=non_default_params.get( - "aws_web_identity_token", None + aws_access_key_id=credentials_dict.get( + "aws_access_key_id", non_default_params.get("aws_access_key_id", None) + ), + aws_secret_access_key=credentials_dict.get( + "aws_secret_access_key", + non_default_params.get("aws_secret_access_key", None), + ), + aws_session_token=credentials_dict.get( + "aws_session_token", non_default_params.get("aws_session_token", None) + ), + aws_region_name=credentials_dict.get( + "aws_region_name", non_default_params.get("aws_region_name", None) + ), + aws_session_name=credentials_dict.get( + "aws_session_name", non_default_params.get("aws_session_name", None) + ), + aws_profile_name=credentials_dict.get( + "aws_profile_name", non_default_params.get("aws_profile_name", None) + ), + aws_role_name=credentials_dict.get( + "aws_role_name", non_default_params.get("aws_role_name", None) + ), + aws_web_identity_token=credentials_dict.get( + "aws_web_identity_token", + non_default_params.get("aws_web_identity_token", None), + ), + aws_sts_endpoint=credentials_dict.get( + "aws_sts_endpoint", non_default_params.get("aws_sts_endpoint", None) ), - aws_sts_endpoint=non_default_params.get("aws_sts_endpoint", None), ) - aws_region_name = self._get_aws_region_name( - optional_params=self.optional_params + aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( + aws_region_name=credentials_dict.get( + "aws_region_name", non_default_params.get("aws_region_name", None) + ), ) # Prepare request data diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index fea2072a0f..337794f162 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -336,6 +336,36 @@ class BaseAWSLLM: return aws_region_name + def get_aws_region_name_for_non_llm_api_calls( + self, + aws_region_name: Optional[str] = None, + ): + """ + Get the AWS region name for non-llm api calls. + + LLM API calls check the model arn and end up using that as the region name. + + For non-llm api calls eg. Guardrails, Vector Stores we just need to check the dynamic param or env vars. + """ + if aws_region_name is None: + # check env # + litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) + + if litellm_aws_region_name is not None and isinstance( + litellm_aws_region_name, str + ): + aws_region_name = litellm_aws_region_name + + standard_aws_region_name = get_secret("AWS_REGION", None) + if standard_aws_region_name is not None and isinstance( + standard_aws_region_name, str + ): + aws_region_name = standard_aws_region_name + + if aws_region_name is None: + aws_region_name = "us-west-2" + return aws_region_name + @tracer.wrap() def _auth_with_web_identity_token( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 1432e44184..fb8e9236f4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -30,7 +30,6 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.secret_managers.main import get_secret from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( @@ -129,23 +128,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): aws_sts_endpoint = self.optional_params.get("aws_sts_endpoint", None) ### SET REGION NAME ### - if aws_region_name is None: - # check env # - litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) - - if litellm_aws_region_name is not None and isinstance( - litellm_aws_region_name, str - ): - aws_region_name = litellm_aws_region_name - - standard_aws_region_name = get_secret("AWS_REGION", None) - if standard_aws_region_name is not None and isinstance( - standard_aws_region_name, str - ): - aws_region_name = standard_aws_region_name - - if aws_region_name is None: - aws_region_name = "us-west-2" + aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( + aws_region_name=aws_region_name, + ) credentials: Credentials = self.get_credentials( aws_access_key_id=aws_access_key_id, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2fee9edf1c..c744cd0097 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -45,6 +45,9 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) +from litellm.proxy.management_helpers.object_permission_utils import ( + handle_update_object_permission_common, +) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) @@ -733,48 +736,20 @@ async def _handle_update_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}, - ) + # Use the common helper to handle the object permission update + object_permission_id = await handle_update_object_permission_common( + data_json=data_json, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, ) - 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, - }, + # Add the object_permission_id to data_json if one was created/updated + if object_permission_id is not None: + data_json["object_permission_id"] = object_permission_id + verbose_proxy_logger.debug( + f"updated object_permission_id: {object_permission_id}" ) - ) - 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 diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index e057ff32a1..309438a17f 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -24,6 +24,9 @@ from litellm.proxy.management_endpoints.budget_management_endpoints import ( new_budget, update_budget, ) +from litellm.proxy.management_helpers.object_permission_utils import ( + handle_update_object_permission_common, +) from litellm.proxy.management_helpers.utils import ( get_new_internal_user_defaults, management_endpoint_wrapper, @@ -309,61 +312,20 @@ async def handle_update_object_permission( - 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}, - ) + # Use the common helper to handle the object permission update + object_permission_id = await handle_update_object_permission_common( + data_json=data_json, + existing_object_permission_id=existing_organization_row.object_permission_id, + prisma_client=prisma_client, ) - # 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 - ) + # Add the object_permission_id to data_json if one was created/updated + if object_permission_id is not None: + data_json["object_permission_id"] = object_permission_id - 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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 1ced6c3218..5bb248670b 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -70,6 +70,9 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) +from litellm.proxy.management_helpers.object_permission_utils import ( + handle_update_object_permission_common, +) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) @@ -752,48 +755,20 @@ async def handle_update_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_team_row.object_permission_id - existing_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": existing_object_permission_id}, - ) + # Use the common helper to handle the object permission update + object_permission_id = await handle_update_object_permission_common( + data_json=data_json, + existing_object_permission_id=existing_team_row.object_permission_id, + prisma_client=prisma_client, ) - existing_object_permissions_dict: Dict = {} + # Add the object_permission_id to data_json if one was created/updated + if object_permission_id is not None: + data_json["object_permission_id"] = object_permission_id + verbose_proxy_logger.debug( + f"updated object_permission_id: {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 - ) - 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 diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py new file mode 100644 index 0000000000..90ae313c6f --- /dev/null +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -0,0 +1,93 @@ +""" +Common utility functions for handling object permission updates across +organizations, teams, and keys. +""" + +import json +import uuid +from typing import Dict, Optional, Union + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.utils import PrismaClient + + +async def handle_update_object_permission_common( + data_json: Dict, + existing_object_permission_id: Optional[str], + prisma_client: Optional[PrismaClient], +) -> Optional[str]: + """ + Common logic for handling object permission updates across organizations, teams, and keys. + + This function: + 1. Extracts `object_permission` from data_json + 2. Looks up existing object permission if it exists + 3. Merges new permissions with existing ones + 4. Upserts to the LiteLLM_ObjectPermissionTable + 5. Returns the object_permission_id + + Args: + data_json: The data dictionary containing the object_permission to update + existing_object_permission_id: The current object_permission_id from the entity (can be None) + prisma_client: The database client + + Returns: + Optional[str]: The object_permission_id after the update/creation, or None if no object_permission to process + + Raises: + ValueError: If prisma_client is None + """ + 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", None) + if new_object_permission is None: + return None + + # Lookup existing object permission ID and update that entry + object_permission_id_to_use: str = existing_object_permission_id or str( + uuid.uuid4() + ) + existing_object_permissions_dict: Dict = {} + + existing_object_permission = ( + await prisma_client.db.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": object_permission_id_to_use}, + ) + ) + + # 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 + ) + + # Handle string JSON object permission + 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": object_permission_id_to_use}, + data={ + "create": existing_object_permissions_dict, + "update": existing_object_permissions_dict, + }, + ) + ) + + verbose_proxy_logger.debug( + f"created_object_permission_row: {created_object_permission_row}" + ) + + return created_object_permission_row.object_permission_id diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index e16799ecce..a045406d78 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -169,8 +169,8 @@ class VectorStoreRegistry: Only add the vector store if it is not already in the registry """ vector_store_id = vector_store.get("vector_store_id") - for vector_store in self.vector_stores: - if vector_store.get("vector_store_id") == vector_store_id: + for _vector_store in self.vector_stores: + if _vector_store.get("vector_store_id") == vector_store_id: return self.vector_stores.append(vector_store) @@ -209,3 +209,18 @@ class VectorStoreRegistry: ) vector_stores_from_db.append(_litellm_managed_vector_store) return vector_stores_from_db + + def get_credentials_for_vector_store(self, vector_store_id: str) -> Dict[str, Any]: + """ + Get the credentials for a vector store + + Returns a dictionary of unpacked credentials for the vector store to use for the request + """ + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + + for vector_store in self.vector_stores: + if vector_store.get("vector_store_id") == vector_store_id: + credentials = vector_store.get("litellm_credential_name") + if credentials: + return CredentialAccessor.get_credential_values(credentials) + return {} 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 f0a2f7cab7..ddecb94ab6 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 @@ -388,12 +388,6 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) # 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() diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index a602a6e338..c05e528846 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -159,11 +159,6 @@ async def test_organization_update_object_permissions_no_existing_permission( 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() 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 7b9acf42e1..c7addf5ded 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -443,11 +443,6 @@ async def test_team_update_object_permissions_no_existing_permission(monkeypatch 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() diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/test_litellm/vector_stores/test_vector_store_registry.py new file mode 100644 index 0000000000..785b55d7e0 --- /dev/null +++ b/tests/test_litellm/vector_stores/test_vector_store_registry.py @@ -0,0 +1,115 @@ +import json +import os +import sys +from unittest.mock import patch + +import httpx +import pytest +import respx +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from datetime import datetime, timezone + +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore +from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + + +def test_get_credentials_for_vector_store(): + """Test that get_credentials_for_vector_store returns correct credentials""" + # Create test vector stores + vector_store_1 = LiteLLM_ManagedVectorStore( + vector_store_id="test_id_1", + custom_llm_provider="openai", + vector_store_name="test_store_1", + litellm_credential_name="test_creds_1", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + vector_store_2 = LiteLLM_ManagedVectorStore( + vector_store_id="test_id_2", + custom_llm_provider="bedrockc", + vector_store_name="test_store_2", + litellm_credential_name="test_creds_2", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + # Create registry with vector stores + registry = VectorStoreRegistry([vector_store_1, vector_store_2]) + + # Mock CredentialAccessor.get_credential_values + with patch( + "litellm.litellm_core_utils.credential_accessor.CredentialAccessor.get_credential_values" + ) as mock_get_creds: + mock_get_creds.return_value = {"api_key": "test_key_1", "env": "test"} + + # Test getting credentials for existing vector store + result = registry.get_credentials_for_vector_store("test_id_1") + + assert result == {"api_key": "test_key_1", "env": "test"} + mock_get_creds.assert_called_once_with("test_creds_1") + + # Test getting credentials for non-existent vector store + result = registry.get_credentials_for_vector_store("non_existent_id") + assert result == {} + + +def test_add_vector_store_to_registry(): + """Test that add_vector_store_to_registry adds vector store correctly when there are pre-existing stores""" + # Create pre-existing vector stores + existing_store_1 = LiteLLM_ManagedVectorStore( + vector_store_id="existing_id_1", + custom_llm_provider="openai", + vector_store_name="existing_store_1", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + existing_store_2 = LiteLLM_ManagedVectorStore( + vector_store_id="existing_id_2", + custom_llm_provider="openai", + vector_store_name="existing_store_2", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + # Create registry with pre-existing stores + registry = VectorStoreRegistry([existing_store_1, existing_store_2]) + assert len(registry.vector_stores) == 2 + + # Add a new vector store + new_store = LiteLLM_ManagedVectorStore( + vector_store_id="new_id", + custom_llm_provider="bedrock", + vector_store_name="new_store", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + registry.add_vector_store_to_registry(new_store) + + # Verify new store was added + assert len(registry.vector_stores) == 3 + assert registry.vector_stores[2]["vector_store_id"] == "new_id" + assert registry.vector_stores[2]["vector_store_name"] == "new_store" + + # Try to add duplicate - should not be added + duplicate_store = LiteLLM_ManagedVectorStore( + vector_store_id="existing_id_1", # Same ID as existing store + custom_llm_provider="different_provider", + vector_store_name="duplicate_store", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + registry.add_vector_store_to_registry(duplicate_store) + + # Verify duplicate was not added + assert len(registry.vector_stores) == 3 + # Original store should still be there unchanged + assert registry.vector_stores[0]["vector_store_name"] == "existing_store_1"