feat(rag): add Milvus vector store ingestion support

.github/workflows/test-unit-misc.yml

@@ -33,6 +33,7 @@ jobs:
         tests/test_litellm/images
         tests/test_litellm/interactions
         tests/test_litellm/passthrough
+        tests/test_litellm/rag
         tests/test_litellm/vector_stores
         tests/test_litellm/test_*.py
       workers: 2

litellm/constants.py

@@ -1751,6 +1751,11 @@
 )
 S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS = ["source_text"]
 
+########################### Milvus RAG Constants ###########################
+MILVUS_DEFAULT_VECTOR_FIELD = "vector"
+MILVUS_DEFAULT_TEXT_FIELD = "text"
+MILVUS_DEFAULT_METRIC_TYPE = "COSINE"
+
 ########################### Microsoft SSO Constants ###########################
 MICROSOFT_USER_EMAIL_ATTRIBUTE = str(
     os.getenv("MICROSOFT_USER_EMAIL_ATTRIBUTE", "userPrincipalName")

litellm/proxy/rag_endpoints/endpoints.py

@@ -28,6 +28,7 @@
     assert_user_can_access_vector_store_id,
 )
 from litellm.repositories.table_repositories import ManagedVectorStoresRepository
+from litellm.rag.main import get_ingestion_class
 
 router = APIRouter()
 
@@ -96,6 +97,32 @@ def _collect_vector_store_ids_from_payload(payload: Any) -> set[str]:
     return vector_store_ids
 
 
+def _normalize_collection_name_as_vector_store_id(
+    ingest_options: dict[str, Any],
+) -> None:
+    """
+    Normalize provider-native write targets to `vector_store_id` for authorization.
+
+    The proxy authorizes ingestion by the `vector_store_id` key, but some
+    providers resolve their real write target from a different field (e.g. Milvus
+    uses `collection_name`). Each ingestion class owns that knowledge via
+    `normalize_authorized_vector_store_id`, so dispatch to it instead of hardcoding
+    provider-specific logic here. This closes the bypass where a caller pairs a
+    write target they cannot access with a `vector_store_id` they can.
+    """
+    vector_store_opts = ingest_options.get("vector_store")
+    if not isinstance(vector_store_opts, dict):
+        return
+    provider = vector_store_opts.get("custom_llm_provider")
+    if not provider:
+        return
+    try:
+        ingestion_class = get_ingestion_class(provider)
+    except ValueError:
+        return
+    ingestion_class.normalize_authorized_vector_store_id(vector_store_opts)
+
+
 async def _authorize_nested_vector_store_ids(
     payload: Any,
     user_api_key_dict: UserAPIKeyAuth,
@@ -107,6 +134,68 @@ async def _authorize_nested_vector_store_ids(
         )
 
 
+def _ingestion_can_auto_create_vector_store(
+    vector_store_opts: dict[str, Any],
+) -> bool:
+    """
+    Whether this ingestion could create a brand-new vector store on write.
+
+    Each ingestion class owns that knowledge via `can_auto_create_vector_store`,
+    so dispatch to it instead of hardcoding provider-specific logic here.
+    """
+    provider = vector_store_opts.get("custom_llm_provider")
+    if not provider:
+        return False
+    try:
+        ingestion_class = get_ingestion_class(provider)
+    except ValueError:
+        return False
+    return ingestion_class.can_auto_create_vector_store(vector_store_opts)
+
+
+async def _assert_view_only_role_cannot_create_vector_store(
+    ingest_options: dict[str, Any],
+    user_api_key_dict: UserAPIKeyAuth,
+) -> None:
+    """
+    INTERNAL_USER_VIEW_ONLY may ingest into an existing vector store but may not
+    create a new one.
+
+    The role must always name a `vector_store_id`. Beyond that, only providers
+    that auto-create the store on ingest (e.g. Milvus with `auto_create_collection`)
+    can let a view-only caller bring a brand-new store into existence; for those,
+    the id must resolve to an existing managed vector store. Providers that only
+    write to a pre-existing store (OpenAI, Bedrock, ...) keep accepting their
+    provider-native ids unchanged.
+    """
+    if user_api_key_dict.user_role != LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value:
+        return
+    vector_store_opts = ingest_options.get("vector_store") or {}
+    vector_store_id = vector_store_opts.get("vector_store_id")
+    if not vector_store_id:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail={
+                "error": "internal_user_viewer role can only ingest files to an existing vector store. "
+                "Provide 'vector_store_id' in ingest_options.vector_store."
+            },
+        )
+    if not _ingestion_can_auto_create_vector_store(vector_store_opts):
+        return
+    existing_vector_store = await assert_user_can_access_vector_store_id(
+        vector_store_id=vector_store_id,
+        user_api_key_dict=user_api_key_dict,
+    )
+    if existing_vector_store is None:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail={
+                "error": "internal_user_viewer role cannot create a new vector store. "
+                f"'{vector_store_id}' does not resolve to an existing managed vector store."
+            },
+        )
+
+
 def _build_file_metadata_entry(
     response: Any,
     file_data: Optional[Tuple[str, bytes, str]] = None,
@@ -418,6 +507,8 @@ async def parse_rag_ingest_request(
                     },
                 )
 
