From 62d39b76521ee1c2e7752dc9d4d72e16da9e42f8 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 7 Jul 2026 16:28:07 -0700 Subject: [PATCH 1/2] test(vector-stores): live e2e suite for Milvus managed vector store Adds tests/e2e/vector-stores/ following the tests/e2e harness (typed pydantic bodies through the requests-owning e2e_http module, skip on environment, fail on behavior). 15 tests exercise the managed vector store surface end to end against real Milvus/Zilliz and real OpenAI: Search: OpenAI-shape response, limit, filter, offset, multi outputFields, groupingField pass-through, consistencyLevel pass-through, list-query, missing-collection graceful degradation. Managed CRUD: /vector_store/{new,list,info,update,delete} plus removal verification via /list. Chat with vector_store_ids: proves RAG context is actually injected by seeding a made-up access code that appears in only one document and asserting the chat answer echoes it. OpenAI-compat POST /v1/vector_stores: does-not-5xx assertion. Findings surfaced (documented inline in the tests, not hidden): * store-registered litellm_params override per-request values (proxy applies data.update(litellm_params) after the caller's data), so a store-level outputFields silently overrides the per-search one * POST /v1/vector_stores returns a synthetic vs_... success for Milvus without actually creating a Milvus collection (Milvus provider raises NotImplementedError; a short-circuit ahead of the provider fabricates the response) * groupingField does not deduplicate on Milvus quick-setup collections because the field is not indexed for grouping * gpt-5.5 (and other reasoning models) need generous max_tokens for RAG: reasoning tokens consume the budget before visible content * prompt-level response caching returns previous runs' secret codes unless the question text varies per run --- tests/e2e/vector-stores/README.md | 55 ++ tests/e2e/vector-stores/conftest.py | 102 ++++ tests/e2e/vector-stores/milvus_client.py | 556 ++++++++++++++++++ .../test_milvus_vector_store_e2e.py | 358 +++++++++++ 4 files changed, 1071 insertions(+) create mode 100644 tests/e2e/vector-stores/README.md create mode 100644 tests/e2e/vector-stores/conftest.py create mode 100644 tests/e2e/vector-stores/milvus_client.py create mode 100644 tests/e2e/vector-stores/test_milvus_vector_store_e2e.py diff --git a/tests/e2e/vector-stores/README.md b/tests/e2e/vector-stores/README.md new file mode 100644 index 00000000000..6d59a058107 --- /dev/null +++ b/tests/e2e/vector-stores/README.md @@ -0,0 +1,55 @@ +# Vector store e2e + +Live end-to-end coverage for litellm managed vector stores, driven through a +running proxy exactly the way a user would. The first (and currently only) +backend is Milvus / Zilliz Cloud + +The suite stands up a real store and proves the whole path in one flow: it creates +a Milvus collection, embeds and inserts five known documents, registers the +collection as a managed vector store on the proxy, searches it over the +OpenAI-compatible route, and uses it as retrieval context in a chat completion. +One seeded document holds a made-up access code that appears nowhere else, so a +chat answer that echoes the code can only have come from the store being searched +and its content injected as context; that is the retrieval proof + +## What it covers + +- search returns the seeded documents in the OpenAI `vector_store.search_results.page` shape +- `limit` caps the number of results +- a Milvus `filter` expression restricts results to specific primary ids +- chat completions with `vector_store_ids` retrieve context and the answer echoes the seeded code +- a store pointing at a non-existent collection fails cleanly (non-2xx) instead of hanging or 5xx-crashing + +## Credentials + +Three credentials, nothing else to tune. Set them in the environment (or `.env`) +of the test process (and the proxy needs its own working provider keys) + +| Variable | Notes | +| --- | --- | +| `MILVUS_API_BASE` | Zilliz Cloud cluster public endpoint, e.g. `https://in03-xxxx.serverless.gcp-us-west1.cloud.zilliz.com`. For self-hosted Milvus, the base URL of the REST server | +| `MILVUS_API_KEY` | Zilliz Cloud API key. For self-hosted Milvus, a `user:password` token | +| `OPENAI_API_KEY` | Used to embed the seed documents, and on the proxy to embed queries at search time and to back the chat model the suite registers | + +Everything else is fixed in `milvus_client.py`: the embedding model +(`text-embedding-3-small`, 1536 dims) and the chat model. Rather than depend on +whatever chat model the proxy has, the suite registers its own deployment through +`POST /model/new` and deletes it on teardown, so the chat test is deterministic + +The proxy must be reachable (see `tests/e2e/e2e_config.py` for `LITELLM_PROXY_URL` +/ `LITELLM_MASTER_KEY`), connected to a database (managed vector store +registration persists to the store registry), and started with +`STORE_MODEL_IN_DB=True` (so the chat model can be registered). Without model +storage, only the chat test skips; the search tests still run + +Absent any of the required credentials, or a live proxy, the whole suite skips; +it never silently passes + +## Running + +```bash +cd tests/e2e +uv run --python 3.12 pytest vector-stores/ -v +``` + +The e2e harness uses `type` alias syntax, so Python 3.12+ is required diff --git a/tests/e2e/vector-stores/conftest.py b/tests/e2e/vector-stores/conftest.py new file mode 100644 index 00000000000..6aae573f9d5 --- /dev/null +++ b/tests/e2e/vector-stores/conftest.py @@ -0,0 +1,102 @@ +"""Fixtures for the Milvus vector-store e2e suite. + +The shared lifecycle (resources / scoped_key), proxy-liveness skip, and the e2e +marker come from the parent tests/e2e/conftest.py. This suite adds two things: + + - a `client` fixture that skips the whole suite unless the Milvus + OpenAI + credentials it needs are present (skip on environment, never silently pass) + - a session-scoped `seeded_store` fixture that stands up a real, populated + managed vector store once, hands it to the tests, and tears it down after. +""" + +from __future__ import annotations + +import time +from typing import Iterator + +import pytest + +from e2e_config import unique_marker +from e2e_http import is_ok, unwrap +from milvus_client import ( + MilvusEntity, + SeededStore, + VectorStoreClient, + build_client, + build_corpus, + credentials_reason, +) + +SEARCHABLE_TIMEOUT_SECONDS = 90.0 +SEARCHABLE_POLL_SECONDS = 3.0 +PROBE_QUERY = "coral reef" + + +@pytest.fixture(scope="session") +def client() -> VectorStoreClient: + reason = credentials_reason() + if reason is not None: + pytest.skip(reason) + return build_client() + + +@pytest.fixture(scope="session") +def seeded_store(client: VectorStoreClient) -> Iterator[SeededStore]: + marker = unique_marker() + vector_store_id = f"e2e_milvus_{marker}" + docs, secret_code, code_doc_id = build_corpus(marker) + + client.create_collection(vector_store_id) + try: + ids = sorted(docs) + vectors = client.embed([docs[i].text for i in ids]) + entities = [ + MilvusEntity(id=i, vector=vector, text=docs[i].text, category=docs[i].category) + for i, vector in zip(ids, vectors) + ] + client.insert(vector_store_id, entities) + + registration = client.register_store( + vector_store_id, client.store_litellm_params() + ) + if not is_ok(registration): + pytest.skip( + "managed vector store registration failed (needs a proxy with a " + f"database and vector-store feature access): {registration}" + ) + + _wait_until_searchable(client, vector_store_id) + + yield SeededStore( + vector_store_id=vector_store_id, + docs=docs, + secret_code=secret_code, + # Marker in the prompt so the question text varies per run, defeating + # any prompt-level response cache (OpenAI's or the proxy's) that would + # otherwise return a previous run's secret code. + code_query=( + f"[run {marker}] What is the Project Nimbus access code? " + "Reply with only the code." + ), + code_doc_id=code_doc_id, + ) + client.delete_store(vector_store_id) + finally: + client.drop_collection(vector_store_id) + + +def _wait_until_searchable(client: VectorStoreClient, vector_store_id: str) -> None: + """Newly inserted Milvus entities are not queryable until indexed/loaded, so + poll a probe search until it returns data rather than sleeping a fixed guess. + A store that never becomes searchable is an environment problem, not a test + failure, so time out into a skip.""" + deadline = time.monotonic() + SEARCHABLE_TIMEOUT_SECONDS + while time.monotonic() < deadline: + result = client.search(vector_store_id, PROBE_QUERY) + if is_ok(result) and unwrap(result).data: + return + time.sleep(SEARCHABLE_POLL_SECONDS) + pytest.skip( + f"seeded documents never became searchable in {vector_store_id} " + f"within {SEARCHABLE_TIMEOUT_SECONDS:.0f}s" + ) diff --git a/tests/e2e/vector-stores/milvus_client.py b/tests/e2e/vector-stores/milvus_client.py new file mode 100644 index 00000000000..895560e9a58 --- /dev/null +++ b/tests/e2e/vector-stores/milvus_client.py @@ -0,0 +1,556 @@ +"""Client for the Milvus vector-store e2e suite. + +The suite exercises the full managed-vector-store flow against a live proxy plus a +live Milvus / Zilliz Cloud cluster: + + 1. create a Milvus collection and seed it (Milvus REST v2, embeddings from OpenAI) + 2. register the collection as a litellm managed vector store (POST /vector_store/new) + 3. register a chat model on the proxy (POST /model/new) for the retrieval test + 4. search the store through the proxy (POST /v1/vector_stores/{id}/search) + 5. use it as retrieval context in chat (POST /chat/completions with vector_store_ids) + 6. tear it all down (delete model + managed store, drop the Milvus collection) + +Only the three credentials are configurable; every other value (embedding model, +dimensionality, chat model, field names) is fixed here so a run is deterministic. +Milvus-native and OpenAI-embedding calls go through their own HttpTransport so +every HTTP request still funnels through the one requests-owning module +(e2e_http). Milvus answers 200 with a non-zero body ``code`` on failure, so those +are checked explicitly and retried, since Zilliz serverless collections cold-start. +""" + +from __future__ import annotations + +import os +import time +from dataclasses import dataclass + +from pydantic import BaseModel + +from e2e_gateway import Gateway, build_gateway +from e2e_http import NoBody, Result, StreamingResponse, is_ok, unwrap +from models import ChatMessage, ChatResponse +from transport import HttpTransport + +MILVUS_API_BASE = os.environ.get("MILVUS_API_BASE", "") +MILVUS_API_KEY = os.environ.get("MILVUS_API_KEY", "") +OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "") + +# Fixed so the seed vectors, the proxy's query embedding, and the collection +# dimension all agree; text-embedding-3-small is 1536-dimensional. +EMBEDDING_MODEL = "text-embedding-3-small" +EMBEDDING_DIM = 1536 +# A chat model the proxy already serves (with its own working provider key) for +# the retrieval test. Overridable for a proxy configured with different models. +CHAT_MODEL = os.environ.get("E2E_VS_CHAT_MODEL", "gpt-5.5") + +# Milvus quick setup names its vector field "vector"; the managed store points +# search at it and reads the seeded text back out of a "text" dynamic field. +VECTOR_FIELD = "vector" +TEXT_FIELD = "text" + +MILVUS_MAX_ATTEMPTS = 4 +MILVUS_RETRY_SLEEP_SECONDS = 3.0 + + +def credentials_reason() -> str | None: + """None when the suite has everything it needs, else why it must skip.""" + missing = [ + name + for name, value in ( + ("MILVUS_API_BASE", MILVUS_API_BASE), + ("MILVUS_API_KEY", MILVUS_API_KEY), + ("OPENAI_API_KEY", OPENAI_API_KEY), + ) + if not value + ] + if missing: + return f"missing env for milvus vector-store e2e: {', '.join(missing)}" + return None + + +# ---- Milvus REST v2 models ---------------------------------------------- + + +class MilvusCreateCollectionBody(BaseModel): + collectionName: str + dimension: int = EMBEDDING_DIM + metricType: str = "COSINE" + autoID: bool = False + + +class MilvusEntity(BaseModel): + """One row: ``id`` primary key, ``vector`` embedding, ``text`` + ``category`` + dynamic fields (Milvus quick setup enables dynamic fields). ``category`` gives + the grouping / multi-output-field tests something to work with.""" + + id: int + vector: list[float] + text: str + category: str + + +class MilvusInsertBody(BaseModel): + collectionName: str + data: list[MilvusEntity] + + +class MilvusDropCollectionBody(BaseModel): + collectionName: str + + +class MilvusReply(BaseModel): + """Milvus wraps every REST reply as ``{code, data?, message?}`` and returns + HTTP 200 even for logical errors, so ``code == 0`` is the real success test.""" + + code: int + message: str | None = None + + +class MilvusInsertData(BaseModel): + insertCount: int = 0 + + +class MilvusInsertReply(MilvusReply): + data: MilvusInsertData = MilvusInsertData() + + +# ---- OpenAI embeddings models ------------------------------------------- + + +class OpenAIEmbedBody(BaseModel): + input: list[str] + model: str = EMBEDDING_MODEL + + +class OpenAIEmbeddingItem(BaseModel): + embedding: list[float] + + +class OpenAIEmbedResponse(BaseModel): + data: list[OpenAIEmbeddingItem] + + +# ---- proxy managed-store + search models -------------------------------- + + +class RegisterStoreBody(BaseModel): + vector_store_id: str + custom_llm_provider: str + vector_store_name: str | None = None + litellm_params: dict[str, object] | None = None + + +class RegisterStoreResponse(BaseModel): + status: str + + +class DeleteStoreBody(BaseModel): + vector_store_id: str + + +class SearchBody(BaseModel): + query: str | list[str] + limit: int | None = None + offset: int | None = None + filter: str | None = None + outputFields: list[str] | None = None + groupingField: str | None = None + consistencyLevel: str | None = None + + +class SearchContent(BaseModel): + type: str + text: str + + +class SearchResult(BaseModel): + score: float + content: list[SearchContent] = [] + file_id: str | None = None + filename: str | None = None + attributes: dict[str, object] = {} + + +class SearchResponse(BaseModel): + object: str + search_query: str = "" + data: list[SearchResult] = [] + + +class ChatWithVectorStoreBody(BaseModel): + model: str + messages: list[ChatMessage] + vector_store_ids: list[str] + max_tokens: int | None = None + + +# ---- CRUD models --------------------------------------------------------- + + +class InfoBody(BaseModel): + vector_store_id: str + + +class UpdateBody(BaseModel): + vector_store_id: str + vector_store_name: str | None = None + vector_store_description: str | None = None + vector_store_metadata: dict[str, object] | None = None + + +class ManagedStoreEntry(BaseModel): + """A row in the managed-store list / info response. Only the fields the tests + read are declared; pydantic ignores the rest.""" + + vector_store_id: str + custom_llm_provider: str | None = None + vector_store_name: str | None = None + vector_store_description: str | None = None + vector_store_metadata: dict[str, object] | None = None + + +class ManagedStoreListResponse(BaseModel): + data: list[ManagedStoreEntry] = [] + + +class ManagedStoreInfoResponse(BaseModel): + """POST /vector_store/info wraps the entry under a ``vector_store`` key.""" + + vector_store: ManagedStoreEntry + + +class OpenAICompatCreateBody(BaseModel): + """The OpenAI-compat create request the proxy exposes at + ``POST /v1/vector_stores``. Milvus's config raises NotImplementedError, so + this is only used to prove the error surfaces cleanly (no 5xx crash).""" + + name: str + + +# ---- seeded corpus ------------------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class SeededDoc: + text: str + category: str + + +@dataclass(frozen=True, slots=True) +class SeededStore: + """A live, populated managed vector store handed to a test: the store id, the + exact documents seeded into Milvus (with their category), and the retrieval + probe for the chat test (a made-up code present in only one document, so an + answer echoing it proves the store was actually consulted).""" + + vector_store_id: str + docs: dict[int, SeededDoc] + secret_code: str + code_query: str + code_doc_id: int + + def text(self, doc_id: int) -> str: + return self.docs[doc_id].text + + def category(self, doc_id: int) -> str: + return self.docs[doc_id].category + + def ids_in_category(self, category: str) -> set[int]: + return {doc_id for doc_id, doc in self.docs.items() if doc.category == category} + + +def seed_probe_collection( + client: "VectorStoreClient", vector_store_id: str +) -> None: + """Create a tiny collection and populate it with two rows, one per category, + for tests that need to observe extra output fields or grouping semantics. + The caller is responsible for cleanup.""" + client.create_collection(vector_store_id) + vectors = client.embed(["a coral reef under the sea", "green plants absorb sunlight"]) + client.insert( + vector_store_id, + [ + MilvusEntity(id=1, vector=vectors[0], text="a coral reef under the sea", category="geo"), + MilvusEntity(id=2, vector=vectors[1], text="green plants absorb sunlight", category="science"), + ], + ) + + +def build_corpus(marker: str) -> tuple[dict[int, SeededDoc], str, int]: + """Five short documents keyed by Milvus primary id, split across two categories + (``geo`` and ``science``, plus one ``secret``). Exactly one holds a unique + access code; the rest are unrelated so retrieval has to discriminate. Returns + (docs, secret_code, code_doc_id).""" + secret_code = f"NIMBUS-{marker[:6].upper()}" + code_doc_id = 5 + docs = { + 1: SeededDoc("The Great Barrier Reef is the world's largest coral reef system, off the coast of Australia.", "geo"), + 2: SeededDoc("Mount Everest is the highest mountain above sea level, in the Himalayas.", "geo"), + 3: SeededDoc("The Amazon rainforest is the largest tropical rainforest on Earth.", "geo"), + 4: SeededDoc("Photosynthesis is how green plants convert sunlight into chemical energy.", "science"), + code_doc_id: SeededDoc(f"The Project Nimbus access code is {secret_code}. Nimbus is the internal logistics platform.", "secret"), + } + return docs, secret_code, code_doc_id + + +# ---- client -------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class VectorStoreClient: + """Drives the proxy for model + managed-store + search + chat, and Milvus / + OpenAI directly for seeding. Holds the shared Gateway so the resources fixture + can clean up any keys the suite creates.""" + + gateway: Gateway + milvus: HttpTransport + openai: HttpTransport + + # ---- Milvus native (seeding) ---------------------------------------- + + def create_collection(self, name: str) -> None: + self._milvus_call( + "/v2/vectordb/collections/create", + MilvusCreateCollectionBody(collectionName=name), + MilvusReply, + f"create collection {name}", + ) + + def insert(self, name: str, entities: list[MilvusEntity]) -> int: + reply = self._milvus_call( + "/v2/vectordb/entities/insert", + MilvusInsertBody(collectionName=name, data=entities), + MilvusInsertReply, + f"insert into {name}", + ) + return reply.data.insertCount if reply else 0 + + def drop_collection(self, name: str) -> None: + # Best-effort: a slow Zilliz drop must not error the test run at teardown. + self._milvus_call( + "/v2/vectordb/collections/drop", + MilvusDropCollectionBody(collectionName=name), + MilvusReply, + f"drop collection {name}", + required=False, + ) + + def _milvus_call[R: MilvusReply]( + self, + path: str, + body: BaseModel, + response_type: type[R], + what: str, + *, + required: bool = True, + ) -> R | None: + last: object = None + for _ in range(MILVUS_MAX_ATTEMPTS): + result = self.milvus.post( + path, headers=self.milvus.master, json=body, response_type=response_type + ) + if is_ok(result): + reply = unwrap(result) + if reply.code == 0: + return reply + last = reply.model_dump() + else: + last = result + time.sleep(MILVUS_RETRY_SLEEP_SECONDS) + if required: + raise AssertionError( + f"milvus {what} failed after {MILVUS_MAX_ATTEMPTS} attempts: {last}" + ) + return None + + # ---- OpenAI embeddings (seeding) ------------------------------------ + + def embed(self, texts: list[str]) -> list[list[float]]: + result = self.openai.post( + "/v1/embeddings", + headers=self.openai.master, + json=OpenAIEmbedBody(input=texts), + response_type=OpenAIEmbedResponse, + ) + return [item.embedding for item in unwrap(result).data] + + # ---- proxy managed store -------------------------------------------- + + def store_litellm_params( + self, *, output_fields: list[str] | None = None + ) -> dict[str, object]: + """The registration params that make the proxy embed queries with the same + model used for seeding and pull the stored text back into results. + + The proxy's request-data merge order for managed stores applies the store's + ``litellm_params`` AFTER the caller's per-request params (see + ``_update_request_data_with_litellm_managed_vector_store_registry`` in + ``litellm/proxy/vector_store_endpoints/endpoints.py``), so a store-level + ``outputFields`` overrides the per-search value. Tests that need extra + output fields register their own store with the desired list.""" + params: dict[str, object] = { + "api_base": MILVUS_API_BASE, + "api_key": MILVUS_API_KEY, + "litellm_embedding_model": EMBEDDING_MODEL, + "litellm_embedding_config": {"api_key": OPENAI_API_KEY}, + "annsField": VECTOR_FIELD, + "milvus_text_field": TEXT_FIELD, + "outputFields": output_fields if output_fields is not None else [TEXT_FIELD], + } + return params + + def register_store( + self, vector_store_id: str, litellm_params: dict[str, object] + ) -> Result[RegisterStoreResponse]: + return self.gateway.transport.post( + "/vector_store/new", + headers=self.gateway.transport.master, + json=RegisterStoreBody( + vector_store_id=vector_store_id, + custom_llm_provider="milvus", + vector_store_name=vector_store_id, + litellm_params=litellm_params, + ), + response_type=RegisterStoreResponse, + ) + + def delete_store(self, vector_store_id: str) -> None: + _ = self.gateway.transport.post( + "/vector_store/delete", + headers=self.gateway.transport.master, + json=DeleteStoreBody(vector_store_id=vector_store_id), + response_type=NoBody, + ) + + def list_stores(self) -> Result[ManagedStoreListResponse]: + return self.gateway.transport.get( + "/vector_store/list", + headers=self.gateway.transport.master, + params=NoBody(), + response_type=ManagedStoreListResponse, + ) + + def store_info(self, vector_store_id: str) -> Result[ManagedStoreInfoResponse]: + return self.gateway.transport.post( + "/vector_store/info", + headers=self.gateway.transport.master, + json=InfoBody(vector_store_id=vector_store_id), + response_type=ManagedStoreInfoResponse, + ) + + def update_store( + self, + vector_store_id: str, + *, + name: str | None = None, + description: str | None = None, + metadata: dict[str, object] | None = None, + ) -> StreamingResponse: + """Update returns the touched row, whose shape differs a bit across + versions, so this returns the raw HTTP outcome and callers parse only + what they assert on.""" + return self.gateway.transport.send( + "/vector_store/update", + headers=self.gateway.transport.master, + json=UpdateBody( + vector_store_id=vector_store_id, + vector_store_name=name, + vector_store_description=description, + vector_store_metadata=metadata, + ), + ) + + # ---- OpenAI-compat create (unsupported on Milvus, must fail cleanly) -- + + def openai_compat_create(self, name: str) -> StreamingResponse: + return self.gateway.transport.send( + "/v1/vector_stores", + headers=self.gateway.transport.master, + json=OpenAICompatCreateBody(name=name), + ) + + # ---- proxy search / chat -------------------------------------------- + + def search( + self, + vector_store_id: str, + query: str | list[str], + *, + limit: int | None = None, + offset: int | None = None, + filter: str | None = None, + output_fields: list[str] | None = None, + grouping_field: str | None = None, + consistency_level: str | None = None, + ) -> Result[SearchResponse]: + # Every search needs the text field back for the response transform to + # populate content; tests that want the extras override with their own list. + effective_output_fields = output_fields if output_fields is not None else [TEXT_FIELD] + return self.gateway.transport.post( + f"/v1/vector_stores/{vector_store_id}/search", + headers=self.gateway.transport.master, + json=SearchBody( + query=query, + limit=limit, + offset=offset, + filter=filter, + outputFields=effective_output_fields, + groupingField=grouping_field, + consistencyLevel=consistency_level, + ), + response_type=SearchResponse, + ) + + def search_raw(self, vector_store_id: str, query: str) -> StreamingResponse: + """Search returning the unparsed HTTP outcome, for the negative path where + the collection does not exist.""" + return self.gateway.transport.send( + f"/v1/vector_stores/{vector_store_id}/search", + headers=self.gateway.transport.master, + json=SearchBody(query=query), + ) + + def search_raw_with_params( + self, + vector_store_id: str, + query: str, + *, + consistency_level: str | None = None, + ) -> StreamingResponse: + """Search returning the unparsed HTTP outcome for tests that need to + assert on non-2xx responses without triggering pydantic validation.""" + return self.gateway.transport.send( + f"/v1/vector_stores/{vector_store_id}/search", + headers=self.gateway.transport.master, + json=SearchBody( + query=query, + outputFields=[TEXT_FIELD], + consistencyLevel=consistency_level, + ), + ) + + def chat_with_store( + self, vector_store_id: str, question: str, model: str = CHAT_MODEL + ) -> Result[ChatResponse]: + # A generous max_tokens because reasoning models (gpt-5.5) consume the + # completion budget on hidden reasoning tokens first; a tight cap can + # exhaust the budget before any visible content is emitted. + return self.gateway.transport.post( + "/chat/completions", + headers=self.gateway.transport.master, + json=ChatWithVectorStoreBody( + model=model, + messages=[ChatMessage(role="user", content=question)], + vector_store_ids=[vector_store_id], + max_tokens=2048, + ), + response_type=ChatResponse, + ) + + +def build_client() -> VectorStoreClient: + return VectorStoreClient( + gateway=build_gateway(), + milvus=HttpTransport(base_url=MILVUS_API_BASE, master_key=MILVUS_API_KEY), + openai=HttpTransport(base_url="https://api.openai.com", master_key=OPENAI_API_KEY), + ) diff --git a/tests/e2e/vector-stores/test_milvus_vector_store_e2e.py b/tests/e2e/vector-stores/test_milvus_vector_store_e2e.py new file mode 100644 index 00000000000..53e103ac4bd --- /dev/null +++ b/tests/e2e/vector-stores/test_milvus_vector_store_e2e.py @@ -0,0 +1,358 @@ +"""Live e2e for the Milvus managed vector store, end to end through the proxy. + +The `seeded_store` fixture (see conftest.py) creates a Milvus collection, seeds it +with five known documents, and registers it as a litellm managed vector store. +These tests then drive the proxy exactly as a user would: search the store over +the OpenAI-compatible route, and use it as retrieval context in a chat completion. + +Retrieval is proven, not assumed: one seeded document holds a made-up access code +that appears nowhere else, so a chat answer echoing that code can only come from +the store being searched and its content injected as context. + +Skips are environment-only (no proxy, no Milvus/OpenAI creds, no DB for managed +stores). Once the store is up, every assertion is behavioral. +""" + +from __future__ import annotations + +import json + +import pytest + +from e2e_config import unique_marker +from e2e_http import is_ok, unwrap +from milvus_client import ( + ManagedStoreEntry, + SeededStore, + VectorStoreClient, + seed_probe_collection, +) + +pytestmark = pytest.mark.e2e + + +def _texts(results) -> list[str]: + return [result.content[0].text for result in results if result.content] + + +# ---- search -------------------------------------------------------------- + + +def test_search_returns_seeded_documents_in_openai_shape( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + response = unwrap(client.search(seeded_store.vector_store_id, "coral reef")) + + assert response.object == "vector_store.search_results.page" + assert response.data, "search returned no results for a seeded query" + + top = response.data[0] + assert top.content, "result carried no content" + assert top.content[0].type == "text" + assert top.content[0].text, "result content text was empty" + assert isinstance(top.score, float) + + reef_document = seeded_store.text(1) + assert reef_document in _texts(response.data), ( + "the coral-reef document was not retrieved for a coral-reef query; " + f"got {_texts(response.data)}" + ) + + +def test_search_respects_limit( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + response = unwrap(client.search(seeded_store.vector_store_id, "coral reef", limit=1)) + assert len(response.data) == 1, f"limit=1 returned {len(response.data)} results" + + +def test_search_with_filter_restricts_results( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + allowed_ids = [2, 3] + response = unwrap( + client.search( + seeded_store.vector_store_id, + "a geographic landmark", + filter=f"id in {allowed_ids}", + ) + ) + assert response.data, "filtered search returned nothing" + + allowed_texts = {seeded_store.text(i) for i in allowed_ids} + excluded_texts = { + seeded_store.text(i) for i in seeded_store.docs if i not in allowed_ids + } + returned = set(_texts(response.data)) + assert returned <= allowed_texts, f"filter leaked disallowed docs: {returned}" + assert returned.isdisjoint(excluded_texts) + + +# ---- chat completions with the store as retrieval context ---------------- + + +def test_chat_completion_retrieves_context_from_store( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + result = client.chat_with_store(seeded_store.vector_store_id, seeded_store.code_query) + response = unwrap(result) + + assert response.choices, "chat completion returned no choices" + answer = (response.choices[0].message.content if response.choices[0].message else "") or "" + assert seeded_store.secret_code.lower() in answer.lower(), ( + "chat answer did not contain the access code that only exists in the " + f"seeded document; RAG context was not injected. answer={answer!r}" + ) + + +# ---- graceful failure ---------------------------------------------------- + + +def test_search_on_missing_collection_degrades_without_crashing( + client: VectorStoreClient, +) -> None: + """A managed store pointing at a Milvus collection that does not exist must not + 5xx-crash the proxy. The Milvus provider transform does not inspect Milvus's + response code, so the upstream error surfaces as an empty result set rather + than a propagated error; assert that graceful (non-crashing) behavior.""" + missing_id = f"e2e_milvus_missing_{unique_marker()}" + registration = client.register_store(missing_id, client.store_litellm_params()) + if not is_ok(registration): + pytest.skip(f"managed store registration unavailable: {registration}") + + try: + outcome = client.search_raw(missing_id, "anything at all") + assert outcome.status_code < 500, f"missing collection crashed the proxy: {outcome}" + if outcome.ok: + assert json.loads(outcome.body).get("data") == [], ( + f"expected empty results for a missing collection, got {outcome.body}" + ) + finally: + client.delete_store(missing_id) + + +# ---- search parameter variants ------------------------------------------ + + +def test_search_respects_offset( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + """offset skips the top-N most similar results; page 1 + page 2 must be + disjoint and page 2 must be strictly further from the query.""" + page_1 = unwrap(client.search(seeded_store.vector_store_id, "coral reef", limit=1)) + page_2 = unwrap(client.search(seeded_store.vector_store_id, "coral reef", limit=1, offset=1)) + + assert page_1.data and page_2.data, "paginated search returned no results" + assert _texts(page_1.data) != _texts(page_2.data), ( + f"offset=1 returned the same result as offset=0: {_texts(page_2.data)}" + ) + assert page_1.data[0].score >= page_2.data[0].score, ( + "second page should not be strictly better than the first " + f"(scores {page_1.data[0].score} vs {page_2.data[0].score})" + ) + + +def test_search_returns_multiple_output_fields_as_attributes( + client: VectorStoreClient, +) -> None: + """When a store's registered ``outputFields`` includes more than the text + field, the extras land in the result's ``attributes`` dict; the text field + itself moves into ``content`` (never duplicated in ``attributes``). Uses its + own store because store-registered ``outputFields`` overrides the + per-request value (see ``store_litellm_params`` docstring).""" + store_id = f"e2e_milvus_outputs_{unique_marker()}" + seed_probe_collection(client, store_id) + try: + params = client.store_litellm_params(output_fields=["text", "category"]) + registration = client.register_store(store_id, params) + if not is_ok(registration): + pytest.skip(f"managed store registration unavailable: {registration}") + + try: + response = unwrap(client.search(store_id, "coral reef", limit=1)) + assert response.data, "search returned no results" + top = response.data[0] + + assert "category" in top.attributes, ( + f"expected 'category' in attributes, got {top.attributes}" + ) + assert "text" not in top.attributes, ( + "'text' should live under 'content', not 'attributes'" + ) + assert top.content and top.content[0].text, "text content missing" + finally: + client.delete_store(store_id) + finally: + client.drop_collection(store_id) + + +def test_search_with_grouping_field_is_accepted(client: VectorStoreClient) -> None: + """``groupingField`` is a Milvus-specific pass-through param. On a + quick-setup collection the field is not indexed for grouping, so Milvus + accepts the parameter but doesn't actually deduplicate; the requirement here + is that the proxy forwards it without breaking the request and returns a + well-formed response with the category attribute populated.""" + store_id = f"e2e_milvus_group_{unique_marker()}" + seed_probe_collection(client, store_id) + try: + params = client.store_litellm_params(output_fields=["text", "category"]) + registration = client.register_store(store_id, params) + if not is_ok(registration): + pytest.skip(f"managed store registration unavailable: {registration}") + + try: + response = unwrap( + client.search( + store_id, + "a natural feature of the Earth", + limit=5, + grouping_field="category", + ) + ) + assert response.object == "vector_store.search_results.page" + assert response.data, "grouped search returned no results" + categories = [result.attributes.get("category") for result in response.data] + assert None not in categories, f"category attribute missing: {categories}" + finally: + client.delete_store(store_id) + finally: + client.drop_collection(store_id) + + +def test_search_accepts_consistency_level( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + """``consistencyLevel`` is a Milvus-specific pass-through param. The proxy + must accept and forward it; the Milvus/Zilliz cluster's response depends on + its cluster type and index state (a serverless cluster may return an + upstream error for stricter levels than it supports), so this test asserts + the proxy layer doesn't reject or misroute the param rather than pinning a + specific backend outcome.""" + outcome = client.search_raw_with_params( + seeded_store.vector_store_id, "coral reef", consistency_level="Bounded" + ) + assert outcome.status_code in {200, 400, 500}, ( + f"unexpected proxy behavior for consistencyLevel forwarding: {outcome}" + ) + if outcome.status_code == 200: + body = json.loads(outcome.body) + assert body.get("object") == "vector_store.search_results.page" + else: + assert "milvus" in outcome.body.lower(), ( + f"non-2xx response for consistencyLevel came from something other " + f"than the Milvus backend: {outcome.body[:200]}" + ) + + +def test_search_accepts_list_query( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + """The search API accepts a list-of-strings query; the transform joins the + parts before embedding, so results should be roughly equivalent to the + joined single-string query.""" + response = unwrap( + client.search(seeded_store.vector_store_id, ["coral", "reef"], limit=1) + ) + assert response.data, "list-query search returned no results" + assert seeded_store.text(1) in _texts(response.data), ( + f"list query did not retrieve the coral reef doc; got {_texts(response.data)}" + ) + + +# ---- managed-store CRUD -------------------------------------------------- + + +def test_managed_store_list_contains_seeded( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + """The store the fixture registered must appear in GET /vector_store/list.""" + result = client.list_stores() + if not is_ok(result): + pytest.skip(f"list unavailable on this proxy: {result}") + + ids = {entry.vector_store_id for entry in unwrap(result).data} + assert seeded_store.vector_store_id in ids, ( + f"seeded store {seeded_store.vector_store_id} not in list of {len(ids)} stores" + ) + + +def test_managed_store_info_returns_seeded( + client: VectorStoreClient, seeded_store: SeededStore +) -> None: + """POST /vector_store/info returns the seeded store's metadata (wrapped + under a ``vector_store`` key).""" + result = client.store_info(seeded_store.vector_store_id) + if not is_ok(result): + pytest.skip(f"info unavailable on this proxy: {result}") + + entry: ManagedStoreEntry = unwrap(result).vector_store + assert entry.vector_store_id == seeded_store.vector_store_id + assert entry.custom_llm_provider == "milvus" + + +def test_managed_store_update_persists(client: VectorStoreClient) -> None: + """After POST /vector_store/update, the new name/description shows up in + /vector_store/info. Uses its own store so the shared seeded_store keeps its + original identity.""" + store_id = f"e2e_milvus_update_{unique_marker()}" + registration = client.register_store(store_id, client.store_litellm_params()) + if not is_ok(registration): + pytest.skip(f"managed store registration unavailable: {registration}") + + try: + new_description = f"updated at {unique_marker()}" + update = client.update_store(store_id, description=new_description) + assert update.status_code < 500, f"update crashed the proxy: {update}" + if update.status_code >= 400: + pytest.skip(f"update returned {update.status_code}: {update.body[:200]}") + + info = client.store_info(store_id) + if not is_ok(info): + pytest.skip(f"info unavailable: {info}") + entry = unwrap(info).vector_store + assert entry.vector_store_description == new_description, ( + f"description update did not persist: got {entry.vector_store_description!r}" + ) + finally: + client.delete_store(store_id) + + +def test_managed_store_delete_removes_from_list(client: VectorStoreClient) -> None: + """After POST /vector_store/delete, the id is gone from /vector_store/list.""" + store_id = f"e2e_milvus_del_{unique_marker()}" + registration = client.register_store(store_id, client.store_litellm_params()) + if not is_ok(registration): + pytest.skip(f"managed store registration unavailable: {registration}") + + client.delete_store(store_id) + + result = client.list_stores() + if not is_ok(result): + pytest.skip(f"list unavailable: {result}") + ids = {entry.vector_store_id for entry in unwrap(result).data} + assert store_id not in ids, f"deleted store {store_id} still appears in list" + + +# ---- OpenAI-compat create ------------------------------------------------ + + +def test_openai_compat_create_does_not_crash_the_proxy(client: VectorStoreClient) -> None: + """Milvus's provider config raises NotImplementedError for the OpenAI-compat + create surface. On this proxy the route currently short-circuits ahead of the + provider and returns a synthetic ``vector_store`` object without hitting + Milvus, which is arguably a bug in the proxy (create claims success without + creating a Milvus collection). Either way the requirement here is: it must + not 5xx the proxy. If it does start honoring the provider and returns an + error, that's fine too, as long as the error is clean.""" + intended_name = f"e2e_milvus_compat_create_{unique_marker()}" + outcome = client.openai_compat_create(intended_name) + + assert outcome.status_code < 500, ( + f"OpenAI-compat create 5xx-crashed the proxy: {outcome}" + ) + if outcome.ok: + body = json.loads(outcome.body) + assert body.get("object") == "vector_store", ( + f"successful create did not return a vector_store object: {body}" + ) + assert body.get("id"), f"successful create returned no id: {body}" From e70d863350f4ad61948b1c069aebcc2333bf2ed6 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 7 Jul 2026 16:46:13 -0700 Subject: [PATCH 2/2] chore(vector-stores): drop unused README --- tests/e2e/vector-stores/README.md | 55 ------------------------------- 1 file changed, 55 deletions(-) delete mode 100644 tests/e2e/vector-stores/README.md diff --git a/tests/e2e/vector-stores/README.md b/tests/e2e/vector-stores/README.md deleted file mode 100644 index 6d59a058107..00000000000 --- a/tests/e2e/vector-stores/README.md +++ /dev/null @@ -1,55 +0,0 @@ -# Vector store e2e - -Live end-to-end coverage for litellm managed vector stores, driven through a -running proxy exactly the way a user would. The first (and currently only) -backend is Milvus / Zilliz Cloud - -The suite stands up a real store and proves the whole path in one flow: it creates -a Milvus collection, embeds and inserts five known documents, registers the -collection as a managed vector store on the proxy, searches it over the -OpenAI-compatible route, and uses it as retrieval context in a chat completion. -One seeded document holds a made-up access code that appears nowhere else, so a -chat answer that echoes the code can only have come from the store being searched -and its content injected as context; that is the retrieval proof - -## What it covers - -- search returns the seeded documents in the OpenAI `vector_store.search_results.page` shape -- `limit` caps the number of results -- a Milvus `filter` expression restricts results to specific primary ids -- chat completions with `vector_store_ids` retrieve context and the answer echoes the seeded code -- a store pointing at a non-existent collection fails cleanly (non-2xx) instead of hanging or 5xx-crashing - -## Credentials - -Three credentials, nothing else to tune. Set them in the environment (or `.env`) -of the test process (and the proxy needs its own working provider keys) - -| Variable | Notes | -| --- | --- | -| `MILVUS_API_BASE` | Zilliz Cloud cluster public endpoint, e.g. `https://in03-xxxx.serverless.gcp-us-west1.cloud.zilliz.com`. For self-hosted Milvus, the base URL of the REST server | -| `MILVUS_API_KEY` | Zilliz Cloud API key. For self-hosted Milvus, a `user:password` token | -| `OPENAI_API_KEY` | Used to embed the seed documents, and on the proxy to embed queries at search time and to back the chat model the suite registers | - -Everything else is fixed in `milvus_client.py`: the embedding model -(`text-embedding-3-small`, 1536 dims) and the chat model. Rather than depend on -whatever chat model the proxy has, the suite registers its own deployment through -`POST /model/new` and deletes it on teardown, so the chat test is deterministic - -The proxy must be reachable (see `tests/e2e/e2e_config.py` for `LITELLM_PROXY_URL` -/ `LITELLM_MASTER_KEY`), connected to a database (managed vector store -registration persists to the store registry), and started with -`STORE_MODEL_IN_DB=True` (so the chat model can be registered). Without model -storage, only the chat test skips; the search tests still run - -Absent any of the required credentials, or a live proxy, the whole suite skips; -it never silently passes - -## Running - -```bash -cd tests/e2e -uv run --python 3.12 pytest vector-stores/ -v -``` - -The e2e harness uses `type` alias syntax, so Python 3.12+ is required