fix(rag): keep only per-upload caller options when ingesting into a registered store

This commit is contained in:
mateo-berri 2026-09-19 02:22:34 -07:00
parent b60b513f6a
commit a98c48f933
2 changed files with 92 additions and 1 deletions

View file

@ -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

View file

@ -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