mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
fix(rag): keep only per-upload caller options when ingesting into a registered store
This commit is contained in:
parent
b60b513f6a
commit
a98c48f933
2 changed files with 92 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue