diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index a9902a0d27c..044315168d2 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -8,6 +8,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( VectorStoreEmbeddingExecutor, ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.types.rag import RAGIngestEmbeddingOptions from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( VECTOR_STORE_OPENAI_PARAMS, @@ -26,6 +27,50 @@ else: _DEFAULT_QUERY_EMBEDDING_MODEL: Final = "text-embedding-3-small" _DEFAULT_TOP_K: Final = 5 +S3_VECTORS_STORE_ID_ERROR: Final = ( + "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " + "or vector_bucket_name must be provided in litellm_params" +) + + +def split_s3_vectors_store_id(vector_store_id: str, fallback_bucket_name: object) -> tuple[str, str]: + id_bucket_name, separator, id_index_name = vector_store_id.partition(":") + bucket_name: Final = id_bucket_name if separator else fallback_bucket_name + index_name: Final = id_index_name if separator else vector_store_id + if not isinstance(bucket_name, str) or not bucket_name or not index_name: + raise ValueError(S3_VECTORS_STORE_ID_ERROR) + return bucket_name, index_name + + +def _non_empty_str(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def s3_vectors_ingest_target(vector_store_config: Mapping[str, object]) -> tuple[str, str | None]: + explicit_bucket_name: Final = _non_empty_str(vector_store_config.get("vector_bucket_name")) + explicit_index_name: Final = _non_empty_str(vector_store_config.get("index_name")) + vector_store_id: Final = _non_empty_str(vector_store_config.get("vector_store_id")) + if vector_store_id is None: + if explicit_bucket_name is None: + raise ValueError(S3_VECTORS_STORE_ID_ERROR) + return explicit_bucket_name, explicit_index_name + derived_bucket_name, derived_index_name = split_s3_vectors_store_id(vector_store_id, explicit_bucket_name) + return explicit_bucket_name or derived_bucket_name, explicit_index_name or derived_index_name + + +def s3_vectors_configured_embedding_model(litellm_params: Mapping[str, object]) -> str | None: + return _non_empty_str(litellm_params.get("litellm_embedding_model") or litellm_params.get("embedding_model")) + + +def s3_vectors_ingest_embedding_options( + vector_store_config: Mapping[str, object], + embedding_options: RAGIngestEmbeddingOptions | None, +) -> RAGIngestEmbeddingOptions | None: + store_embedding_model: Final = s3_vectors_configured_embedding_model(vector_store_config) + if store_embedding_model is None: + return embedding_options + store_embedding_options: Final[RAGIngestEmbeddingOptions] = {"model": store_embedding_model} + return store_embedding_options class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM): @@ -69,21 +114,11 @@ class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM @staticmethod def query_embedding_model(litellm_params: Mapping[str, object]) -> str: - configured: Final = litellm_params.get("litellm_embedding_model") or litellm_params.get("embedding_model") - return configured if isinstance(configured, str) and configured else _DEFAULT_QUERY_EMBEDDING_MODEL + return s3_vectors_configured_embedding_model(litellm_params) or _DEFAULT_QUERY_EMBEDDING_MODEL @staticmethod def _query_target(vector_store_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]: - if ":" in vector_store_id: - bucket_name, index_name = vector_store_id.split(":", 1) - return bucket_name, index_name - bucket_name_from_params: Final = litellm_params.get("vector_bucket_name") - if not isinstance(bucket_name_from_params, str) or not bucket_name_from_params: - raise ValueError( - "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " - "or vector_bucket_name must be provided in litellm_params" - ) - return bucket_name_from_params, vector_store_id + return split_s3_vectors_store_id(vector_store_id, litellm_params.get("vector_bucket_name")) @staticmethod def _query_request( diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index c09f9c755ed..4f0c9f42421 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.vector_store_endpoints.endpoints import ( from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) +from litellm.rag.main import get_ingestion_class from litellm.repositories.table_repositories import ManagedVectorStoresRepository from litellm.types.utils import ModelResponse @@ -154,6 +155,53 @@ async def _authorize_nested_vector_store_ids( ) +def _ingest_provider_error(vector_store_config: Mapping[str, object]) -> str | None: + provider: Final = vector_store_config.get("custom_llm_provider", "openai") + if not isinstance(provider, str): + return "custom_llm_provider must be a string" + try: + get_ingestion_class(provider) + except ValueError as error: + return str(error) + return None + + +_MANAGED_STORE_CALLER_OPTIONS: Final = frozenset( + { + "vector_store_id", + "data_source_id", + "wait_for_ingestion", + "ingestion_timeout", + "custom_metadata", + "file_description", + "max_embedding_requests_per_min", + } +) + + +def _caller_vector_store_options( + request_vector_store_config: Mapping[str, object], + managed_store: LiteLLM_ManagedVectorStore | None, +) -> Mapping[str, object]: + if managed_store is None: + return request_vector_store_config + return MappingProxyType( + {key: value for key, value in request_vector_store_config.items() if key in _MANAGED_STORE_CALLER_OPTIONS} + ) + + +def _managed_store_overrides(managed_store: LiteLLM_ManagedVectorStore | None) -> Mapping[str, object]: + if managed_store is None: + return MappingProxyType({}) + return MappingProxyType( + { + key: value + for key, value in build_request_data_from_managed_vector_store(managed_store).items() + if value is not None + } + ) + + def _build_file_metadata_entry( response: object, file_data: tuple[str, bytes, str] | None = None, @@ -213,6 +261,8 @@ async def _save_vector_store_to_db_from_rag_ingest( user_api_key_dict: UserAPIKeyAuth, file_data: tuple[str, bytes, str] | None = None, file_url: str | None = None, + *, + store_is_managed: bool = False, ) -> None: """ Helper function to save a newly created vector store from RAG ingest to the database. @@ -220,7 +270,7 @@ async def _save_vector_store_to_db_from_rag_ingest( This function: - Extracts vector store ID and config from the ingest response - Checks if the vector store already exists in the database - - Creates a new database entry if it doesn't exist + - Creates a new database entry if it doesn't exist and the store is not registry-managed - Adds the vector store to the registry - Tracks team_id and user_id for access control @@ -229,6 +279,8 @@ async def _save_vector_store_to_db_from_rag_ingest( ingest_options: The ingest options containing vector store config prisma_client: The Prisma database client user_api_key_dict: User API key authentication info + store_is_managed: True when the requested id resolved to a managed store, so a missing row means + the store is config-registered and must not get a database row """ from litellm.proxy.vector_store_endpoints.management_endpoints import ( create_vector_store_in_db, @@ -277,6 +329,10 @@ async def _save_vector_store_to_db_from_rag_ingest( where={"vector_store_id": vector_store_id} ) + if existing_vector_store is None and store_is_managed: + verbose_proxy_logger.info("Vector store %s is config-registered, skipping database save", vector_store_id) + return + # Only create if it doesn't exist if existing_vector_store is None: verbose_proxy_logger.info("Saving newly created vector store %s to database", vector_store_id) @@ -545,14 +601,15 @@ async def rag_ingest( }, ) - await _authorize_nested_vector_store_ids( + resolved_stores: Final = await _authorize_nested_vector_store_ids( payload=ingest_options, user_api_key_dict=user_api_key_dict, ) + request_vector_store_config: Final = ingest_options.get("vector_store", {}) try: is_request_body_safe( - request_body=ingest_options.get("vector_store", {}), + request_body=request_vector_store_config, general_settings=general_settings, llm_router=llm_router, model="", @@ -560,6 +617,23 @@ async def rag_ingest( except ValueError as e: raise HTTPException(status_code=400, detail={"error": str(e)}) + managed_store: Final = resolved_stores.get(request_vector_store_config.get("vector_store_id")) + merged_vector_store_config: Final = { # mutable-ok: ingestion classes mutate it when loading credentials + **_caller_vector_store_options(request_vector_store_config, managed_store), + **_managed_store_overrides(managed_store), + } + merged_ingest_options: Final = { # mutable-ok: litellm.aingest takes a plain dict payload + **ingest_options, + "vector_store": merged_vector_store_config, + } + + provider_error: Final = _ingest_provider_error(merged_vector_store_config) + if provider_error is not None: + raise HTTPException( + status_code=400, + detail={"error": provider_error}, # mutable-ok: FastAPI serializes the detail as JSON + ) + # Add litellm data request_data: dict[str, Any] = {} request_data = await add_litellm_data_to_request( @@ -571,11 +645,15 @@ async def rag_ingest( proxy_config=proxy_config, ) - verbose_proxy_logger.debug("RAG Ingest - options: %s", ingest_options) + verbose_proxy_logger.debug( + "RAG Ingest - options: %s, custom_llm_provider: %s", + ingest_options, + merged_vector_store_config.get("custom_llm_provider", "openai"), + ) # Call ingest response: Final = await litellm.aingest( - ingest_options=ingest_options, + ingest_options=merged_ingest_options, file_data=file_data, file_url=file_url, file_id=file_id, @@ -599,6 +677,7 @@ async def rag_ingest( user_api_key_dict=user_api_key_dict, file_data=file_data, file_url=file_url, + store_is_managed=managed_store is not None, ) else: verbose_proxy_logger.warning( diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 2a9bda08325..e2aa5555eec 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -33,6 +33,10 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.llms.s3_vectors.vector_stores.transformation import ( + s3_vectors_ingest_embedding_options, + s3_vectors_ingest_target, +) from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion if TYPE_CHECKING: @@ -73,8 +77,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): 4. Store vectors with PutVectors API Configuration: - - vector_bucket_name: S3 vector bucket name (required) - - index_name: Vector index name (auto-creates if not provided) + - vector_store_id: "bucket_name:index_name" of an existing index, or an index name when vector_bucket_name is set + - vector_bucket_name: S3 vector bucket name (required unless vector_store_id carries it) + - index_name: Vector index name (auto-creates if neither it nor vector_store_id is provided) - dimension: Vector dimension (default: S3_VECTORS_DEFAULT_DIMENSION) - distance_metric: "cosine" or "euclidean" (default: S3_VECTORS_DEFAULT_DISTANCE_METRIC) - non_filterable_metadata_keys: List of metadata keys to exclude from filtering @@ -88,9 +93,8 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): BaseRAGIngestion.__init__(self, ingest_options=ingest_options, router=router) BaseAWSLLM.__init__(self) - # Extract config - self.vector_bucket_name: str = self.vector_store_config["vector_bucket_name"] - self.index_name: str | None = self.vector_store_config.get("index_name") + self.vector_bucket_name, self.index_name = s3_vectors_ingest_target(self.vector_store_config) + self.embedding_config = s3_vectors_ingest_embedding_options(self.vector_store_config, self.embedding_config) self.distance_metric: str = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC) self.non_filterable_metadata_keys: Sequence[str] = self.vector_store_config.get( "non_filterable_metadata_keys", diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index e313b749d06..781e92ea7d9 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -236,6 +236,22 @@ class TestS3VectorsVectorStoreConfig: assert executor.calls == [] + @pytest.mark.parametrize("vector_store_id", ["test-bucket:", ":test-index"]) + def test_transform_search_request_rejects_an_empty_bucket_or_index_in_the_id(self, vector_store_id): + config = S3VectorsVectorStoreConfig() + executor = _RecordingExecutor() + + with pytest.raises(ValueError, match="vector_store_id must be in format 'bucket_name:index_name'"): + config.transform_search_vector_store_request( + **_search_kwargs( + vector_store_id=vector_store_id, + litellm_params={"vector_bucket_name": "test-bucket"}, + embedding_executor=executor, + ) + ) + + assert executor.calls == [] + def test_transform_search_request_bucket_from_litellm_params(self): config = S3VectorsVectorStoreConfig() diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 1cceaf95b09..2d654ea28ec 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -6,6 +6,7 @@ Covers: """ import io +import json from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -239,6 +240,448 @@ class TestRagIngestSSRFBlocked: ) +S3_REGISTRY_STORE = { + "vector_store_id": "s3-store", + "custom_llm_provider": "s3_vectors", + "litellm_params": {"aws_region_name": "eu-west-1", "vector_bucket_name": "bkt", "index_name": "docs"}, +} +DB_MANAGED_STORE = { + "vector_store_id": "db-store", + "custom_llm_provider": "openai", + "litellm_credential_name": None, + "litellm_params": {"ttl_days": 7}, +} +AZURE_REGISTRY_STORE = { + "vector_store_id": "my-azure-index", + "custom_llm_provider": "azure_ai", + "litellm_params": { + "api_key": "azure-search-key", + "api_base": "https://search.example.net", + "api_version": "2024-07-01", + }, +} +BEDROCK_REGISTRY_STORE = { + "vector_store_id": "kb-store", + "custom_llm_provider": "bedrock", + "litellm_params": { + "aws_region_name": "eu-west-1", + "aws_access_key_id": "AKIA-registry", + "aws_secret_access_key": "registry-secret", + }, +} +CREDENTIALED_REGISTRY_STORE = { + "vector_store_id": "cred-store", + "custom_llm_provider": "openai", + "litellm_credential_name": "registry-openai", + "litellm_params": {}, +} +VERTEX_REGISTRY_STORE = { + "vector_store_id": "projects/registry-project/locations/us-central1/ragCorpora/42", + "custom_llm_provider": "vertex_ai", + "litellm_params": {"vertex_project": "registry-project", "vertex_location": "us-central1"}, +} +UNSUPPORTED_INGEST_PROVIDER_ERROR = ( + "Provider '{provider}' is not supported for RAG ingestion. " + "Supported providers: openai, bedrock, gemini, s3_vectors, vertex_ai" +) + + +def _registry_with(store): + registry = MagicMock() + registry.get_litellm_managed_vector_store_from_registry.return_value = store + return registry + + +def _ingest_form(vector_store): + return { + "files": {"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, + "data": {"request": json.dumps({"ingest_options": {"vector_store": vector_store}})}, + } + + +def _patched_ingest_boundary(registry_store, aingest_response): + return ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; tests assert the forwarded options + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new=AsyncMock(return_value=aingest_response), + ), + patch.object( # test-quality-ok: seeds the managed-store registry the merge under test reads + litellm, + "vector_store_registry", + _registry_with(registry_store), + ), + ) + + +def _patched_prisma_client(prisma_client): + return patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.prisma_client", + prisma_client, + ) + + +def test_rag_ingest_resolves_registry_store_provider_and_params(client_internal_user): + aingest_patch, registry_patch = _patched_ingest_boundary( + S3_REGISTRY_STORE, {"vector_store_id": "s3-store", "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(None), + ): + response = client_internal_user.post("/v1/rag/ingest", **_ingest_form({"vector_store_id": "s3-store"})) + + assert response.status_code == 200, response.json() + mock_aingest.assert_awaited_once() + forwarded = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["vector_store_id"] == "s3-store" + assert forwarded["custom_llm_provider"] == "s3_vectors" + assert forwarded["aws_region_name"] == "eu-west-1" + assert forwarded["vector_bucket_name"] == "bkt" + assert forwarded["index_name"] == "docs" + + +def test_rag_ingest_registry_store_wins_over_request_provider_and_params(client_internal_user): + aingest_patch, registry_patch = _patched_ingest_boundary( + S3_REGISTRY_STORE, {"vector_store_id": "s3-store", "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(None), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form( + {"vector_store_id": "s3-store", "custom_llm_provider": "openai", "aws_region_name": "us-east-1"} + ), + ) + + assert response.status_code == 200, response.json() + forwarded = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["custom_llm_provider"] == "s3_vectors" + assert forwarded["aws_region_name"] == "eu-west-1" + + +def test_rag_ingest_registry_store_drops_caller_destinations_and_keeps_upload_options(client_internal_user): + aingest_patch, registry_patch = _patched_ingest_boundary( + BEDROCK_REGISTRY_STORE, {"vector_store_id": "kb-store", "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(None), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form( + { + "vector_store_id": "kb-store", + "s3_bucket": "someone-elses-bucket", + "s3_prefix": "other-kb/", + "vector_bucket_name": "someone-elses-vectors", + "index_name": "other-index", + "vertex_project": "other-project", + "data_source_id": "DS2", + "wait_for_ingestion": True, + "ingestion_timeout": 60, + } + ), + ) + + assert response.status_code == 200, response.json() + forwarded = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded == { + "vector_store_id": "kb-store", + "custom_llm_provider": "bedrock", + "aws_region_name": "eu-west-1", + "aws_access_key_id": "AKIA-registry", + "aws_secret_access_key": "registry-secret", + "data_source_id": "DS2", + "wait_for_ingestion": True, + "ingestion_timeout": 60, + } + + +def test_rag_ingest_unmanaged_store_keeps_the_callers_full_config(client_internal_user): + caller_config = { + "vector_store_id": "KB-unmanaged", + "custom_llm_provider": "bedrock", + "s3_bucket": "callers-bucket", + "s3_prefix": "docs/", + } + aingest_patch, registry_patch = _patched_ingest_boundary( + None, {"vector_store_id": "KB-unmanaged", "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(None), + ): + response = client_internal_user.post("/v1/rag/ingest", **_ingest_form(caller_config)) + + assert response.status_code == 200, response.json() + assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == caller_config + + +def test_rag_ingest_db_managed_store_drops_the_callers_credential_name(client_internal_user): + aingest_patch, registry_patch = _patched_ingest_boundary( + DB_MANAGED_STORE, {"vector_store_id": "db-store", "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(None), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form({"vector_store_id": "db-store", "litellm_credential_name": "team-openai"}), + ) + + assert response.status_code == 200, response.json() + forwarded = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert "litellm_credential_name" not in forwarded + assert forwarded["custom_llm_provider"] == "openai" + assert forwarded["ttl_days"] == 7 + + +def test_rag_ingest_registry_store_credential_name_beats_the_callers(client_internal_user): + aingest_patch, registry_patch = _patched_ingest_boundary( + CREDENTIALED_REGISTRY_STORE, {"vector_store_id": "cred-store", "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(None), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form({"vector_store_id": "cred-store", "litellm_credential_name": "team-openai"}), + ) + + assert response.status_code == 200, response.json() + forwarded = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["litellm_credential_name"] == "registry-openai" + + +def test_rag_ingest_registry_store_keeps_the_callers_vertex_embedding_throttle(client_internal_user): + aingest_patch, registry_patch = _patched_ingest_boundary( + VERTEX_REGISTRY_STORE, {"vector_store_id": VERTEX_REGISTRY_STORE["vector_store_id"], "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(None), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form( + { + "vector_store_id": VERTEX_REGISTRY_STORE["vector_store_id"], + "max_embedding_requests_per_min": 500, + "vector_db_config": {"pinecone": {"index_name": "attacker-index"}}, + } + ), + ) + + assert response.status_code == 200, response.json() + assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == { + "vector_store_id": VERTEX_REGISTRY_STORE["vector_store_id"], + "custom_llm_provider": "vertex_ai", + "vertex_project": "registry-project", + "vertex_location": "us-central1", + "max_embedding_requests_per_min": 500, + } + + +def test_rag_ingest_rejects_registry_store_provider_without_ingestion_support(client_internal_user): + aingest_patch, registry_patch = _patched_ingest_boundary( + AZURE_REGISTRY_STORE, {"vector_store_id": "my-azure-index", "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(None), + ): + response = client_internal_user.post("/v1/rag/ingest", **_ingest_form({"vector_store_id": "my-azure-index"})) + + assert response.status_code == 400, response.json() + assert response.json()["detail"]["error"] == UNSUPPORTED_INGEST_PROVIDER_ERROR.format(provider="azure_ai") + mock_aingest.assert_not_awaited() + + +def test_rag_ingest_rejects_request_provider_without_ingestion_support(client_internal_user): + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts it is never reached + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new=AsyncMock(return_value={"vector_store_id": "vs_new", "file_id": "file-test"}), + ) as mock_aingest, + patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam + ): + response = client_internal_user.post( + "/v1/rag/ingest", + json={"file_id": "file-test", "ingest_options": {"vector_store": {"custom_llm_provider": "milvus"}}}, + ) + + assert response.status_code == 400, response.json() + assert response.json() == {"detail": {"error": UNSUPPORTED_INGEST_PROVIDER_ERROR.format(provider="milvus")}} + mock_aingest.assert_not_awaited() + + +def test_rag_ingest_rejects_non_string_provider(client_internal_user): + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts it is never reached + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new=AsyncMock(return_value={"vector_store_id": "vs_new", "file_id": "file-test"}), + ) as mock_aingest, + patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam + ): + response = client_internal_user.post( + "/v1/rag/ingest", + json={ + "file_id": "file-test", + "ingest_options": {"vector_store": {"custom_llm_provider": {"provider": "milvus"}}}, + }, + ) + + assert response.status_code == 400, response.json() + assert response.json() == {"detail": {"error": "custom_llm_provider must be a string"}} + mock_aingest.assert_not_awaited() + + +def test_rag_ingest_never_creates_db_row_for_registry_store(client_internal_user): + prisma_client = MagicMock() + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) + create_in_db = AsyncMock() + aingest_patch, registry_patch = _patched_ingest_boundary( + S3_REGISTRY_STORE, {"vector_store_id": "s3-store", "file_id": "file_123"} + ) + with ( + aingest_patch, + registry_patch, + _patched_prisma_client(prisma_client), + patch( # test-quality-ok: the DB write boundary the guard under test must never reach + "litellm.proxy.vector_store_endpoints.management_endpoints.create_vector_store_in_db", + new=create_in_db, + ), + ): + response = client_internal_user.post("/v1/rag/ingest", **_ingest_form({"vector_store_id": "s3-store"})) + + assert response.status_code == 200, response.json() + prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_awaited_once() + create_in_db.assert_not_awaited() + prisma_client.db.litellm_managedvectorstorestable.update.assert_not_called() + + +def test_rag_ingest_fresh_store_creates_db_row_with_the_requesters_params(client_internal_user): + prisma_client = MagicMock() + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) + create_in_db = AsyncMock() + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; persistence is what the test asserts + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new=AsyncMock(return_value={"vector_store_id": "vs_new", "file_id": "file_123"}), + ), + patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam + _patched_prisma_client(prisma_client), + patch( # test-quality-ok: the DB write boundary whose inputs the test asserts + "litellm.proxy.vector_store_endpoints.management_endpoints.create_vector_store_in_db", + new=create_in_db, + ), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form({"custom_llm_provider": "bedrock", "aws_region_name": "us-east-1"}), + ) + + assert response.status_code == 200, response.json() + create_in_db.assert_awaited_once() + created = create_in_db.await_args.kwargs + assert created["vector_store_id"] == "vs_new" + assert created["custom_llm_provider"] == "bedrock" + assert created["litellm_params"] == {"aws_region_name": "us-east-1"} + + +def test_rag_ingest_hands_persistence_the_requesters_options_not_registry_credentials(client_internal_user): + save_helper = AsyncMock() + aingest_patch, registry_patch = _patched_ingest_boundary( + BEDROCK_REGISTRY_STORE, {"vector_store_id": "kb-store", "file_id": "file_123"} + ) + with ( + aingest_patch as mock_aingest, + registry_patch, + _patched_prisma_client(MagicMock()), + patch( # test-quality-ok: the persistence seam whose inputs the test asserts + "litellm.proxy.rag_endpoints.endpoints._save_vector_store_to_db_from_rag_ingest", + new=save_helper, + ), + ): + response = client_internal_user.post("/v1/rag/ingest", **_ingest_form({"vector_store_id": "kb-store"})) + + assert response.status_code == 200, response.json() + forwarded = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["aws_secret_access_key"] == "registry-secret" + save_helper.assert_awaited_once() + assert save_helper.await_args.kwargs["ingest_options"]["vector_store"] == {"vector_store_id": "kb-store"} + assert save_helper.await_args.kwargs["store_is_managed"] is True + + +async def test_save_vector_store_from_rag_ingest_appends_file_to_db_managed_store(): + from litellm.proxy.rag_endpoints.endpoints import _save_vector_store_to_db_from_rag_ingest + + existing_row = MagicMock() + existing_row.vector_store_metadata = {"ingested_files": [{"file_id": "file_old"}]} + prisma_client = MagicMock() + table = prisma_client.db.litellm_managedvectorstorestable + table.find_unique = AsyncMock(return_value=existing_row) + table.update = AsyncMock() + create_in_db = AsyncMock() + + with patch( # test-quality-ok: the DB write boundary the append branch must not reach + "litellm.proxy.vector_store_endpoints.management_endpoints.create_vector_store_in_db", + new=create_in_db, + ): + await _save_vector_store_to_db_from_rag_ingest( + response={"vector_store_id": "vs_db_managed", "file_id": "file_new"}, + ingest_options={"vector_store": {"vector_store_id": "vs_db_managed"}}, + prisma_client=prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1", team_id="team-1"), + store_is_managed=True, + ) + + create_in_db.assert_not_awaited() + table.update.assert_awaited_once() + stored_metadata = json.loads(table.update.await_args.kwargs["data"]["vector_store_metadata"]) + assert [entry["file_id"] for entry in stored_metadata["ingested_files"]] == ["file_old", "file_new"] + + +async def test_save_vector_store_from_rag_ingest_still_creates_row_for_fresh_store(): + from litellm.proxy.rag_endpoints.endpoints import _save_vector_store_to_db_from_rag_ingest + + prisma_client = MagicMock() + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) + create_in_db = AsyncMock() + + with patch( # test-quality-ok: the DB write boundary whose inputs the test asserts + "litellm.proxy.vector_store_endpoints.management_endpoints.create_vector_store_in_db", + new=create_in_db, + ): + await _save_vector_store_to_db_from_rag_ingest( + response={"vector_store_id": "vs_new", "file_id": "file_new"}, + ingest_options={"vector_store": {"custom_llm_provider": "bedrock", "aws_region_name": "us-east-1"}}, + prisma_client=prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1", team_id="team-1"), + store_is_managed=False, + ) + + create_in_db.assert_awaited_once() + created = create_in_db.await_args.kwargs + assert created["vector_store_id"] == "vs_new" + assert created["custom_llm_provider"] == "bedrock" + assert created["litellm_params"] == {"aws_region_name": "us-east-1"} + assert created["team_id"] == "team-1" + + def test_rag_query_returns_response_cost_header(client_internal_user): """ /v1/rag/query must surface the completion cost via the diff --git a/tests/test_litellm/rag/ingestion/__init__.py b/tests/test_litellm/rag/ingestion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py new file mode 100644 index 00000000000..07fd2b765f3 --- /dev/null +++ b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py @@ -0,0 +1,110 @@ +from types import SimpleNamespace + +import pytest + +from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion + +STORE_ID_FORMAT_ERROR = "vector_store_id must be in format 'bucket_name:index_name'" +REQUEST_EMBEDDING_MODEL = "text-embedding-3-small" +STORE_EMBEDDING_MODEL = "text-embedding-3-large" +REQUEST_EMBEDDING = {"model": REQUEST_EMBEDDING_MODEL} + + +class _RecordingRouter: + def __init__(self): + self.embedding_models = [] + + async def aembedding(self, model, input): + self.embedding_models.append(model) + return SimpleNamespace(data=[{"embedding": [0.1, 0.2]} for _ in input]) + + +def _ingestion(embedding=REQUEST_EMBEDDING, router=None, **vector_store): + vector_store_options = {"custom_llm_provider": "s3_vectors", "aws_region_name": "us-west-2", **vector_store} + ingest_options = {"vector_store": vector_store_options} if embedding is None else { + "embedding": embedding, + "vector_store": vector_store_options, + } + return S3VectorsRAGIngestion(ingest_options=ingest_options, router=router) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("store_model_key", ["embedding_model", "litellm_embedding_model"]) +async def test_a_registered_store_embedding_model_wins_over_the_request_on_ingest(store_model_key): + router = _RecordingRouter() + ingestion = _ingestion( + router=router, vector_store_id="my-embeddings:my-index", **{store_model_key: STORE_EMBEDDING_MODEL} + ) + + await ingestion.embed(["chunk one", "chunk two"]) + + assert router.embedding_models == [STORE_EMBEDDING_MODEL] + + +@pytest.mark.asyncio +async def test_a_registered_store_embedding_model_is_used_when_the_request_names_none(): + router = _RecordingRouter() + ingestion = _ingestion( + embedding=None, router=router, vector_store_id="my-embeddings:my-index", embedding_model=STORE_EMBEDDING_MODEL + ) + + await ingestion.embed(["chunk"]) + + assert router.embedding_models == [STORE_EMBEDDING_MODEL] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("store_model", [{}, {"embedding_model": ""}]) +async def test_the_request_embedding_model_is_kept_when_the_store_names_none(store_model): + router = _RecordingRouter() + ingestion = _ingestion(router=router, vector_store_id="my-embeddings:my-index", **store_model) + + await ingestion.embed(["chunk"]) + + assert router.embedding_models == [REQUEST_EMBEDDING_MODEL] + + +def test_store_id_alone_names_the_bucket_and_index(): + ingestion = _ingestion(vector_store_id="my-embeddings:my-index") + + assert (ingestion.vector_bucket_name, ingestion.index_name) == ("my-embeddings", "my-index") + + +def test_store_id_without_a_colon_is_the_index_inside_the_given_bucket(): + ingestion = _ingestion(vector_store_id="my-index", vector_bucket_name="my-embeddings") + + assert (ingestion.vector_bucket_name, ingestion.index_name) == ("my-embeddings", "my-index") + + +def test_explicit_bucket_and_index_win_over_the_store_id(): + ingestion = _ingestion(vector_store_id="id-bucket:id-index", vector_bucket_name="my-bucket", index_name="docs") + + assert (ingestion.vector_bucket_name, ingestion.index_name) == ("my-bucket", "docs") + + +def test_bucket_alone_leaves_the_index_to_be_generated(): + ingestion = _ingestion(vector_bucket_name="my-embeddings") + + assert (ingestion.vector_bucket_name, ingestion.index_name) == ("my-embeddings", None) + + +@pytest.mark.parametrize( + "vector_store", + [{}, {"vector_store_id": "my-index"}, {"vector_store_id": "my-index", "vector_bucket_name": ""}], +) +def test_no_bucket_anywhere_is_rejected(vector_store): + with pytest.raises(ValueError, match=STORE_ID_FORMAT_ERROR): + _ingestion(**vector_store) + + +@pytest.mark.parametrize( + "vector_store", + [ + {"vector_store_id": "my-embeddings:"}, + {"vector_store_id": ":my-index"}, + {"vector_store_id": "my-embeddings:", "vector_bucket_name": "my-embeddings"}, + ], +) +def test_an_empty_bucket_or_index_in_the_store_id_is_rejected_instead_of_generating_an_index(vector_store): + with pytest.raises(ValueError, match=STORE_ID_FORMAT_ERROR): + _ingestion(**vector_store)