diff --git a/.github/workflows/integration-vector-io-tests.yml b/.github/workflows/integration-vector-io-tests.yml index 5be1607c3..c11720b4b 100644 --- a/.github/workflows/integration-vector-io-tests.yml +++ b/.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 diff --git a/llama_stack/providers/registry/vector_io.py b/llama_stack/providers/registry/vector_io.py index 5e233b94c..de7e08445 100644 --- a/llama_stack/providers/registry/vector_io.py +++ b/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=""" diff --git a/tests/integration/vector_io/test_vector_io.py b/tests/integration/vector_io/test_vector_io.py index 95fcb8db5..9cd4fc18c 100644 --- a/tests/integration/vector_io/test_vector_io.py +++ b/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