From b60b513f6a4634803f5fc42bc9905425480c9f11 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 01:37:50 -0700 Subject: [PATCH 1/7] fix(rag): resolve registry stores on /v1/rag/ingest and reject providers without ingestion POST /v1/rag/ingest authorized the managed vector store the request named but then handed the raw request options to the ingestion pipeline, which defaults to OpenAI. A request naming only a registered store id uploaded the document to OpenAI Files, got an OpenAI 400, and answered HTTP 200 with status "failed"; naming azure_ai explicitly escaped as a 500. The store's provider and litellm_params now merge into the request the way /v1/rag/query already does (store wins, None values dropped), the merged provider is checked against the ingestion registry before any upload so unsupported providers get a 400 naming the supported ones, and persistence keeps reading the caller's original options so registry credentials never reach the database. A registry store with no database row is no longer written as a new row. --- litellm/proxy/rag_endpoints/endpoints.py | 65 +++- .../proxy/rag_endpoints/test_rag_endpoints.py | 350 ++++++++++++++++++ 2 files changed, 410 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index c09f9c755ed..3ece5399232 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,29 @@ 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 + + +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 +237,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 +246,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 +255,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 +305,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 +577,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 +593,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 + **request_vector_store_config, + **_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 +621,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 +653,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/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 1cceaf95b09..55d46f621c7 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,355 @@ 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", + }, +} +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): + """ + Regression for LIT-7956: naming only a registry store id must ingest into + that store's provider with its litellm_params, the way /v1/rag/query and + /v1/vector_stores/{id}/search resolve it. Pre-fix the resolved store was + thrown away and the pipeline defaulted to OpenAI Files. + """ + 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): + """A caller cannot steer a registry store to another provider or region by repeating the keys in the request.""" + 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_db_managed_store_keeps_the_callers_credential_name(client_internal_user): + """ + A store synced from the database carries litellm_credential_name=None; that + null is the absence of a store-side value, not an override, so the credential + the caller named must survive the merge exactly as it did before the fix. + """ + 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 forwarded["litellm_credential_name"] == "team-openai" + assert forwarded["custom_llm_provider"] == "openai" + assert forwarded["ttl_days"] == 7 + + +def test_rag_ingest_rejects_registry_store_provider_without_ingestion_support(client_internal_user): + """ + Regression for LIT-7956: a registry store on a provider with no ingestion + implementation must be rejected with 400 before anything is uploaded. + Pre-fix the document went to OpenAI Files and the proxy answered 200 with + status "failed". + """ + 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): + """A request-supplied provider outside the ingestion registry is a 400, never a 500 from inside the pipeline.""" + 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): + """ + A config-registered store has no DB row; ingesting into it must not create + one, since that row would outlive the config and carry request-side params. + """ + 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): + """A request naming no store id creates a brand new one, whose row must still be written as before the fix.""" + 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): + """ + Persistence only ever sees what the requester sent: the merged options carry + the registry's credentials, which must never be written back as litellm_params. + """ + 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 From a98c48f9336fe84703373fcf6cf5436245fa51d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 02:22:34 -0700 Subject: [PATCH 2/7] fix(rag): keep only per-upload caller options when ingesting into a registered store --- litellm/proxy/rag_endpoints/endpoints.py | 26 ++++++- .../proxy/rag_endpoints/test_rag_endpoints.py | 67 +++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 3ece5399232..cd7657b3536 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -166,6 +166,30 @@ def _ingest_provider_error(vector_store_config: Mapping[str, object]) -> str | N return None +_MANAGED_STORE_CALLER_OPTIONS: Final = frozenset( + { + "vector_store_id", + "litellm_credential_name", + "data_source_id", + "wait_for_ingestion", + "ingestion_timeout", + "custom_metadata", + "file_description", + } +) + + +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({}) @@ -595,7 +619,7 @@ async def rag_ingest( 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 - **request_vector_store_config, + **_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 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 55d46f621c7..b2b6496f542 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -359,6 +359,73 @@ def test_rag_ingest_registry_store_wins_over_request_provider_and_params(client_ assert forwarded["aws_region_name"] == "eu-west-1" +def test_rag_ingest_registry_store_drops_caller_destinations_and_keeps_upload_options(client_internal_user): + """ + The store's registered credentials ride along on the upload, so a caller authorized + on the store must not be able to point them at a bucket, index or project the store + does not define. Per-upload options still pass through. + """ + 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): + """A store id the proxy does not manage carries no server-side config, so the caller's config is all there is.""" + 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_keeps_the_callers_credential_name(client_internal_user): """ A store synced from the database carries litellm_credential_name=None; that From e2d118aaf8e570a30288aa625913539fe2230aea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:38:20 -0700 Subject: [PATCH 3/7] fix(rag): read a registered S3 Vectors store's bucket and index from its id A registered S3 Vectors store usually carries only its "bucket:index" id, and the previous commit stopped forwarding the caller's bucket and index for a managed store, so ingesting into one raised KeyError 'vector_bucket_name'. The ingestion now derives both from vector_store_id with the rule the search side already uses, explicit keys still winning. The caller's litellm_credential_name is dropped for a managed store too, since it expands into api_key and api_base, and max_embedding_requests_per_min joins the per-upload options a caller may still set. --- .../vector_stores/transformation.py | 24 ++++--- litellm/proxy/rag_endpoints/endpoints.py | 2 +- litellm/rag/ingestion/s3_vectors_ingestion.py | 29 ++++++-- .../proxy/rag_endpoints/test_rag_endpoints.py | 71 +++++++++++++++++-- tests/test_litellm/rag/ingestion/__init__.py | 0 .../ingestion/test_s3_vectors_ingestion.py | 52 ++++++++++++++ 6 files changed, 157 insertions(+), 21 deletions(-) create mode 100644 tests/test_litellm/rag/ingestion/__init__.py create mode 100644 tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index a9902a0d27c..04f561aa2ca 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -26,6 +26,19 @@ 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]: + if ":" in vector_store_id: + bucket_name, index_name = vector_store_id.split(":", 1) + return bucket_name, index_name + if not isinstance(fallback_bucket_name, str) or not fallback_bucket_name: + raise ValueError(S3_VECTORS_STORE_ID_ERROR) + return fallback_bucket_name, vector_store_id class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM): @@ -74,16 +87,7 @@ class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM @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 cd7657b3536..4f0c9f42421 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -169,12 +169,12 @@ def _ingest_provider_error(vector_store_config: Mapping[str, object]) -> str | N _MANAGED_STORE_CALLER_OPTIONS: Final = frozenset( { "vector_store_id", - "litellm_credential_name", "data_source_id", "wait_for_ingestion", "ingestion_timeout", "custom_metadata", "file_description", + "max_embedding_requests_per_min", } ) diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 2a9bda08325..15f0a89cf95 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_STORE_ID_ERROR, + split_s3_vectors_store_id, +) from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion if TYPE_CHECKING: @@ -62,6 +66,22 @@ class S3VectorsQueryResponse(TypedDict, total=False): vectors: Sequence[S3VectorsQueryMatch] +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 + + class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): """ S3 Vectors RAG ingestion using httpx + AWS SigV4 signing. @@ -73,8 +93,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 +109,7 @@ 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.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/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index b2b6496f542..4b8efa14c2b 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -269,6 +269,17 @@ BEDROCK_REGISTRY_STORE = { "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" @@ -426,11 +437,12 @@ def test_rag_ingest_unmanaged_store_keeps_the_callers_full_config(client_interna assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == caller_config -def test_rag_ingest_db_managed_store_keeps_the_callers_credential_name(client_internal_user): +def test_rag_ingest_db_managed_store_drops_the_callers_credential_name(client_internal_user): """ - A store synced from the database carries litellm_credential_name=None; that - null is the absence of a store-side value, not an override, so the credential - the caller named must survive the merge exactly as it did before the fix. + litellm_credential_name expands into api_key and api_base at ingest time, so a + caller naming one would point a managed store's upload at a different endpoint. + A store synced from the database carries litellm_credential_name=None, and that + null must not resurrect the caller's choice either. """ aingest_patch, registry_patch = _patched_ingest_boundary( DB_MANAGED_STORE, {"vector_store_id": "db-store", "file_id": "file_123"} @@ -447,11 +459,60 @@ def test_rag_ingest_db_managed_store_keeps_the_callers_credential_name(client_in assert response.status_code == 200, response.json() forwarded = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] - assert forwarded["litellm_credential_name"] == "team-openai" + 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): """ Regression for LIT-7956: a registry store on a provider with no ingestion 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..30f9adf07b3 --- /dev/null +++ b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py @@ -0,0 +1,52 @@ +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'" + + +def _ingestion(**vector_store): + return S3VectorsRAGIngestion( + ingest_options={ + "embedding": {"model": "text-embedding-3-small"}, + "vector_store": {"custom_llm_provider": "s3_vectors", "aws_region_name": "us-west-2", **vector_store}, + } + ) + + +def test_store_id_alone_names_the_bucket_and_index(): + """ + Regression for LIT-7956: a registered S3 Vectors store carries only its + "bucket:index" id, and the proxy no longer forwards the caller's bucket and + index for a managed store, so the ingestion must read both from the id. + """ + 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) From e4d01d1d781ac1efe41f5356a14f1adc1ee250a6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:52:48 -0700 Subject: [PATCH 4/7] fix(s3_vectors): reject a store id with an empty bucket or index part A "bucket:" or ":index" id split into an empty name, so ingestion silently generated a fresh index and search sent the empty name to AWS. Both sides now raise the existing format error through the shared helper. --- .../s3_vectors/vector_stores/transformation.py | 10 +++++----- .../test_s3_vectors_transformation.py | 16 ++++++++++++++++ .../rag/ingestion/test_s3_vectors_ingestion.py | 13 +++++++++++++ 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 04f561aa2ca..b02734e316d 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -33,12 +33,12 @@ S3_VECTORS_STORE_ID_ERROR: Final = ( def split_s3_vectors_store_id(vector_store_id: str, fallback_bucket_name: object) -> tuple[str, str]: - if ":" in vector_store_id: - bucket_name, index_name = vector_store_id.split(":", 1) - return bucket_name, index_name - if not isinstance(fallback_bucket_name, str) or not fallback_bucket_name: + 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 fallback_bucket_name, vector_store_id + return bucket_name, index_name class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM): 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/rag/ingestion/test_s3_vectors_ingestion.py b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py index 30f9adf07b3..3256de48ef9 100644 --- a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py +++ b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py @@ -50,3 +50,16 @@ def test_bucket_alone_leaves_the_index_to_be_generated(): 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) From ccb48eb52843e6f56683d36dc01a9fc67809e60a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 04:04:59 -0700 Subject: [PATCH 5/7] refactor(s3_vectors): keep the ingest target derivation under llms/s3_vectors The ingest-side bucket and index precedence now sits next to the shared store id split instead of under litellm/rag/, where provider-specific parsing does not belong. --- .../vector_stores/transformation.py | 16 ++++++++++++++ litellm/rag/ingestion/s3_vectors_ingestion.py | 21 +------------------ 2 files changed, 17 insertions(+), 20 deletions(-) diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index b02734e316d..e074d1ebce2 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -41,6 +41,22 @@ def split_s3_vectors_store_id(vector_store_id: str, fallback_bucket_name: object 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 + + class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS S3 Vectors.""" diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 15f0a89cf95..8f362c146c3 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -33,10 +33,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.llms.s3_vectors.vector_stores.transformation import ( - S3_VECTORS_STORE_ID_ERROR, - split_s3_vectors_store_id, -) +from litellm.llms.s3_vectors.vector_stores.transformation import s3_vectors_ingest_target from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion if TYPE_CHECKING: @@ -66,22 +63,6 @@ class S3VectorsQueryResponse(TypedDict, total=False): vectors: Sequence[S3VectorsQueryMatch] -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 - - class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): """ S3 Vectors RAG ingestion using httpx + AWS SigV4 signing. From e74a5e0c21cdfd4ad9590e963c7a216517f4e1c8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 04:10:49 -0700 Subject: [PATCH 6/7] test(rag): drop the docstrings from the registered-store ingest tests --- .../proxy/rag_endpoints/test_rag_endpoints.py | 35 ------------------- .../ingestion/test_s3_vectors_ingestion.py | 5 --- 2 files changed, 40 deletions(-) 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 4b8efa14c2b..2d654ea28ec 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -321,12 +321,6 @@ def _patched_prisma_client(prisma_client): def test_rag_ingest_resolves_registry_store_provider_and_params(client_internal_user): - """ - Regression for LIT-7956: naming only a registry store id must ingest into - that store's provider with its litellm_params, the way /v1/rag/query and - /v1/vector_stores/{id}/search resolve it. Pre-fix the resolved store was - thrown away and the pipeline defaulted to OpenAI Files. - """ aingest_patch, registry_patch = _patched_ingest_boundary( S3_REGISTRY_STORE, {"vector_store_id": "s3-store", "file_id": "file_123"} ) @@ -348,7 +342,6 @@ def test_rag_ingest_resolves_registry_store_provider_and_params(client_internal_ def test_rag_ingest_registry_store_wins_over_request_provider_and_params(client_internal_user): - """A caller cannot steer a registry store to another provider or region by repeating the keys in the request.""" aingest_patch, registry_patch = _patched_ingest_boundary( S3_REGISTRY_STORE, {"vector_store_id": "s3-store", "file_id": "file_123"} ) @@ -371,11 +364,6 @@ def test_rag_ingest_registry_store_wins_over_request_provider_and_params(client_ def test_rag_ingest_registry_store_drops_caller_destinations_and_keeps_upload_options(client_internal_user): - """ - The store's registered credentials ride along on the upload, so a caller authorized - on the store must not be able to point them at a bucket, index or project the store - does not define. Per-upload options still pass through. - """ aingest_patch, registry_patch = _patched_ingest_boundary( BEDROCK_REGISTRY_STORE, {"vector_store_id": "kb-store", "file_id": "file_123"} ) @@ -416,7 +404,6 @@ def test_rag_ingest_registry_store_drops_caller_destinations_and_keeps_upload_op def test_rag_ingest_unmanaged_store_keeps_the_callers_full_config(client_internal_user): - """A store id the proxy does not manage carries no server-side config, so the caller's config is all there is.""" caller_config = { "vector_store_id": "KB-unmanaged", "custom_llm_provider": "bedrock", @@ -438,12 +425,6 @@ def test_rag_ingest_unmanaged_store_keeps_the_callers_full_config(client_interna def test_rag_ingest_db_managed_store_drops_the_callers_credential_name(client_internal_user): - """ - litellm_credential_name expands into api_key and api_base at ingest time, so a - caller naming one would point a managed store's upload at a different endpoint. - A store synced from the database carries litellm_credential_name=None, and that - null must not resurrect the caller's choice either. - """ aingest_patch, registry_patch = _patched_ingest_boundary( DB_MANAGED_STORE, {"vector_store_id": "db-store", "file_id": "file_123"} ) @@ -514,12 +495,6 @@ def test_rag_ingest_registry_store_keeps_the_callers_vertex_embedding_throttle(c def test_rag_ingest_rejects_registry_store_provider_without_ingestion_support(client_internal_user): - """ - Regression for LIT-7956: a registry store on a provider with no ingestion - implementation must be rejected with 400 before anything is uploaded. - Pre-fix the document went to OpenAI Files and the proxy answered 200 with - status "failed". - """ aingest_patch, registry_patch = _patched_ingest_boundary( AZURE_REGISTRY_STORE, {"vector_store_id": "my-azure-index", "file_id": "file_123"} ) @@ -536,7 +511,6 @@ def test_rag_ingest_rejects_registry_store_provider_without_ingestion_support(cl def test_rag_ingest_rejects_request_provider_without_ingestion_support(client_internal_user): - """A request-supplied provider outside the ingestion registry is a 400, never a 500 from inside the pipeline.""" 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", @@ -576,10 +550,6 @@ def test_rag_ingest_rejects_non_string_provider(client_internal_user): def test_rag_ingest_never_creates_db_row_for_registry_store(client_internal_user): - """ - A config-registered store has no DB row; ingesting into it must not create - one, since that row would outlive the config and carry request-side params. - """ prisma_client = MagicMock() prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) create_in_db = AsyncMock() @@ -604,7 +574,6 @@ def test_rag_ingest_never_creates_db_row_for_registry_store(client_internal_user def test_rag_ingest_fresh_store_creates_db_row_with_the_requesters_params(client_internal_user): - """A request naming no store id creates a brand new one, whose row must still be written as before the fix.""" prisma_client = MagicMock() prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) create_in_db = AsyncMock() @@ -634,10 +603,6 @@ def test_rag_ingest_fresh_store_creates_db_row_with_the_requesters_params(client def test_rag_ingest_hands_persistence_the_requesters_options_not_registry_credentials(client_internal_user): - """ - Persistence only ever sees what the requester sent: the merged options carry - the registry's credentials, which must never be written back as litellm_params. - """ save_helper = AsyncMock() aingest_patch, registry_patch = _patched_ingest_boundary( BEDROCK_REGISTRY_STORE, {"vector_store_id": "kb-store", "file_id": "file_123"} diff --git a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py index 3256de48ef9..24dfc392bbe 100644 --- a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py +++ b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py @@ -15,11 +15,6 @@ def _ingestion(**vector_store): def test_store_id_alone_names_the_bucket_and_index(): - """ - Regression for LIT-7956: a registered S3 Vectors store carries only its - "bucket:index" id, and the proxy no longer forwards the caller's bucket and - index for a managed store, so the ingestion must read both from the id. - """ ingestion = _ingestion(vector_store_id="my-embeddings:my-index") assert (ingestion.vector_bucket_name, ingestion.index_name) == ("my-embeddings", "my-index") From aef209963a388dfe0448ce7404341cc9c3019b69 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 04:31:15 -0700 Subject: [PATCH 7/7] fix(s3_vectors): embed registered-store ingests with the store's embedding model The S3 Vectors ingestion embedded every chunk with the request's embedding.model or the default, never the embedding_model the store was registered with, while search on the same store embeds with the registered model. A registered store uploaded to by id alone therefore embedded with the wrong model and AWS rejected the vectors on the dimension mismatch. The store's embedding model now wins for S3 Vectors ingestion through a helper next to the one search already uses --- .../vector_stores/transformation.py | 19 +++++- litellm/rag/ingestion/s3_vectors_ingestion.py | 6 +- .../ingestion/test_s3_vectors_ingestion.py | 62 +++++++++++++++++-- 3 files changed, 78 insertions(+), 9 deletions(-) diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index e074d1ebce2..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, @@ -57,6 +58,21 @@ def s3_vectors_ingest_target(vector_store_config: Mapping[str, object]) -> tuple 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): """Vector store configuration for AWS S3 Vectors.""" @@ -98,8 +114,7 @@ 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]: diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 8f362c146c3..e2aa5555eec 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -33,7 +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_target +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: @@ -91,6 +94,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): BaseAWSLLM.__init__(self) 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/rag/ingestion/test_s3_vectors_ingestion.py b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py index 24dfc392bbe..07fd2b765f3 100644 --- a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py +++ b/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py @@ -1,18 +1,68 @@ +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} -def _ingestion(**vector_store): - return S3VectorsRAGIngestion( - ingest_options={ - "embedding": {"model": "text-embedding-3-small"}, - "vector_store": {"custom_llm_provider": "s3_vectors", "aws_region_name": "us-west-2", **vector_store}, - } +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")