Merge e70d863350 into 30ff3723b2

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"
+    )

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),
+    )

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