add milvus support (#187)

Author: ArnavAgrawal03

core/config.py

@@ -23,6 +23,9 @@ class Settings(BaseSettings):
     OPENAI_API_KEY: Optional[str] = None
     ANTHROPIC_API_KEY: Optional[str] = None
     ASSEMBLYAI_API_KEY: Optional[str] = None
+    # Milvus configuration
+    MILVUS_URI: Optional[str] = None
+    MILVUS_API_KEY: Optional[str] = None
 
     # API configuration
     HOST: str
@@ -107,9 +110,13 @@ class Settings(BaseSettings):
     S3_BUCKET: Optional[str] = None
 
     # Vector store configuration
-    VECTOR_STORE_PROVIDER: Literal["pgvector"]
+    VECTOR_STORE_PROVIDER: Literal["pgvector", "milvus"]
     VECTOR_STORE_DATABASE_NAME: Optional[str] = None
 
+    # Multi-vector store configuration
+    MULTIVECTOR_PROVIDER: Literal["postgres", "milvus"] = "postgres"
+    MILVUS_BATCH_SIZE: int = 500  # Batch size for Milvus multivector insertions
+
     # Colpali configuration
     ENABLE_COLPALI: bool
     # Colpali embedding mode: off, local, or api
@@ -288,13 +295,54 @@ def get_settings() -> Settings:
 
     # load vector store config
     vector_store_config = {"VECTOR_STORE_PROVIDER": config["vector_store"]["provider"]}
-    if vector_store_config["VECTOR_STORE_PROVIDER"] != "pgvector":
+    if vector_store_config["VECTOR_STORE_PROVIDER"] not in ["pgvector", "milvus"]:
         prov = vector_store_config["VECTOR_STORE_PROVIDER"]
         raise ValueError(f"Unknown vector store provider selected: '{prov}'")
 
-    if "POSTGRES_URI" not in os.environ:
-        msg = em.format(missing_value="POSTGRES_URI", field="vector_store.provider", value="pgvector")
-        raise ValueError(msg)
+    # Validate required environment variables based on vector store provider
+    if vector_store_config["VECTOR_STORE_PROVIDER"] == "pgvector":
+        if "POSTGRES_URI" not in os.environ:
+            msg = em.format(missing_value="POSTGRES_URI", field="vector_store.provider", value="pgvector")
+            raise ValueError(msg)
+    elif vector_store_config["VECTOR_STORE_PROVIDER"] == "milvus":
+        if "MILVUS_URI" not in os.environ:
+            msg = em.format(missing_value="MILVUS_URI", field="vector_store.provider", value="milvus")
+            raise ValueError(msg)
+        vector_store_config.update(
+            {
+                "MILVUS_URI": os.environ["MILVUS_URI"],
+                "MILVUS_API_KEY": os.environ.get("MILVUS_API_KEY"),  # API key is optional for some Milvus setups
+            }
+        )
+
+    # load multivector store config
+    multivector_store_config = {}
+    if "multivector_store" in config:
+        multivector_store_config = {
+            "MULTIVECTOR_PROVIDER": config["multivector_store"]["provider"],
+            "MILVUS_BATCH_SIZE": config["multivector_store"].get("milvus_batch_size", 500),  # Default to 500
+        }
+        if multivector_store_config["MULTIVECTOR_PROVIDER"] not in ["postgres", "milvus"]:
+            prov = multivector_store_config["MULTIVECTOR_PROVIDER"]
+            raise ValueError(f"Unknown multivector store provider selected: '{prov}'")
+
+        # Validate required environment variables based on multivector store provider
+        if multivector_store_config["MULTIVECTOR_PROVIDER"] == "postgres":
+            if "POSTGRES_URI" not in os.environ:
+                msg = em.format(missing_value="POSTGRES_URI", field="multivector_store.provider", value="postgres")
+                raise ValueError(msg)
+        elif multivector_store_config["MULTIVECTOR_PROVIDER"] == "milvus":
+            if "MILVUS_URI" not in os.environ:
+                msg = em.format(missing_value="MILVUS_URI", field="multivector_store.provider", value="milvus")
+                raise ValueError(msg)
+            # Add Milvus credentials to config if not already added
+            if "MILVUS_URI" not in vector_store_config:
+                multivector_store_config.update(
+                    {
+                        "MILVUS_URI": os.environ["MILVUS_URI"],
+                        "MILVUS_API_KEY": os.environ.get("MILVUS_API_KEY"),
+                    }
+                )
 
     # load rules config
     rules_config = {
@@ -382,6 +430,7 @@ def get_settings() -> Settings:
             reranker_config,
             storage_config,
             vector_store_config,
+            multivector_store_config,
             rules_config,
             morphik_config,
             redis_config,

core/services_init.py

@@ -28,6 +28,8 @@
 from core.services.document_service import DocumentService
 from core.storage.local_storage import LocalStorage
 from core.storage.s3_storage import S3Storage
+from core.vector_store.milvus_multivector_store import MilvusMultiVectorStore
+from core.vector_store.milvus_vector_store import MilvusVectorStore
 from core.vector_store.multi_vector_store import MultiVectorStore
 from core.vector_store.pgvector_store import PGVectorStore
 
@@ -49,8 +51,18 @@
 database = PostgresDatabase(uri=settings.POSTGRES_URI)
 logger.debug("Created PostgresDatabase singleton")
 
-vector_store = PGVectorStore(uri=settings.POSTGRES_URI)
-logger.debug("Created PGVectorStore singleton")
+# Initialize vector store based on configuration
+match settings.VECTOR_STORE_PROVIDER:
+    case "pgvector":
+        vector_store = PGVectorStore(uri=settings.POSTGRES_URI)
+        logger.info("Using PGVector for main vector storage")
+    case "milvus":
+        vector_store = MilvusVectorStore()
+        logger.info("Using Milvus for main vector storage")
+    case _:
+        raise ValueError(f"Unsupported vector store provider: {settings.VECTOR_STORE_PROVIDER}")
+
+logger.debug("Created vector store singleton")
 
 # ---------------------------------------------------------------------------
 # Object storage
@@ -121,16 +133,31 @@
 # ColPali multi-vector support
 # ---------------------------------------------------------------------------
 
+colpali_embedding_model = None
+colpali_vector_store = None
+
 match settings.COLPALI_MODE:
     case "off":
         colpali_embedding_model = None
         colpali_vector_store = None
     case "local":
         colpali_embedding_model = ColpaliEmbeddingModel()
-        colpali_vector_store = MultiVectorStore(uri=settings.POSTGRES_URI)
+        # Check if we should use Milvus or PostgreSQL for multi-vector storage
+        if settings.MULTIVECTOR_PROVIDER.lower() == "milvus":
+            colpali_vector_store = MilvusMultiVectorStore(batch_size=settings.MILVUS_BATCH_SIZE)
+            logger.info("Using Milvus for ColPali multi-vector storage")
+        else:
+            colpali_vector_store = MultiVectorStore(uri=settings.POSTGRES_URI)
+            logger.info("Using PostgreSQL for ColPali multi-vector storage")
     case "api":
         colpali_embedding_model = ColpaliApiEmbeddingModel()
-        colpali_vector_store = MultiVectorStore(uri=settings.POSTGRES_URI)
+        # Check if we should use Milvus or PostgreSQL for multi-vector storage
+        if settings.MULTIVECTOR_PROVIDER.lower() == "milvus":
+            colpali_vector_store = MilvusMultiVectorStore(batch_size=settings.MILVUS_BATCH_SIZE)
+            logger.info("Using Milvus for ColPali multi-vector storage (API mode)")
+        else:
+            colpali_vector_store = MultiVectorStore(uri=settings.POSTGRES_URI)
+            logger.info("Using PostgreSQL for ColPali multi-vector storage (API mode)")
     case _:
         raise ValueError(f"Unsupported COLPALI_MODE: {settings.COLPALI_MODE}")
 

core/tests/unit/test_multivector.py

@@ -4,11 +4,10 @@
 import numpy as np
 import pytest
 import torch
-from pgvector.psycopg import Bit
 
 from core.models.chunk import DocumentChunk
 from core.tests import setup_test_logging
-from core.vector_store.multi_vector_store import MultiVectorStore
+from core.vector_store.milvus_multivector_store import MilvusMultiVectorStore
 
 # Set up test logging
 setup_test_logging()
@@ -50,102 +49,115 @@ def get_sample_document_chunks(num_chunks=3, num_vectors=3, dim=128):
     return chunks
 
 
-# Fixtures
+# For Milvus
 @pytest.fixture(scope="function")
 async def vector_store():
     """Create a real MultiVectorStore instance connected to the test database"""
-    # Create the store
-    store = MultiVectorStore(uri=TEST_DB_URI)
-
-    try:
-        # Try to initialize the database
-        store.initialize()
-
-        # Clean up any existing data
-        store.conn.execute("TRUNCATE TABLE multi_vector_embeddings RESTART IDENTITY")
-
-        # Drop the function if it exists
-        try:
-            store.conn.execute("DROP FUNCTION IF EXISTS max_sim(bit[], bit[])")
-        except Exception as e:
-            print(f"Error dropping function: {e}")
-    except Exception as e:
-        print(f"Error setting up database: {e}")
-
+    store = MilvusMultiVectorStore(collection_name="test_collection")
+    store.client.drop_collection(collection_name="test_collection")
+    store._create_collection()
+    store.client.load_collection(collection_name="test_collection")
     yield store
-
-    # Clean up after tests
-    try:
-        store.conn.execute("TRUNCATE TABLE multi_vector_embeddings RESTART IDENTITY")
-    except Exception as e:
-        print(f"Error cleaning up: {e}")
-
-    # Close connection
     store.close()
 
 
-# Glassbox Tests - Testing internal implementation details
-@pytest.mark.asyncio
-async def test_binary_quantize():
-    """Test the _binary_quantize method correctly converts embeddings"""
-    store = MultiVectorStore(uri=TEST_DB_URI)
-
-    # Test with torch tensor
-    torch_embeddings = torch.tensor([[0.1, -0.2, 0.3], [-0.1, 0.2, -0.3]])
-    binary_result = store._binary_quantize(torch_embeddings)
-    assert len(binary_result) == 2
-
-    # Check results match expected patterns
-    assert binary_result[0].to_text() == Bit("101").to_text()  # Positive values (>0) become 1, negative/zero become 0
-    assert binary_result[1].to_text() == Bit("010").to_text()  # First row: [0.1 (>0), -0.2 (<0), 0.3 (>0)] → "101"
-    # Second row: [-0.1 (<0), 0.2 (>0), -0.3 (<0)] → "010"
-
-    # Test with numpy array
-    numpy_embeddings = np.array([[0.1, -0.2, 0.3], [-0.1, 0.2, -0.3]])
-    binary_result = store._binary_quantize(numpy_embeddings)
-    assert len(binary_result) == 2
-
-    assert binary_result[0].to_text() == Bit("101").to_text()
-    assert binary_result[1].to_text() == Bit("010").to_text()
-
-
-@pytest.mark.asyncio
-async def test_initialize_creates_tables_and_function(vector_store):
-    """Test that initialize creates the necessary tables and functions"""
-    vector_store.initialize()
-    # Check if the table exists
-    result = vector_store.conn.execute(
-        "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'multi_vector_embeddings')"
-    ).fetchone()
-    table_exists = result[0]
-    assert table_exists is True
-
-    logger.info("Table exists!")
-
-    # Check if the max_sim function exists
-    result = vector_store.conn.execute("SELECT EXISTS (SELECT FROM pg_proc WHERE proname = 'max_sim')").fetchone()
-    function_exists = result[0]
-    logger.info(f"Function exists {function_exists}")
-    assert function_exists is True
-
-
-@pytest.mark.asyncio
-async def test_database_schema(vector_store):
-    """Test that the database schema matches our expectations"""
-    # Check columns in the table
-    result = vector_store.conn.execute(
-        "SELECT column_name, data_type FROM information_schema.columns " "WHERE table_name = 'multi_vector_embeddings'"
-    ).fetchall()
-
-    # Convert to a dict for easier checking
-    column_dict = {col[0]: col[1] for col in result}
-
-    # Check required columns
-    assert "id" in column_dict
-    assert "document_id" in column_dict
-    assert "chunk_number" in column_dict
-    assert "content" in column_dict
-    assert "embeddings" in column_dict
+# For Postgres
+# # Fixtures
+# @pytest.fixture(scope="function")
+# async def vector_store():
+#     """Create a real MultiVectorStore instance connected to the test database"""
+#     # Create the store
+#     store = MultiVectorStore(uri=TEST_DB_URI)
+
+#     try:
+#         # Try to initialize the database
+#         store.initialize()
+
+#         # Clean up any existing data
+#         store.conn.execute("TRUNCATE TABLE multi_vector_embeddings RESTART IDENTITY")
+
+#         # Drop the function if it exists
+#         try:
+#             store.conn.execute("DROP FUNCTION IF EXISTS max_sim(bit[], bit[])")
+#         except Exception as e:
+#             print(f"Error dropping function: {e}")
+#     except Exception as e:
+#         print(f"Error setting up database: {e}")
+
+#     yield store
+
+#     # Clean up after tests
+#     try:
+#         store.conn.execute("TRUNCATE TABLE multi_vector_embeddings RESTART IDENTITY")
+#     except Exception as e:
+#         print(f"Error cleaning up: {e}")
+
+#     # Close connection
+#     store.close()
+
+
+# # Glassbox Tests - Testing internal implementation details
+# @pytest.mark.asyncio
+# async def test_binary_quantize():
+#     """Test the _binary_quantize method correctly converts embeddings"""
+#     store = MultiVectorStore(uri=TEST_DB_URI)
+
+#     # Test with torch tensor
+#     torch_embeddings = torch.tensor([[0.1, -0.2, 0.3], [-0.1, 0.2, -0.3]])
+#     binary_result = store._binary_quantize(torch_embeddings)
+#     assert len(binary_result) == 2
+
+#     # Check results match expected patterns
+#     assert binary_result[0].to_text() == Bit("101").to_text()  # Positive values (>0) become 1, negative/zero become 0
+#     assert binary_result[1].to_text() == Bit("010").to_text()  # First row: [0.1 (>0), -0.2 (<0), 0.3 (>0)] → "101"
+#     # Second row: [-0.1 (<0), 0.2 (>0), -0.3 (<0)] → "010"
+
+#     # Test with numpy array
+#     numpy_embeddings = np.array([[0.1, -0.2, 0.3], [-0.1, 0.2, -0.3]])
+#     binary_result = store._binary_quantize(numpy_embeddings)
+#     assert len(binary_result) == 2
+
+#     assert binary_result[0].to_text() == Bit("101").to_text()
+#     assert binary_result[1].to_text() == Bit("010").to_text()
+
+
+# @pytest.mark.asyncio
+# async def test_initialize_creates_tables_and_function(vector_store):
+#     """Test that initialize creates the necessary tables and functions"""
+#     vector_store.initialize()
+#     # Check if the table exists
+#     result = vector_store.conn.execute(
+#         "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'multi_vector_embeddings')"
+#     ).fetchone()
+#     table_exists = result[0]
+#     assert table_exists is True
+
+#     logger.info("Table exists!")
+
+#     # Check if the max_sim function exists
+#     result = vector_store.conn.execute("SELECT EXISTS (SELECT FROM pg_proc WHERE proname = 'max_sim')").fetchone()
+#     function_exists = result[0]
+#     logger.info(f"Function exists {function_exists}")
+#     assert function_exists is True
+
+
+# @pytest.mark.asyncio
+# async def test_database_schema(vector_store):
+#     """Test that the database schema matches our expectations"""
+#     # Check columns in the table
+#     result = vector_store.conn.execute(
+#         "SELECT column_name, data_type FROM information_schema.columns " "WHERE table_name = 'multi_vector_embeddings'"
+#     ).fetchall()
+
+#     # Convert to a dict for easier checking
+#     column_dict = {col[0]: col[1] for col in result}
+
+#     # Check required columns
+#     assert "id" in column_dict
+#     assert "document_id" in column_dict
+#     assert "chunk_number" in column_dict
+#     assert "content" in column_dict
+#     assert "embeddings" in column_dict
 
 
 # Blackbox Tests - Testing the public API