feat(rag): add Milvus vector store ingestion support - #30388
Conversation
|
|
Congrats! CodSpeed is installed 🎉
You will start to see performance impacts in the reports once the benchmarks are run from your default branch.
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
a1ca56b to
569542b
Compare
|
|
||
| from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple | ||
|
|
||
| import litellm |
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 7 · PR risk: 0/10 |
569542b to
67059f3
Compare
67059f3 to
c27d0de
Compare
Greptile SummaryAdds Milvus as a supported RAG ingestion provider by registering
Confidence Score: 5/5Safe to merge; the new Milvus ingestion path is well-isolated, uses litellm's existing HTTP infrastructure, and the authorization refactor is backward-compatible for all existing providers. The security-sensitive changes (authorization normalization, credential-hydration protection, VIEW_ONLY role guard) are all covered by targeted tests, and the logic correctly short-circuits the managed-store check for non-auto-create providers so existing OpenAI/Bedrock users are unaffected. The only observation is a minor debugging ergonomics issue in No files require special attention;
|
| Filename | Overview |
|---|---|
| litellm/rag/ingestion/milvus_ingestion.py | New Milvus ingestion class: correctly uses litellm's httpx infrastructure, credential-protected fields prevent auth bypass, api_key env fallback is scoped to env-supplied api_base, defaults wired in __init__. Minor: _collection_exists swallows all exceptions. |
| litellm/proxy/rag_endpoints/endpoints.py | Authorization flow refactored: provider-specific write-target normalization is dispatched polymorphically via get_ingestion_class, VIEW_ONLY managed-store check is now conditional on provider auto-create capability (backward-compatible for OpenAI/Bedrock), addresses all previously flagged bypass vectors. |
| litellm/rag/ingestion/base_ingestion.py | Adds normalize_authorized_vector_store_id, credential_protected_fields, and can_auto_create_vector_store classmethods with safe defaults; credential hydration now skips protected fields, preventing post-authorization write-target override. |
| litellm/types/rag.py | Adds MilvusVectorStoreOptions TypedDict and includes it in the RAGIngestVectorStoreOptions union; shape matches the ingestion class config. |
| litellm/constants.py | Adds three Milvus RAG constants (vector field name, text field name, metric type) following the existing constant pattern for other providers. |
| litellm/rag/main.py | Registers MilvusRAGIngestion under the "milvus" key in INGESTION_REGISTRY, enabling get_ingestion_class("milvus") lookups used by the authorization dispatch. |
| tests/test_litellm/rag/test_milvus_ingestion.py | 16 unit tests covering config validation, auth header logic, auto-create/skip-create flows, error propagation, env-only db/partition, and credential hydration protection — all mocked, no real network calls. |
| tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py | Extended with tests covering collection_name normalization/authorization bypass prevention, VIEW_ONLY auto-create guard for Milvus, backward-compatibility for OpenAI/Bedrock provider-native IDs, and provider capability detection — all properly mocked. |
| tests/vector_store_tests/rag/test_rag_milvus.py | Env-gated integration test (skips unless MILVUS_API_BASE is set); follows BaseRAGTest pattern; uses httpx directly for verification queries, consistent with integration test conventions for this folder. |
Reviews (5): Last reviewed commit: "fix(rag): block credential hydration fro..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
@maycuatroi1 Please resolve all the greptile comments and get the score to 5/5. Also add working proof of this. Thanks! |
|
@Sameerlite thanks for your comment. Pushed Working proof - live
|
|
Thanks for the contribution! We're triggering a Greptile code review on this PR — we'll take a closer look once the results are in! |
|
@maycuatroi1 Can you please rebase the PR to litellm_internal_staging? Thanks |
Adds write/ingest support for self-hosted Milvus to complement the existing Milvus search provider. /rag/ingest now accepts custom_llm_provider=milvus. - MilvusRAGIngestion implements the store() step via the Milvus REST API v2 (entities/insert), reusing the base upload/ocr/chunk/embed pipeline - Auto-creates the collection via quick setup (dynamic fields) when missing - Embeddings generated through litellm embedding API (any provider) - api_key optional for auth-less self-hosted Milvus; supports db_name/partition - Registered in INGESTION_REGISTRY; MilvusVectorStoreOptions added to types - 16 unit tests (mocked REST) + env-gated integration test
Milvus ingestion writes to collection_name (falling back to vector_store_id), but /rag/ingest only authorized fields named vector_store_id. A request with custom_llm_provider=milvus and collection_name set to another team's managed collection bypassed assert_user_can_access_vector_store_id. Normalize collection_name into vector_store_id before authorization.
Resolves the Greptile review on the Milvus RAG ingestion path: - P0 (security): vector-store-id normalization for authorization now always mirrors collection_name onto vector_store_id for Milvus, not only when vector_store_id is absent. A request pairing a collection_name the caller cannot access with a vector_store_id they can no longer bypasses assert_user_can_access_vector_store_id. Adds a test for the both-fields case. - P1: removes the provider-specific `custom_llm_provider == "milvus"` branch from proxy/rag_endpoints/endpoints.py. BaseRAGIngestion now exposes a normalize_authorized_vector_store_id classmethod (no-op by default) that MilvusRAGIngestion overrides; the proxy dispatches generically via get_ingestion_class. - P2: removes the embed() side-effect that mutated self.embedding_config on first call. The default model is set once in MilvusRAGIngestion.__init__ and the class inherits BaseRAGIngestion.embed. Drops the now-unused top-level `import litellm` (also clears the CodeQL import/import-from warning).
Require INTERNAL_USER_VIEW_ONLY ingest targets to resolve to an existing managed vector store. Presence of vector_store_id was insufficient: Milvus normalization mirrors collection_name onto vector_store_id and unknown ids pass authorization as provider-native targets, letting a view-only caller trigger Milvus auto_create_collection for a brand-new collection.
Milvus db_name selects the write target's database namespace but the proxy only authorizes collection_name/vector_store_id. A caller with access to a managed collection could set db_name to redirect writes/auto-create into another Milvus database using the server's credentials, outside the per-collection authorization boundary. Resolve db_name from MILVUS_DB_NAME (server-side) only; never from the request. Drop db_name from MilvusVectorStoreOptions and add a regression test asserting a request-supplied db_name is ignored.
|
and the score is still 4/5 |
The view-only ingest guard required every vector_store_id to resolve to a litellm-managed store, which broke INTERNAL_USER_VIEW_ONLY callers writing to provider-native ids (e.g. OpenAI vs_*) that are not in the managed registry Only providers that can create a store on ingest (Milvus with auto_create_collection) let a view-only caller bring a brand-new store into existence, so the managed-store requirement now applies only to those. Each ingestion class declares this via can_auto_create_vector_store and the proxy dispatches to it instead of hardcoding provider logic. Providers that only write to a pre-existing store keep accepting their provider-native ids unchanged Also drops the banned typing imports from the new milvus_ingestion module so it stays within the strict-rule budget gate after the rebase onto litellm_internal_staging
|
@Sameerlite pushed The strict guard required every So I scoped the guard to that case. Each ingestion class now declares whether it can auto-create via This also answers your other question about the changed test: with the guard scoped, The rebase onto |
A named credential can carry api_base while leaving api_key unset, which slips a request-controlled endpoint past the proxy's api_base block. The constructor then fell back to MILVUS_API_KEY independently, sending the server token to that endpoint. Only fall back to the env token when api_base also comes from MILVUS_API_BASE.
…s of auto_create flag can_auto_create_vector_store read the request-supplied auto_create_collection flag, so a view-only key could set it to false, name any existing unmanaged collection, and skip the managed-store resolution check in _assert_view_only_role_cannot_create_vector_store. Report the provider's capability instead: Milvus can always auto-create, so a view-only target must always resolve to a managed vector store.
… feat/milvus-vector-store-ingest
|
TL;DR (update): Heads up on the With the gate actually running now, it reports 87 LIT009 findings on this PR's changed lines. A good chunk of them look like gate false positives that no concrete typing can resolve, so I wanted to ask how you'd like them handled before adding any suppressions. The first group is TypedDict field declarations, 9 findings in The rest are genuine |
Use PEP 585/604 builtins (dict, tuple, X | None) in the Milvus ingestion and RAG endpoint helpers so the strict-rule budget delta (UP006/UP035/UP045) stays under the lowered ceiling pulled in from staging.
|
Hey thanks for the heads up. I think you can just ignore it for now. I am going to make any-discipline much less strict. It seems to flag on way too many prs currently |
|
The any-discipline should be more lax now. Let me know if it's a blocker still |
|
@maycuatroi1 Still failing lint |
|
Thanks for the PR! A couple of things to get this over the finish line:
Once those are addressed we'll take another look — appreciate the contribution! |
|
@Sameerlite please check. The code now clean and beauty 😄 |
|
Thanks for your contribution! Triggering a code review now. |
|
@Sameerlite gentle nudge on this one. All 74 CI checks are green now, lint is clean, and the Greptile concerns from the earlier rounds have all been addressed. Whenever you get a chance, could you take another look and let me know if anything else is needed before merge? Happy to make further changes. Thanks! |
e5a339e
into
BerriAI:litellm_oss_staging_230626
* feat(rag): add Milvus vector store ingestion support Adds write/ingest support for self-hosted Milvus to complement the existing Milvus search provider. /rag/ingest now accepts custom_llm_provider=milvus. - MilvusRAGIngestion implements the store() step via the Milvus REST API v2 (entities/insert), reusing the base upload/ocr/chunk/embed pipeline - Auto-creates the collection via quick setup (dynamic fields) when missing - Embeddings generated through litellm embedding API (any provider) - api_key optional for auth-less self-hosted Milvus; supports db_name/partition - Registered in INGESTION_REGISTRY; MilvusVectorStoreOptions added to types - 16 unit tests (mocked REST) + env-gated integration test * fix(rag): authorize Milvus collection_name as vector_store_id on ingest Milvus ingestion writes to collection_name (falling back to vector_store_id), but /rag/ingest only authorized fields named vector_store_id. A request with custom_llm_provider=milvus and collection_name set to another team's managed collection bypassed assert_user_can_access_vector_store_id. Normalize collection_name into vector_store_id before authorization. * fix(rag): close Milvus collection_name authz bypass and address review Resolves the Greptile review on the Milvus RAG ingestion path: - P0 (security): vector-store-id normalization for authorization now always mirrors collection_name onto vector_store_id for Milvus, not only when vector_store_id is absent. A request pairing a collection_name the caller cannot access with a vector_store_id they can no longer bypasses assert_user_can_access_vector_store_id. Adds a test for the both-fields case. - P1: removes the provider-specific `custom_llm_provider == "milvus"` branch from proxy/rag_endpoints/endpoints.py. BaseRAGIngestion now exposes a normalize_authorized_vector_store_id classmethod (no-op by default) that MilvusRAGIngestion overrides; the proxy dispatches generically via get_ingestion_class. - P2: removes the embed() side-effect that mutated self.embedding_config on first call. The default model is set once in MilvusRAGIngestion.__init__ and the class inherits BaseRAGIngestion.embed. Drops the now-unused top-level `import litellm` (also clears the CodeQL import/import-from warning). * fix(rag): block view-only role from auto-creating Milvus collections Require INTERNAL_USER_VIEW_ONLY ingest targets to resolve to an existing managed vector store. Presence of vector_store_id was insufficient: Milvus normalization mirrors collection_name onto vector_store_id and unknown ids pass authorization as provider-native targets, letting a view-only caller trigger Milvus auto_create_collection for a brand-new collection. * fix(rag): authorize Milvus db_name via server env only Milvus db_name selects the write target's database namespace but the proxy only authorizes collection_name/vector_store_id. A caller with access to a managed collection could set db_name to redirect writes/auto-create into another Milvus database using the server's credentials, outside the per-collection authorization boundary. Resolve db_name from MILVUS_DB_NAME (server-side) only; never from the request. Drop db_name from MilvusVectorStoreOptions and add a regression test asserting a request-supplied db_name is ignored. * fix(rag): authorize Milvus partition_name via server env only * fix(rag): scope view-only ingest guard to auto-creating providers The view-only ingest guard required every vector_store_id to resolve to a litellm-managed store, which broke INTERNAL_USER_VIEW_ONLY callers writing to provider-native ids (e.g. OpenAI vs_*) that are not in the managed registry Only providers that can create a store on ingest (Milvus with auto_create_collection) let a view-only caller bring a brand-new store into existence, so the managed-store requirement now applies only to those. Each ingestion class declares this via can_auto_create_vector_store and the proxy dispatches to it instead of hardcoding provider logic. Providers that only write to a pre-existing store keep accepting their provider-native ids unchanged Also drops the banned typing imports from the new milvus_ingestion module so it stays within the strict-rule budget gate after the rebase onto litellm_internal_staging * fix(rag): bind Milvus api_key fallback to server-resolved api_base A named credential can carry api_base while leaving api_key unset, which slips a request-controlled endpoint past the proxy's api_base block. The constructor then fell back to MILVUS_API_KEY independently, sending the server token to that endpoint. Only fall back to the env token when api_base also comes from MILVUS_API_BASE. * fix(rag): require managed store for view-only Milvus ingest regardless of auto_create flag can_auto_create_vector_store read the request-supplied auto_create_collection flag, so a view-only key could set it to false, name any existing unmanaged collection, and skip the managed-store resolution check in _assert_view_only_role_cannot_create_vector_store. Report the provider's capability instead: Milvus can always auto-create, so a view-only target must always resolve to a managed vector store. * style(rag): apply black formatting to Milvus ingest files * style(rag): modernize typing to satisfy ruff strict-rule budget Use PEP 585/604 builtins (dict, tuple, X | None) in the Milvus ingestion and RAG endpoint helpers so the strict-rule budget delta (UP006/UP035/UP045) stays under the lowered ceiling pulled in from staging. * style(rag): drop redundant quoted annotations to satisfy UP037 budget * chore(rag): retrigger CI after transient artifact-download 403 * fix(rag): block credential hydration from overriding authorized write target
* feat(rag): add Milvus vector store ingestion support Adds write/ingest support for self-hosted Milvus to complement the existing Milvus search provider. /rag/ingest now accepts custom_llm_provider=milvus. - MilvusRAGIngestion implements the store() step via the Milvus REST API v2 (entities/insert), reusing the base upload/ocr/chunk/embed pipeline - Auto-creates the collection via quick setup (dynamic fields) when missing - Embeddings generated through litellm embedding API (any provider) - api_key optional for auth-less self-hosted Milvus; supports db_name/partition - Registered in INGESTION_REGISTRY; MilvusVectorStoreOptions added to types - 16 unit tests (mocked REST) + env-gated integration test * fix(rag): authorize Milvus collection_name as vector_store_id on ingest Milvus ingestion writes to collection_name (falling back to vector_store_id), but /rag/ingest only authorized fields named vector_store_id. A request with custom_llm_provider=milvus and collection_name set to another team's managed collection bypassed assert_user_can_access_vector_store_id. Normalize collection_name into vector_store_id before authorization. * fix(rag): close Milvus collection_name authz bypass and address review Resolves the Greptile review on the Milvus RAG ingestion path: - P0 (security): vector-store-id normalization for authorization now always mirrors collection_name onto vector_store_id for Milvus, not only when vector_store_id is absent. A request pairing a collection_name the caller cannot access with a vector_store_id they can no longer bypasses assert_user_can_access_vector_store_id. Adds a test for the both-fields case. - P1: removes the provider-specific `custom_llm_provider == "milvus"` branch from proxy/rag_endpoints/endpoints.py. BaseRAGIngestion now exposes a normalize_authorized_vector_store_id classmethod (no-op by default) that MilvusRAGIngestion overrides; the proxy dispatches generically via get_ingestion_class. - P2: removes the embed() side-effect that mutated self.embedding_config on first call. The default model is set once in MilvusRAGIngestion.__init__ and the class inherits BaseRAGIngestion.embed. Drops the now-unused top-level `import litellm` (also clears the CodeQL import/import-from warning). * fix(rag): block view-only role from auto-creating Milvus collections Require INTERNAL_USER_VIEW_ONLY ingest targets to resolve to an existing managed vector store. Presence of vector_store_id was insufficient: Milvus normalization mirrors collection_name onto vector_store_id and unknown ids pass authorization as provider-native targets, letting a view-only caller trigger Milvus auto_create_collection for a brand-new collection. * fix(rag): authorize Milvus db_name via server env only Milvus db_name selects the write target's database namespace but the proxy only authorizes collection_name/vector_store_id. A caller with access to a managed collection could set db_name to redirect writes/auto-create into another Milvus database using the server's credentials, outside the per-collection authorization boundary. Resolve db_name from MILVUS_DB_NAME (server-side) only; never from the request. Drop db_name from MilvusVectorStoreOptions and add a regression test asserting a request-supplied db_name is ignored. * fix(rag): authorize Milvus partition_name via server env only * fix(rag): scope view-only ingest guard to auto-creating providers The view-only ingest guard required every vector_store_id to resolve to a litellm-managed store, which broke INTERNAL_USER_VIEW_ONLY callers writing to provider-native ids (e.g. OpenAI vs_*) that are not in the managed registry Only providers that can create a store on ingest (Milvus with auto_create_collection) let a view-only caller bring a brand-new store into existence, so the managed-store requirement now applies only to those. Each ingestion class declares this via can_auto_create_vector_store and the proxy dispatches to it instead of hardcoding provider logic. Providers that only write to a pre-existing store keep accepting their provider-native ids unchanged Also drops the banned typing imports from the new milvus_ingestion module so it stays within the strict-rule budget gate after the rebase onto litellm_internal_staging * fix(rag): bind Milvus api_key fallback to server-resolved api_base A named credential can carry api_base while leaving api_key unset, which slips a request-controlled endpoint past the proxy's api_base block. The constructor then fell back to MILVUS_API_KEY independently, sending the server token to that endpoint. Only fall back to the env token when api_base also comes from MILVUS_API_BASE. * fix(rag): require managed store for view-only Milvus ingest regardless of auto_create flag can_auto_create_vector_store read the request-supplied auto_create_collection flag, so a view-only key could set it to false, name any existing unmanaged collection, and skip the managed-store resolution check in _assert_view_only_role_cannot_create_vector_store. Report the provider's capability instead: Milvus can always auto-create, so a view-only target must always resolve to a managed vector store. * style(rag): apply black formatting to Milvus ingest files * style(rag): modernize typing to satisfy ruff strict-rule budget Use PEP 585/604 builtins (dict, tuple, X | None) in the Milvus ingestion and RAG endpoint helpers so the strict-rule budget delta (UP006/UP035/UP045) stays under the lowered ceiling pulled in from staging. * style(rag): drop redundant quoted annotations to satisfy UP037 budget * chore(rag): retrigger CI after transient artifact-download 403 * fix(rag): block credential hydration from overriding authorized write target
5 participants
Relevant issues
Adds write/ingest support for Milvus to the
/rag/ingestpipeline. Milvus is currently supported only for search/retrieval (litellm/llms/milvus/vector_stores); this PR brings it to parity by letting users ingest documents into a self-hosted Milvus throughcustom_llm_provider="milvus".Related: #26771. That issue reports the same underlying limitation from the other end - a self-hosted vector store (pg_vector) failing ingest with
Provider 'X' is not supported for RAG ingestion. Supported providers: openai, bedrock, gemini, s3_vectors, vertex_ai, raised from the exact same code path (litellm/rag/main.py->get_ingestion_class->INGESTION_REGISTRY). I hit the same wall while wiring up a self-hosted Milvus for RAG ingest, and this PR fixes it for Milvus by extending that registry.Pre-Submission checklist
Type
🆕 New Feature
Changes
litellm/rag/ingestion/milvus_ingestion.py(MilvusRAGIngestion) implementing thestore()step on top of the existingBaseRAGIngestionupload/OCR/chunk/embed pipeline./v2/vectordb/entities/insert) via httpx - nopymilvusdependency.api_keyis optional (self-hosted Milvus without auth); supportsdb_nameandpartition_name."milvus"inINGESTION_REGISTRY(litellm/rag/main.py).MilvusVectorStoreOptionsTypedDict and included it inRAGIngestVectorStoreOptions(litellm/types/rag.py).MILVUS_DEFAULT_VECTOR_FIELD/MILVUS_DEFAULT_TEXT_FIELD/MILVUS_DEFAULT_METRIC_TYPEconstants.Usage
Tests
tests/test_litellm/rag/test_milvus_ingestion.py- 16 unit tests (mocked Milvus REST): config/defaults, optional auth header, auto-create vs existing collection,auto_create_collection=False, insert body shape (vector/text/metadata), error propagation on non-zero Milvus code, db/partition propagation, embedding via litellm vs router, and the collection-exists error fallback.tests/vector_store_tests/rag/test_rag_milvus.py- env-gated integration test (skipped unlessMILVUS_API_BASEset), follows the existingBaseRAGTestpattern.Working proof - live
/rag/ingestagainstmilvusdb/milvus:v2.5.4Auto-created collection is queryable (
POST /v2/vectordb/entities/query):{"code":0,"data":[{"chunk_index":0,"filename":"devproof.txt","text":"LiteLLM PR 30388 proof. Photosynthesis ..."}]}Shows up in the dashboard as: