scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

feat(rag): add Milvus vector store ingestion support (#30388)

Browse source 
* 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
This commit is contained in:
Nguyễn Anh Bình 2026-06-23 20:25:21 +07:00 • committed by Sameer Kankute
parent de7d22283e
commit 49fd7a6fb2
No known key found for this signature in database
11 changed files with 1355 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit-misc.yml vendored
View file

@ -34,6 +34,7 @@ jobs:
tests/test_litellm/interactions
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag
tests/test_litellm/sandbox
tests/test_litellm/vector_stores
tests/test_litellm/test_*.py

5
litellm/constants.py
View file

@ -1589,6 +1589,11 @@ S3_VECTORS_DEFAULT_DIMENSION = int(os.getenv("S3_VECTORS_DEFAULT_DIMENSION", 102
S3_VECTORS_DEFAULT_DISTANCE_METRIC = str(os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine"))
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"))
MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "displayName"))

107
litellm/proxy/rag_endpoints/endpoints.py
View file

@ -28,6 +28,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store_id,
)
from litellm.repositories.table_repositories import ManagedVectorStoresRepository
from litellm.rag.main import get_ingestion_class
router = APIRouter()
@ -92,6 +93,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,
@ -103,6 +130,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,
@ -390,6 +479,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
@ -455,17 +546,11 @@ async def rag_ingest(
# Parse request
ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(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,

42
litellm/rag/ingestion/base_ingestion.py
View file

@ -83,7 +83,10 @@ class BaseRAGIngestion(ABC):
credential_values = CredentialAccessor.get_credential_values(credential_name)
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",
@ -98,6 +101,45 @@ class BaseRAGIngestion(ABC):
"""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,

269
litellm/rag/ingestion/milvus_ingestion.py Normal file
View file

@ -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": <non-zero>, "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

2
litellm/rag/main.py
View file

@ -30,6 +30,7 @@ import litellm
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 @@ INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = {
"gemini": GeminiRAGIngestion,
"s3_vectors": S3VectorsRAGIngestion,
"vertex_ai": VertexAIRAGIngestion,
"milvus": MilvusRAGIngestion,
}

36
litellm/types/rag.py
View file

@ -173,12 +173,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,
]

393
tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py
View file

@ -8,9 +8,10 @@ Covers:
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, (
f"internal_user_viewer with vector_store_id should pass role check. "
f"Response: {response.json()}"
"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, (
"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 @@ class TestRagIngestSSRFBlocked:
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 @@ class TestRagIngestSSRFBlocked:
},
},
)
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
)

0
tests/test_litellm/rag/__init__.py Normal file
View file

436
tests/test_litellm/rag/test_milvus_ingestion.py Normal file
View file

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

86
tests/vector_store_tests/rag/test_rag_milvus.py Normal file
View file

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