+    _normalize_collection_name_as_vector_store_id(ingest_options)
+
     return ingest_options, file_data, file_url, file_id
 
 
@@ -485,19 +576,11 @@ async def rag_ingest(
             request
         )
 
-        # INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores only
-        if (
-            user_api_key_dict.user_role
-            == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
-            and not ingest_options.get("vector_store", {}).get("vector_store_id")
-        ):
-            raise HTTPException(
-                status_code=status.HTTP_403_FORBIDDEN,
-                detail={
-                    "error": "internal_user_viewer role can only ingest files to an existing vector store. "
-                    "Provide 'vector_store_id' in ingest_options.vector_store."
-                },
-            )
+        # INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores, but cannot create new ones
+        await _assert_view_only_role_cannot_create_vector_store(
+            ingest_options=ingest_options,
+            user_api_key_dict=user_api_key_dict,
+        )
 
         await _authorize_nested_vector_store_ids(
             payload=ingest_options,

litellm/rag/ingestion/base_ingestion.py

@@ -85,7 +85,10 @@ def _load_credentials_from_config(self) -> None:
             )
             if not credential_values:
                 return
+            protected_fields = self.credential_protected_fields()
             for key, value in credential_values.items():
+                if key in protected_fields:
+                    continue
                 self.vector_store_config[key] = value
             for key in (
                 "api_base",
@@ -100,6 +103,45 @@ def custom_llm_provider(self) -> str:
         """Get the vector store provider."""
         return self.vector_store_config.get("custom_llm_provider", "openai")
 
+    @classmethod
+    def normalize_authorized_vector_store_id(
+        cls, vector_store_opts: dict[str, object]
+    ) -> None:
+        """
+        Rewrite the vector_store config so `vector_store_id` matches the actual
+        write target before the proxy authorizes it.
+
+        The proxy authorizes ingestion by the `vector_store_id` key. Providers
+        whose real write target is a different field (e.g. Milvus uses
+        `collection_name`) must override this so authorization covers the target
+        that will actually be written to. Default: no-op.
+        """
+        return None
+
+    @classmethod
+    def credential_protected_fields(cls) -> frozenset[str]:
+        """
+        Vector-store config keys that credential hydration must never override.
+
+        The proxy authorizes ingestion against `vector_store_id`, so letting a
+        stored credential redefine the write target after authorization would
+        bypass the access check. Providers whose real write target is a different
+        field (e.g. Milvus `collection_name`) must extend this set.
+        """
+        return frozenset({"vector_store_id"})
+
+    @classmethod
+    def can_auto_create_vector_store(cls, vector_store_opts: dict[str, object]) -> bool:
+        """
+        Whether ingesting can bring a brand-new vector store into existence.
+
+        Providers that only write to a pre-existing store return False. Providers
+        that create the store on demand (e.g. Milvus `auto_create_collection`)
+        must override this so the proxy can stop a view-only caller from creating
+        one. Default: False.
+        """
+        return False
+
     async def upload(
         self,
         file_data: Optional[Tuple[str, bytes, str]] = None,