diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index a7363ac3b43..a99bc2e5009 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.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 diff --git a/litellm/constants.py b/litellm/constants.py index a3ea68c7949..1ba5183a263 100644 --- a/litellm/constants.py +++ b/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") diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 7ff54ac4c5a..da7f3210de5 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/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, diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index 1c770d0a992..517ccd8b284 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/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, diff --git a/litellm/rag/ingestion/milvus_ingestion.py b/litellm/rag/ingestion/milvus_ingestion.py new file mode 100644 index 00000000000..f29afe56593 --- /dev/null +++ b/litellm/rag/ingestion/milvus_ingestion.py @@ -0,0 +1,269 @@ +""" +Milvus-specific RAG Ingestion implementation. + +Milvus is an open-source, self-hostable vector database. This implementation +adds write/ingest support to complement the existing Milvus vector store +search provider (litellm/llms/milvus/vector_stores). + +This implementation: +1. Generates embeddings using LiteLLM's embedding API (supports any provider) +2. Auto-creates the target collection via the Milvus REST "quick setup" API + when it does not exist (dynamic fields enabled so chunk text + metadata are + stored alongside the vector) +3. Inserts chunks + embeddings via the Milvus REST `entities/insert` API + +It talks to Milvus over the REST API v2 (`/v2/vectordb/...`) using httpx, so it +does not add a `pymilvus` dependency. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from litellm._logging import verbose_logger +from litellm.constants import ( + MILVUS_DEFAULT_METRIC_TYPE, + MILVUS_DEFAULT_TEXT_FIELD, + MILVUS_DEFAULT_VECTOR_FIELD, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm import Router + from litellm.types.rag import RAGIngestOptions + + +class MilvusRAGIngestion(BaseRAGIngestion): + """ + Milvus RAG ingestion using the Milvus REST API v2 + httpx. + + Workflow: + 1. Generate embeddings using LiteLLM (supports any embedding provider) + 2. Auto-create the collection if needed (quick setup, dynamic fields on) + 3. Insert chunks + embeddings via `entities/insert` + + Configuration (vector_store config): + - collection_name / vector_store_id: target Milvus collection (required) + - api_base: Milvus REST base URL (or MILVUS_API_BASE env) + - api_key: Milvus token/credential. The MILVUS_API_KEY env fallback applies + only when api_base also comes from the MILVUS_API_BASE env, so the server + token is never sent to a config/credential-supplied endpoint. Optional for + a self-hosted Milvus without auth. + - vector_field: embedding field name (default: "vector") + - text_field: chunk text field name (default: "text") + - metric_type: distance metric for auto-created collection (default: "COSINE") + - db_name: Milvus database namespace, server-side only via MILVUS_DB_NAME env. + Not accepted from the request: it selects the write target's database and is + outside the per-collection authorization boundary. + - partition_name: target partition, server-side only via MILVUS_PARTITION_NAME + env. Not accepted from the request: it selects the write target's partition + and is outside the per-collection authorization boundary. + - auto_create_collection: create the collection if missing (default: True) + """ + + def __init__( + self, + ingest_options: RAGIngestOptions, + router: Router | None = None, + ): + BaseRAGIngestion.__init__(self, ingest_options=ingest_options, router=router) + + if not self.embedding_config: + self.embedding_config = {"model": "text-embedding-3-small"} + + self.collection_name = self.vector_store_config.get( + "collection_name" + ) or self.vector_store_config.get("vector_store_id") + if not self.collection_name: + raise ValueError( + "Milvus RAG ingestion requires 'collection_name' (or 'vector_store_id') in the vector_store config." + ) + + config_api_base = self.vector_store_config.get("api_base") + self.api_base = config_api_base or get_secret_str("MILVUS_API_BASE") + if not self.api_base: + raise ValueError( + "Milvus API base URL is required. Set the MILVUS_API_BASE environment " + "variable or pass 'api_base' in the vector_store config." + ) + self.api_base = self.api_base.rstrip("/") + + config_api_key = self.vector_store_config.get("api_key") + self.api_key = ( + config_api_key + if config_api_base + else config_api_key or get_secret_str("MILVUS_API_KEY") + ) + self.vector_field = self.vector_store_config.get( + "vector_field", MILVUS_DEFAULT_VECTOR_FIELD + ) + self.text_field = self.vector_store_config.get( + "text_field", MILVUS_DEFAULT_TEXT_FIELD + ) + self.metric_type = self.vector_store_config.get( + "metric_type", MILVUS_DEFAULT_METRIC_TYPE + ) + self.db_name = get_secret_str("MILVUS_DB_NAME") + self.partition_name = get_secret_str("MILVUS_PARTITION_NAME") + self.auto_create_collection = self.vector_store_config.get( + "auto_create_collection", True + ) + + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.RAG + ) + + @classmethod + def normalize_authorized_vector_store_id( + cls, vector_store_opts: dict[str, object] + ) -> None: + """ + Milvus resolves its write target from `collection_name` first (falling + back to `vector_store_id`). Always mirror `collection_name` onto + `vector_store_id` so authorization covers the collection that will be + written to - even when the caller supplies a different `vector_store_id` + they happen to have access to. + """ + collection_name = vector_store_opts.get("collection_name") + if collection_name: + vector_store_opts["vector_store_id"] = collection_name + + @classmethod + def credential_protected_fields(cls) -> frozenset[str]: + """ + Milvus selects its write target from `collection_name` (mirrored onto + `vector_store_id` for authorization), so both must be shielded from + credential hydration to keep the authorized target intact. + """ + return super().credential_protected_fields() | {"collection_name"} + + @classmethod + def can_auto_create_vector_store(cls, vector_store_opts: dict[str, object]) -> bool: + """ + Milvus is capable of creating the target collection on ingest, so the + view-only guard must always require the target to resolve to a managed + vector store. This reports the provider's capability, not the + request-supplied `auto_create_collection` flag: that flag is caller + controlled, and trusting it would let a view-only key set it to false, + name any existing collection, and skip the managed-store check. + """ + return True + + def _headers(self) -> dict[str, str]: + headers = {"Content-Type": "application/json", "Accept": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + return headers + + def _with_db(self, body: dict[str, object]) -> dict[str, object]: + if self.db_name: + body["dbName"] = self.db_name + return body + + async def _post(self, path: str, body: dict[str, object]) -> dict[str, object]: + url = f"{self.api_base}{path}" + response = await self.async_httpx_client.post( + url, json=body, headers=self._headers() + ) + response.raise_for_status() + data = response.json() + # Milvus REST returns {"code": 0, "data": ...} on success, + # {"code": , "message": "..."} on error. + if isinstance(data, dict) and data.get("code") not in (0, None): + raise RuntimeError( + f"Milvus API call to {path} failed: code={data.get('code')} message={data.get('message')}" + ) + return data + + async def _collection_exists(self) -> bool: + try: + data = await self._post( + "/v2/vectordb/collections/has", + self._with_db({"collectionName": self.collection_name}), + ) + inner = data.get("data") + return bool(inner.get("has")) if isinstance(inner, dict) else False + except Exception as e: + verbose_logger.debug(f"Milvus collection 'has' check failed: {e}") + return False + + async def _ensure_collection_exists(self, dimension: int) -> None: + if not self.auto_create_collection: + return + if await self._collection_exists(): + return + + verbose_logger.debug( + f"Creating Milvus collection '{self.collection_name}' (dimension={dimension}, metric={self.metric_type})" + ) + # Quick-setup create: enables a dynamic field so chunk text + metadata + # are stored alongside the vector without declaring a full schema. + body = self._with_db( + { + "collectionName": self.collection_name, + "dimension": dimension, + "metricType": self.metric_type, + "vectorFieldName": self.vector_field, + "autoId": True, + "enableDynamicField": True, + } + ) + await self._post("/v2/vectordb/collections/create", body) + + async def store( + self, + file_content: bytes | None, + filename: str | None, + content_type: str | None, + chunks: list[str], + embeddings: list[list[float]] | None, + ) -> tuple[str | None, str | None]: + """ + Insert chunks + embeddings into a Milvus collection. + + Steps: + 1. Validate chunks/embeddings were produced + 2. Ensure the collection exists (auto-create quick setup if needed) + 3. Insert rows via `entities/insert` + + Returns: + Tuple of (collection_name, filename) + """ + if not embeddings or not chunks: + raise ValueError( + "No text content could be extracted from the file for embedding. " + "Possible causes:\n" + " 1. PDF files require OCR - add an 'ocr' config with a vision model " + "(e.g., 'anthropic/claude-3-5-sonnet-20241022')\n" + " 2. Binary files cannot be processed - convert to text first\n" + " 3. File is empty or contains no extractable text" + ) + + await self._ensure_collection_exists(dimension=len(embeddings[0])) + + rows: list[dict[str, object]] = [] + for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): + row: dict[str, object] = { + self.vector_field: embedding, + self.text_field: chunk, + "chunk_index": i, + } + if filename: + row["filename"] = filename + rows.append(row) + + body = self._with_db({"collectionName": self.collection_name, "data": rows}) + if self.partition_name: + body["partitionName"] = self.partition_name + + await self._post("/v2/vectordb/entities/insert", body) + verbose_logger.info( + f"Inserted {len(rows)} vectors into Milvus collection '{self.collection_name}'" + ) + + return self.collection_name, filename diff --git a/litellm/rag/main.py b/litellm/rag/main.py index e3d354b6c33..5740c60f11b 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -30,6 +30,7 @@ from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion +from litellm.rag.ingestion.milvus_ingestion import MilvusRAGIngestion from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion from litellm.rag.ingestion.vertex_ai_ingestion import VertexAIRAGIngestion @@ -52,6 +53,7 @@ "gemini": GeminiRAGIngestion, "s3_vectors": S3VectorsRAGIngestion, "vertex_ai": VertexAIRAGIngestion, + "milvus": MilvusRAGIngestion, } diff --git a/litellm/types/rag.py b/litellm/types/rag.py index 29e35d5fe8d..b89a55eb5bd 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -193,12 +193,48 @@ class S3VectorsVectorStoreOptions(TypedDict, total=False): aws_external_id: Optional[str] +class MilvusVectorStoreOptions(TypedDict, total=False): + """ + Milvus (self-hostable open-source vector database) configuration. + + Example (auto-create collection): + {"custom_llm_provider": "milvus", "collection_name": "my_docs", + "api_base": "http://localhost:19530"} + + Example (existing collection, no auth): + {"custom_llm_provider": "milvus", "collection_name": "my_docs", + "api_base": "http://localhost:19530", "auto_create_collection": False} + + Embeddings are generated using LiteLLM's embedding API (supports any provider). + When the collection does not exist it is created with the Milvus REST + "quick setup" API (dynamic fields enabled so chunk text + metadata are + stored alongside the vector). + """ + + custom_llm_provider: Literal["milvus"] + collection_name: str # Target Milvus collection (alias: vector_store_id) + vector_store_id: Optional[str] # Alternative to collection_name + + # Connection + api_base: Optional[str] # Milvus REST base URL (or MILVUS_API_BASE env) + api_key: Optional[str] # Milvus token (or MILVUS_API_KEY env); optional if no auth + + # Schema / collection config (used for auto-creation) + vector_field: Optional[str] # Embedding field name (default: "vector") + text_field: Optional[str] # Chunk text field name (default: "text") + metric_type: Optional[str] # Distance metric (default: "COSINE") + auto_create_collection: Optional[ + bool + ] # Create collection if missing (default: True) + + # Union type for vector store options RAGIngestVectorStoreOptions = Union[ OpenAIVectorStoreOptions, BedrockVectorStoreOptions, VertexAIVectorStoreOptions, S3VectorsVectorStoreOptions, + MilvusVectorStoreOptions, ] 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 656e1406f07..192ce6b00c7 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -8,9 +8,10 @@ import io import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient sys.path.insert( @@ -83,7 +84,7 @@ def test_internal_user_viewer_rag_ingest_with_vector_store_id_passes_check( client_internal_user_viewer, ): """ - internal_user_viewer with vector_store_id passes the role check. + internal_user_viewer with a vector_store_id passes the role check. (Actual ingest may fail due to missing API keys, but we get past 403.) """ with patch( @@ -100,9 +101,167 @@ def test_internal_user_viewer_rag_ingest_with_vector_store_id_passes_check( ) # Should not be 403 (role check passed) + assert ( + response.status_code != 403 + ), f"internal_user_viewer with vector_store_id should pass role check. Response: {response.json()}" + + +def test_internal_user_viewer_provider_native_vector_store_id_allowed( + client_internal_user_viewer, +): + """ + Regression: a view-only caller may still ingest into a provider-native + vector store id that is not in litellm's managed registry, as long as the + provider does not auto-create the store. Only auto-creating providers + (e.g. Milvus) require the id to resolve to a managed store, so an OpenAI id + must not be rejected just for being unregistered. + """ + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "vs_provider_native", "file_id": "f"}, + ), + ): + response = client_internal_user_viewer.post( + "/v1/rag/ingest", + files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, + data={ + "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai","vector_store_id":"vs_provider_native"}}}' + }, + ) + + assert response.status_code != 403, ( + "view-only ingest into an unregistered provider-native (non-auto-create) " + f"vector store id must be allowed. Got: {response.status_code} {response.json()}" + ) + + +def test_internal_user_viewer_milvus_collection_name_auto_create_rejected( + client_internal_user_viewer, +): + """ + internal_user_viewer must not be able to auto-create a new Milvus collection. + + The Milvus normalization mirrors collection_name onto vector_store_id, but an + unknown id resolves to no managed vector store, so the view-only caller is + denied before Milvus auto_create_collection can fire. + """ + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "brand_new_collection", "file_id": "f"}, + ) as mock_aingest, + ): + response = client_internal_user_viewer.post( + "/v1/rag/ingest", + json={ + "file_url": "https://example.com/doc.pdf", + "ingest_options": { + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "brand_new_collection", + } + }, + }, + ) + + assert response.status_code == 403, ( + "internal_user_viewer creating a new Milvus collection via collection_name " + f"must be denied. Got: {response.status_code} {response.json()}" + ) + mock_aingest.assert_not_called() + + +def test_internal_user_viewer_milvus_auto_create_disabled_still_requires_managed_store( + client_internal_user_viewer, +): + """ + Regression: auto_create_collection is request-controlled, so a view-only key + must not be able to set it to false, name any existing unmanaged collection, + and skip the managed-store check. Milvus can always auto-create, so the + target must still resolve to a managed vector store regardless of the flag; + an unmanaged collection (assert_user_can_access_vector_store_id returns None) + is denied before aingest fires. + """ + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "existing_collection", "file_id": "f"}, + ) as mock_aingest, + ): + response = client_internal_user_viewer.post( + "/v1/rag/ingest", + json={ + "file_url": "https://example.com/doc.pdf", + "ingest_options": { + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "existing_collection", + "auto_create_collection": False, + } + }, + }, + ) + + assert response.status_code == 403, ( + "view-only Milvus ingest with auto_create_collection disabled must still " + "require the collection to resolve to a managed vector store. " + f"Got: {response.status_code} {response.json()}" + ) + mock_aingest.assert_not_called() + + +def test_internal_user_viewer_milvus_managed_collection_passes_with_auto_create_disabled( + client_internal_user_viewer, +): + """ + The capability-based managed-store requirement must not over-block: when the + Milvus collection does resolve to a managed vector store the caller can + access, a view-only ingest with auto_create_collection disabled passes. + """ + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", + new=AsyncMock(return_value=object()), + ), + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "managed_collection", "file_id": "f"}, + ), + ): + response = client_internal_user_viewer.post( + "/v1/rag/ingest", + json={ + "file_url": "https://example.com/doc.pdf", + "ingest_options": { + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "managed_collection", + "auto_create_collection": False, + } + }, + }, + ) + assert response.status_code != 403, ( - f"internal_user_viewer with vector_store_id should pass role check. " - f"Response: {response.json()}" + "view-only Milvus ingest into a managed collection must pass the role " + f"check. Got: {response.status_code} {response.json()}" ) @@ -124,10 +283,9 @@ def test_internal_user_rag_ingest_without_vector_store_id_allowed(client_interna ) # Should not be 403 - assert response.status_code != 403, ( - f"internal_user should be allowed to create new vector stores. " - f"Response: {response.json()}" - ) + assert ( + response.status_code != 403 + ), f"internal_user should be allowed to create new vector stores. Response: {response.json()}" @pytest.mark.parametrize( @@ -181,6 +339,8 @@ def test_rag_ingest_blocks_clientside_credentials(client_internal_user, blocked_ assert blocked_field in str( body ), f"Response should mention '{blocked_field}': {body}" + + class TestRagIngestSSRFBlocked: """ aws_sts_endpoint and related credential-redirect fields must be rejected @@ -222,7 +382,9 @@ def test_ssrf_field_in_vector_store_config_rejected( error_text = ( detail.get("error", "") if isinstance(detail, dict) else str(detail) ) - assert field in error_text, f"Error should name the offending field: {error_text}" + assert ( + field in error_text + ), f"Error should name the offending field: {error_text}" def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user): with patch( @@ -239,6 +401,215 @@ def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user): }, }, ) - assert response.status_code != 400, ( - f"Clean Bedrock ingest_options should not be rejected: {response.json()}" + assert ( + response.status_code != 400 + ), f"Clean Bedrock ingest_options should not be rejected: {response.json()}" + + +class TestMilvusCollectionNameAuthorization: + """ + Milvus ingestion writes to `collection_name` (falling back to + `vector_store_id`). The proxy authorizes write targets by the + `vector_store_id` key, so a request that sets only `collection_name` must be + normalized so it is authorized as the vector store id - otherwise an + authenticated user could write into another team's managed collection with + the server's Milvus credentials. + """ + + def test_normalize_copies_collection_name_to_vector_store_id(self): + from litellm.proxy.rag_endpoints.endpoints import ( + _normalize_collection_name_as_vector_store_id, + ) + + ingest_options = { + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "other_team_collection", + } + } + _normalize_collection_name_as_vector_store_id(ingest_options) + assert ( + ingest_options["vector_store"]["vector_store_id"] == "other_team_collection" + ) + + def test_normalize_overrides_vector_store_id_with_collection_name(self): + """ + Sending both fields must not let a caller authorize against a + `vector_store_id` they can access while writing to a different + `collection_name`. The collection_name (the real write target) always + wins so authorization covers it. + """ + from litellm.proxy.rag_endpoints.endpoints import ( + _normalize_collection_name_as_vector_store_id, + ) + + ingest_options = { + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "other_team_collection", + "vector_store_id": "vs_caller_can_access", + } + } + _normalize_collection_name_as_vector_store_id(ingest_options) + assert ( + ingest_options["vector_store"]["vector_store_id"] == "other_team_collection" + ) + + def test_normalize_ignores_non_milvus_providers(self): + from litellm.proxy.rag_endpoints.endpoints import ( + _normalize_collection_name_as_vector_store_id, + ) + + ingest_options = { + "vector_store": { + "custom_llm_provider": "openai", + "collection_name": "col_a", + } + } + _normalize_collection_name_as_vector_store_id(ingest_options) + assert "vector_store_id" not in ingest_options["vector_store"] + + def test_milvus_collection_name_is_authorized(self, client_internal_user): + async def fake_assert(vector_store_id, user_api_key_dict, **kwargs): + if vector_store_id == "other_team_collection": + raise HTTPException( + status_code=403, + detail={"error": "Access denied"}, + ) + return None + + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", + new=AsyncMock(side_effect=fake_assert), + ), + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={ + "vector_store_id": "other_team_collection", + "file_id": "f", + }, + ), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + json={ + "file_url": "https://example.com/doc.pdf", + "ingest_options": { + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "other_team_collection", + } + }, + }, + ) + assert response.status_code == 403, ( + "Milvus ingest targeting another team's collection via collection_name " + f"must be authorized and denied. Got: {response.status_code} {response.json()}" + ) + + def test_milvus_collection_name_bypass_with_both_fields_is_denied( + self, client_internal_user + ): + async def fake_assert(vector_store_id, user_api_key_dict, **kwargs): + if vector_store_id == "other_team_collection": + raise HTTPException( + status_code=403, + detail={"error": "Access denied"}, + ) + return None + + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", + new=AsyncMock(side_effect=fake_assert), + ), + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={ + "vector_store_id": "other_team_collection", + "file_id": "f", + }, + ), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + json={ + "file_url": "https://example.com/doc.pdf", + "ingest_options": { + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "other_team_collection", + "vector_store_id": "vs_caller_can_access", + } + }, + }, + ) + assert response.status_code == 403, ( + "Pairing an authorized vector_store_id with an unauthorized " + "collection_name must still be denied. Got: " + f"{response.status_code} {response.json()}" + ) + + +class TestIngestionAutoCreateDetection: + """ + The view-only guard only requires managed-store resolution for ingestions + that can create a store on write. That decision is owned per-provider via + `can_auto_create_vector_store` and dispatched by + `_ingestion_can_auto_create_vector_store`. + """ + + def test_milvus_auto_creates_by_default(self): + from litellm.proxy.rag_endpoints.endpoints import ( + _ingestion_can_auto_create_vector_store, + ) + + assert ( + _ingestion_can_auto_create_vector_store( + {"custom_llm_provider": "milvus", "collection_name": "c"} + ) + is True + ) + + def test_milvus_auto_create_flag_is_not_request_trusted(self): + from litellm.proxy.rag_endpoints.endpoints import ( + _ingestion_can_auto_create_vector_store, + ) + + assert ( + _ingestion_can_auto_create_vector_store( + { + "custom_llm_provider": "milvus", + "collection_name": "c", + "auto_create_collection": False, + } + ) + is True + ) + + def test_openai_never_auto_creates(self): + from litellm.proxy.rag_endpoints.endpoints import ( + _ingestion_can_auto_create_vector_store, + ) + + assert ( + _ingestion_can_auto_create_vector_store( + {"custom_llm_provider": "openai", "vector_store_id": "vs_x"} + ) + is False + ) + + def test_unknown_provider_is_not_auto_create(self): + from litellm.proxy.rag_endpoints.endpoints import ( + _ingestion_can_auto_create_vector_store, + ) + + assert ( + _ingestion_can_auto_create_vector_store( + {"custom_llm_provider": "not_a_real_provider"} + ) + is False ) diff --git a/tests/test_litellm/rag/__init__.py b/tests/test_litellm/rag/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rag/test_milvus_ingestion.py b/tests/test_litellm/rag/test_milvus_ingestion.py new file mode 100644 index 00000000000..697c4c0490c --- /dev/null +++ b/tests/test_litellm/rag/test_milvus_ingestion.py @@ -0,0 +1,436 @@ +""" +Unit tests for Milvus RAG ingestion (litellm/rag/ingestion/milvus_ingestion.py). + +These tests mock the Milvus REST API via the async httpx client, so they run in +CI without a live Milvus instance. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.rag.ingestion.milvus_ingestion import MilvusRAGIngestion + + +def _make_ingestion(**vector_store_overrides): + vector_store = { + "custom_llm_provider": "milvus", + "collection_name": "test_collection", + "api_base": "http://localhost:19530", + "api_key": "root:Milvus", + } + vector_store.update(vector_store_overrides) + ingestion = MilvusRAGIngestion( + ingest_options={ + "embedding": {"model": "text-embedding-3-small"}, + "vector_store": vector_store, + } + ) + return ingestion + + +def _json_response(payload): + resp = MagicMock() + resp.raise_for_status = MagicMock() + resp.json = MagicMock(return_value=payload) + return resp + + +def test_requires_collection_name(): + with pytest.raises(ValueError, match="collection_name"): + MilvusRAGIngestion( + ingest_options={ + "vector_store": { + "custom_llm_provider": "milvus", + "api_base": "http://localhost:19530", + } + } + ) + + +def test_requires_api_base(monkeypatch): + monkeypatch.delenv("MILVUS_API_BASE", raising=False) + with pytest.raises(ValueError, match="API base"): + MilvusRAGIngestion( + ingest_options={ + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "c", + } + } + ) + + +def test_config_defaults(): + ingestion = _make_ingestion() + assert ingestion.collection_name == "test_collection" + assert ingestion.api_base == "http://localhost:19530" + assert ingestion.vector_field == "vector" + assert ingestion.text_field == "text" + assert ingestion.metric_type == "COSINE" + assert ingestion.auto_create_collection is True + + +def test_vector_store_id_alias(): + ingestion = _make_ingestion(collection_name=None, vector_store_id="aliased") + assert ingestion.collection_name == "aliased" + + +def test_headers_include_auth_when_api_key_set(): + ingestion = _make_ingestion() + headers = ingestion._headers() + assert headers["Authorization"] == "Bearer root:Milvus" + + +def test_headers_omit_auth_when_no_api_key(monkeypatch): + monkeypatch.delenv("MILVUS_API_KEY", raising=False) + ingestion = _make_ingestion(api_key=None) + headers = ingestion._headers() + assert "Authorization" not in headers + + +def test_server_api_key_not_sent_to_config_supplied_api_base(monkeypatch): + monkeypatch.setenv("MILVUS_API_KEY", "server-secret") + ingestion = _make_ingestion(api_base="https://attacker.example", api_key=None) + assert ingestion.api_key is None + assert "Authorization" not in ingestion._headers() + + +def test_server_api_key_used_only_with_server_api_base(monkeypatch): + monkeypatch.setenv("MILVUS_API_KEY", "server-secret") + monkeypatch.setenv("MILVUS_API_BASE", "https://milvus.internal") + ingestion = _make_ingestion(api_base=None, api_key=None) + assert ingestion.api_base == "https://milvus.internal" + assert ingestion.api_key == "server-secret" + assert ingestion._headers()["Authorization"] == "Bearer server-secret" + + +def test_config_supplied_api_key_used_with_config_api_base(monkeypatch): + monkeypatch.setenv("MILVUS_API_KEY", "server-secret") + ingestion = _make_ingestion( + api_base="https://tenant.milvus.example", api_key="tenant-token" + ) + assert ingestion.api_key == "tenant-token" + + +@pytest.mark.asyncio +async def test_store_raises_without_embeddings(): + ingestion = _make_ingestion() + with pytest.raises(ValueError, match="No text content"): + await ingestion.store( + file_content=None, + filename="doc.txt", + content_type="text/plain", + chunks=[], + embeddings=None, + ) + + +@pytest.mark.asyncio +async def test_store_auto_creates_collection_and_inserts(): + ingestion = _make_ingestion() + + post_mock = AsyncMock() + post_mock.side_effect = [ + _json_response({"code": 0, "data": {"has": False}}), # has -> not exists + _json_response({"code": 0, "data": {}}), # create + _json_response({"code": 0, "data": {"insertCount": 2}}), # insert + ] + ingestion.async_httpx_client = MagicMock() + ingestion.async_httpx_client.post = post_mock + + vector_store_id, filename = await ingestion.store( + file_content=None, + filename="doc.txt", + content_type="text/plain", + chunks=["chunk a", "chunk b"], + embeddings=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], + ) + + assert vector_store_id == "test_collection" + assert filename == "doc.txt" + assert post_mock.call_count == 3 + + called_paths = [c.args[0] for c in post_mock.call_args_list] + assert called_paths[0].endswith("/v2/vectordb/collections/has") + assert called_paths[1].endswith("/v2/vectordb/collections/create") + assert called_paths[2].endswith("/v2/vectordb/entities/insert") + + # auto-create body uses detected dimension + configured fields + create_body = post_mock.call_args_list[1].kwargs["json"] + assert create_body["collectionName"] == "test_collection" + assert create_body["dimension"] == 3 + assert create_body["vectorFieldName"] == "vector" + assert create_body["metricType"] == "COSINE" + + # insert body carries vectors, chunk text, and metadata + insert_body = post_mock.call_args_list[2].kwargs["json"] + assert insert_body["collectionName"] == "test_collection" + rows = insert_body["data"] + assert len(rows) == 2 + assert rows[0]["vector"] == [0.1, 0.2, 0.3] + assert rows[0]["text"] == "chunk a" + assert rows[0]["chunk_index"] == 0 + assert rows[0]["filename"] == "doc.txt" + + +@pytest.mark.asyncio +async def test_store_skips_create_when_collection_exists(): + ingestion = _make_ingestion() + + post_mock = AsyncMock() + post_mock.side_effect = [ + _json_response({"code": 0, "data": {"has": True}}), # has -> exists + _json_response({"code": 0, "data": {"insertCount": 1}}), # insert + ] + ingestion.async_httpx_client = MagicMock() + ingestion.async_httpx_client.post = post_mock + + await ingestion.store( + file_content=None, + filename=None, + content_type=None, + chunks=["only chunk"], + embeddings=[[1.0, 2.0]], + ) + + called_paths = [c.args[0] for c in post_mock.call_args_list] + assert not any(p.endswith("/collections/create") for p in called_paths) + assert called_paths[-1].endswith("/v2/vectordb/entities/insert") + + +@pytest.mark.asyncio +async def test_store_respects_auto_create_false(): + ingestion = _make_ingestion(auto_create_collection=False) + + post_mock = AsyncMock() + post_mock.return_value = _json_response({"code": 0, "data": {"insertCount": 1}}) + ingestion.async_httpx_client = MagicMock() + ingestion.async_httpx_client.post = post_mock + + await ingestion.store( + file_content=None, + filename=None, + content_type=None, + chunks=["c"], + embeddings=[[1.0]], + ) + + # only the insert call - no has/create probing + assert post_mock.call_count == 1 + assert post_mock.call_args_list[0].args[0].endswith("/v2/vectordb/entities/insert") + + +@pytest.mark.asyncio +async def test_post_raises_on_milvus_error_code(): + ingestion = _make_ingestion(auto_create_collection=False) + + post_mock = AsyncMock() + post_mock.return_value = _json_response( + {"code": 1100, "message": "collection not found"} + ) + ingestion.async_httpx_client = MagicMock() + ingestion.async_httpx_client.post = post_mock + + with pytest.raises(RuntimeError, match="collection not found"): + await ingestion.store( + file_content=None, + filename=None, + content_type=None, + chunks=["c"], + embeddings=[[1.0]], + ) + + +@pytest.mark.asyncio +async def test_db_name_and_partition_from_env_propagated(monkeypatch): + monkeypatch.setenv("MILVUS_DB_NAME", "mydb") + monkeypatch.setenv("MILVUS_PARTITION_NAME", "p1") + ingestion = _make_ingestion(auto_create_collection=False) + + post_mock = AsyncMock() + post_mock.return_value = _json_response({"code": 0, "data": {"insertCount": 1}}) + ingestion.async_httpx_client = MagicMock() + ingestion.async_httpx_client.post = post_mock + + await ingestion.store( + file_content=None, + filename=None, + content_type=None, + chunks=["c"], + embeddings=[[1.0]], + ) + + insert_body = post_mock.call_args_list[0].kwargs["json"] + assert insert_body["dbName"] == "mydb" + assert insert_body["partitionName"] == "p1" + + +@pytest.mark.asyncio +async def test_partition_name_from_request_is_ignored(monkeypatch): + monkeypatch.delenv("MILVUS_PARTITION_NAME", raising=False) + ingestion = _make_ingestion( + auto_create_collection=False, partition_name="victim_partition" + ) + + assert ingestion.partition_name is None + + post_mock = AsyncMock() + post_mock.return_value = _json_response({"code": 0, "data": {"insertCount": 1}}) + ingestion.async_httpx_client = MagicMock() + ingestion.async_httpx_client.post = post_mock + + await ingestion.store( + file_content=None, + filename=None, + content_type=None, + chunks=["c"], + embeddings=[[1.0]], + ) + + insert_body = post_mock.call_args_list[0].kwargs["json"] + assert "partitionName" not in insert_body + + +@pytest.mark.asyncio +async def test_db_name_from_request_is_ignored(monkeypatch): + monkeypatch.delenv("MILVUS_DB_NAME", raising=False) + ingestion = _make_ingestion(auto_create_collection=False, db_name="victim_db") + + assert ingestion.db_name is None + + post_mock = AsyncMock() + post_mock.return_value = _json_response({"code": 0, "data": {"insertCount": 1}}) + ingestion.async_httpx_client = MagicMock() + ingestion.async_httpx_client.post = post_mock + + await ingestion.store( + file_content=None, + filename=None, + content_type=None, + chunks=["c"], + embeddings=[[1.0]], + ) + + insert_body = post_mock.call_args_list[0].kwargs["json"] + assert "dbName" not in insert_body + + +@pytest.mark.asyncio +async def test_embed_returns_none_for_empty_chunks(): + ingestion = _make_ingestion() + assert await ingestion.embed([]) is None + + +@pytest.mark.asyncio +async def test_embed_uses_litellm_aembedding(monkeypatch): + ingestion = _make_ingestion() + captured = {} + + async def fake_aembedding(model, input): + captured["model"] = model + captured["input"] = input + resp = MagicMock() + resp.data = [{"embedding": [0.1, 0.2]} for _ in input] + return resp + + monkeypatch.setattr( + "litellm.rag.ingestion.base_ingestion.litellm.aembedding", + fake_aembedding, + ) + + result = await ingestion.embed(["a", "b"]) + assert result == [[0.1, 0.2], [0.1, 0.2]] + assert captured["model"] == "text-embedding-3-small" + assert captured["input"] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_embed_uses_router_when_present(): + router = MagicMock() + resp = MagicMock() + resp.data = [{"embedding": [1.0]}] + router.aembedding = AsyncMock(return_value=resp) + ingestion = MilvusRAGIngestion( + ingest_options={ + "embedding": {"model": "custom-embed"}, + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": "c", + "api_base": "http://localhost:19530", + }, + }, + router=router, + ) + + result = await ingestion.embed(["x"]) + assert result == [[1.0]] + router.aembedding.assert_awaited_once_with(model="custom-embed", input=["x"]) + + +@pytest.mark.asyncio +async def test_collection_exists_false_on_error(): + ingestion = _make_ingestion() + post_mock = AsyncMock(side_effect=RuntimeError("boom")) + ingestion.async_httpx_client = MagicMock() + ingestion.async_httpx_client.post = post_mock + assert await ingestion._collection_exists() is False + + +def test_can_auto_create_vector_store_default_true(): + assert ( + MilvusRAGIngestion.can_auto_create_vector_store( + {"custom_llm_provider": "milvus", "collection_name": "c"} + ) + is True + ) + + +def test_can_auto_create_vector_store_ignores_request_supplied_disabled_flag(): + assert ( + MilvusRAGIngestion.can_auto_create_vector_store( + { + "custom_llm_provider": "milvus", + "collection_name": "c", + "auto_create_collection": False, + } + ) + is True + ) + + +def test_credential_hydration_cannot_override_authorized_target(monkeypatch): + import litellm + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="milvus_cred", + credential_info={}, + credential_values={ + "collection_name": "victim_collection", + "vector_store_id": "victim_collection", + "api_base": "http://localhost:19530", + "api_key": "hydrated:key", + }, + ) + ], + ) + + ingestion = _make_ingestion( + collection_name="authorized_collection", + api_key="request:key", + litellm_credential_name="milvus_cred", + ) + + assert ingestion.collection_name == "authorized_collection" + assert ingestion.vector_store_config["collection_name"] == "authorized_collection" + assert "victim_collection" not in ingestion.vector_store_config.values() + assert ingestion.api_key == "hydrated:key" diff --git a/tests/vector_store_tests/rag/test_rag_milvus.py b/tests/vector_store_tests/rag/test_rag_milvus.py new file mode 100644 index 00000000000..783cd48f01f --- /dev/null +++ b/tests/vector_store_tests/rag/test_rag_milvus.py @@ -0,0 +1,86 @@ +""" +Milvus RAG ingestion tests. + +Requires environment variables: +- MILVUS_API_BASE (e.g. http://localhost:19530) + +Optional: +- MILVUS_API_KEY (token, e.g. "root:Milvus"); omit for a Milvus without auth +- MILVUS_COLLECTION_NAME (default: litellm_rag_test) + +These tests are skipped unless MILVUS_API_BASE is set. +""" + +import os +import sys +from typing import Any, Dict, Optional + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm.types.rag import RAGIngestOptions +from tests.vector_store_tests.rag.base_rag_tests import BaseRAGTest + +COLLECTION_NAME = os.environ.get("MILVUS_COLLECTION_NAME", "litellm_rag_test") + + +class TestRAGMilvus(BaseRAGTest): + """Test RAG Ingest with self-hosted Milvus.""" + + @pytest.fixture(autouse=True) + def check_env_vars(self): + if not os.environ.get("MILVUS_API_BASE"): + pytest.skip("Skipping Milvus test: MILVUS_API_BASE required") + + def get_base_ingest_options(self) -> RAGIngestOptions: + return { + "chunking_strategy": {"chunk_size": 512, "chunk_overlap": 100}, + "embedding": {"model": "text-embedding-3-small"}, + "vector_store": { + "custom_llm_provider": "milvus", + "collection_name": COLLECTION_NAME, + "api_base": os.environ["MILVUS_API_BASE"], + "api_key": os.environ.get("MILVUS_API_KEY"), + "vector_field": "vector", + "text_field": "text", + "metric_type": "COSINE", + }, + } + + async def query_vector_store( + self, + vector_store_id: str, + query: str, + ) -> Optional[Dict[str, Any]]: + """Vector-search the Milvus collection via the REST API to verify ingestion.""" + embedding_response = await litellm.aembedding( + model="text-embedding-3-small", input=[query] + ) + query_vector = embedding_response.data[0]["embedding"] + + headers = {"Content-Type": "application/json"} + api_key = os.environ.get("MILVUS_API_KEY") + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + url = f"{os.environ['MILVUS_API_BASE'].rstrip('/')}/v2/vectordb/entities/search" + body = { + "collectionName": vector_store_id, + "data": [query_vector], + "annsField": "vector", + "limit": 5, + "outputFields": ["text", "filename", "chunk_index"], + } + + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post(url, json=body, headers=headers) + response.raise_for_status() + data = response.json() + + if data.get("code") not in (0, None): + return None + results = data.get("data") or [] + return data if results else None