chore: Enabling VectorIO Integration tests for Milvus (#2546)

.github/workflows/integration-vector-io-tests.yml

@@ -22,7 +22,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        vector-io-provider: ["inline::faiss", "inline::sqlite-vec", "remote::chromadb", "remote::pgvector"]
+        vector-io-provider: ["inline::faiss", "inline::sqlite-vec", "inline::milvus", "remote::chromadb", "remote::pgvector"]
         python-version: ["3.12", "3.13"]
       fail-fast: false # we want to run all tests regardless of failure
 

llama_stack/providers/registry/vector_io.py

@@ -520,7 +520,7 @@ Please refer to the inline provider documentation.
             Api.vector_io,
             AdapterSpec(
                 adapter_type="milvus",
-                pip_packages=["pymilvus"],
+                pip_packages=["pymilvus[marshmallow<3.13.0]"],
                 module="llama_stack.providers.remote.vector_io.milvus",
                 config_class="llama_stack.providers.remote.vector_io.milvus.MilvusVectorIOConfig",
                 description="""

tests/integration/vector_io/test_vector_io.py

@@ -123,6 +123,9 @@ def test_insert_chunks(client_with_empty_registry, embedding_model_id, embedding
 
 
 def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, embedding_model_id, embedding_dimension):
+    vector_io_provider_params_dict = {
+        "inline::milvus": {"score_threshold": -1.0},
+    }
     vector_db_id = "test_precomputed_embeddings_db"
     client_with_empty_registry.vector_dbs.register(
         vector_db_id=vector_db_id,
@@ -133,7 +136,7 @@ def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, e
     chunks_with_embeddings = [
         Chunk(
             content="This is a test chunk with precomputed embedding.",
-            metadata={"document_id": "doc1", "source": "precomputed"},
+            metadata={"document_id": "doc1", "source": "precomputed", "chunk_id": "chunk1"},
             embedding=[0.1] * int(embedding_dimension),
         ),
     ]
@@ -143,22 +146,29 @@ def test_insert_chunks_with_precomputed_embeddings(client_with_empty_registry, e
         chunks=chunks_with_embeddings,
     )
 
-    # Query for the first document
+    provider = [p.provider_id for p in client_with_empty_registry.providers.list() if p.api == "vector_io"][0]
     response = client_with_empty_registry.vector_io.query(
         vector_db_id=vector_db_id,
         query="precomputed embedding test",
+        params=vector_io_provider_params_dict.get(provider, None),
     )
 
     # Verify the top result is the expected document
     assert response is not None
-    assert len(response.chunks) > 0
+    assert len(response.chunks) > 0, (
+        f"provider params for {provider} = {vector_io_provider_params_dict.get(provider, None)}"
+    )
     assert response.chunks[0].metadata["document_id"] == "doc1"
     assert response.chunks[0].metadata["source"] == "precomputed"
 
 
+# expect this test to fail
 def test_query_returns_valid_object_when_identical_to_embedding_in_vdb(
     client_with_empty_registry, embedding_model_id, embedding_dimension
 ):
+    vector_io_provider_params_dict = {
+        "inline::milvus": {"score_threshold": 0.0},
+    }
     vector_db_id = "test_precomputed_embeddings_db"
     client_with_empty_registry.vector_dbs.register(
         vector_db_id=vector_db_id,
@@ -179,9 +189,11 @@ def test_query_returns_valid_object_when_identical_to_embedding_in_vdb(
         chunks=chunks_with_embeddings,
     )
 
+    provider = [p.provider_id for p in client_with_empty_registry.providers.list() if p.api == "vector_io"][0]
     response = client_with_empty_registry.vector_io.query(
         vector_db_id=vector_db_id,
         query="duplicate",
+        params=vector_io_provider_params_dict.get(provider, None),
     )
 
     # Verify the top result is the expected document