feat(vector_store)!: replace ChromaDB with Milvus Lite

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -56,7 +56,7 @@ body:
         - ai4rag version:
         - Python version:
         - openai SDK version:
-        - Vector Store (Milvus/ChromaDB):
+        - Vector Store (Milvus/pgvector):
         - Operating System:
       render: markdown
     validations:

README.md

@@ -29,7 +29,7 @@ It accepts a variety of RAG Templates and a search space definition, then return
 
 > [!IMPORTANT]
 > `ai4rag` is **provider-agnostic**. It reaches foundation and embedding models through the stock [`openai`](https://github.com/openai/openai-python) SDK, so any **OpenAI-compatible endpoint** works — a hosted API, a self-managed server (vLLM, TGI, Ollama, …), or an [OpenShift AI Models-as-a-Service (MaaS)](https://www.redhat.com/en/products/ai) deployment, the integration `ai4rag` ships helpers for out of the box. You can also plug in your **own** foundation model, embedding model, or vector store by implementing the matching `Base*` interface.
-> To run an experiment you'll need one foundation model and one embedding model (from any of the above), plus a vector store (Chroma, Milvus, or PostgreSQL/pgvector) connected directly via `ai4rag.rag.vector_store`.
+> To run an experiment you'll need one foundation model and one embedding model (from any of the above), plus a vector store (remote Milvus, embedded Milvus Lite, or PostgreSQL/pgvector) connected directly via `ai4rag.rag.vector_store`.
 
 ## Model providers
 
@@ -48,21 +48,24 @@ When using the MaaS backend, ai4rag relies on:
 - **Embeddings** — Text embeddings via the `embeddings` endpoint (e.g. for indexing and query encoding). Because `models.list()` carries no metadata, embedding dimension and context length are auto-detected at construction (or supplied via `params`).
 - **Chat / completions** — Foundation model integration for answer generation when evaluating RAG patterns.
 
-Vector storage is independent of MaaS: `ai4rag` connects directly to Chroma, Milvus, or PostgreSQL/pgvector via the config classes in `ai4rag.rag.vector_store` (see [Vector stores](#vector-stores) below).
+Vector storage is independent of MaaS: `ai4rag` connects directly to remote Milvus, embedded Milvus Lite, or PostgreSQL/pgvector via the config classes in `ai4rag.rag.vector_store` (see [Vector stores](#vector-stores) below).
 
 ## Vector stores
 
 ai4RAG talks to the vector store directly through provider-specific clients — no MaaS deployment is required for this part. Pick a provider and pass its config to `AI4RAGExperiment` as `vector_store_config`:
 
-- **`ChromaConfig`** — Chroma. Ephemeral in-memory by default; persistent (via `persist_directory`) or client/server (via `host`/`port`) modes are also supported. Vector-only search.
-- **`MilvusConfig`** — Milvus. Requires a `uri`; supports TLS (`https://` scheme) and self-signed CAs via `server_cert`. Hybrid search (dense + BM25).
+- **`MilvusConfig`** — remote Milvus server or Zilliz Cloud only. `uri` must be a `http(s)://` URL (TLS and self-signed CAs via `server_cert`); anything else (a bare host, a file path, an empty string) raises `ValueError`. This is a deliberate safety check: a mistyped or unreachable `MILVUS_URI` now fails loudly instead of silently falling back to a throwaway local database. Supports hybrid search (dense + BM25).
+- **`MilvusLiteConfig`** — the embedded, zero-server **Milvus Lite** engine, backed by a local `db_path` file (default `"./ai4rag_milvus_lite.db"`) — no setup required, ideal for local development and small-scale workloads. Also supports hybrid search (dense + BM25); rejects `http(s)://` values (use `MilvusConfig` for those).
 - **`PGVectorConfig`** — PostgreSQL with the `pgvector` extension. Hybrid search (dense + `tsvector` full-text).
 
-Each config is a frozen dataclass with a `.from_env()` constructor and an `env_vars` attribute listing the environment variables it reads (e.g. `MILVUS_URI`, `PGVECTOR_HOST`).
+Each config is a frozen dataclass with a `.from_env()` constructor and an `env_vars` attribute listing the environment variables it reads (e.g. `MILVUS_URI` for `MilvusConfig`, `MILVUS_LITE_DB_PATH` for `MilvusLiteConfig`, `PGVECTOR_HOST` for `PGVectorConfig`).
+
+> [!note]
+> Milvus Lite is intended for local development, tests, and small-scale workloads (prototyping, up to roughly 1M vectors) — not production serving. For production or large corpora, use a remote Milvus server (`MilvusConfig`), Zilliz Cloud, or pgvector.
 
 ## Document processing
 
-ai4RAG uses [`docling-core`](https://github.com/docling-project/docling-core) for document representation and chunking. Documents are represented as `DoclingDocument` instances, and the `DoclingChunker` leverages docling's `HybridChunker` for structure-aware, token-aware chunking. `docling-core`, `openai`, and the vector store clients (`chromadb`, `pymilvus`, `pgvector`, `asyncpg`) are all installed automatically with `ai4rag`.
+ai4RAG uses [`docling-core`](https://github.com/docling-project/docling-core) for document representation and chunking. Documents are represented as `DoclingDocument` instances, and the `DoclingChunker` leverages docling's `HybridChunker` for structure-aware, token-aware chunking. `docling-core`, `openai`, and the vector store clients (`pymilvus` with Milvus Lite, `pgvector`, `asyncpg`) are all installed automatically with `ai4rag`.
 
 
 ## Quick start
@@ -232,8 +235,8 @@ Using the information from the previous steps, create an experiment and run the
 
 > [!note]
 > Select the vector store by passing a `vector_store_config` to `AI4RAGExperiment`:
-> `ChromaConfig()` for a zero-config in-memory store (vector-only search), or
-> `MilvusConfig.from_env()` / `PGVectorConfig.from_env()` for a server-backed store with hybrid (dense + keyword) search.
+> `MilvusLiteConfig()` (or `MilvusLiteConfig(db_path="./ai4rag.db")`) for a zero-config, local Milvus Lite store, or
+> `MilvusConfig.from_env()` / `PGVectorConfig.from_env()` for a server-backed store. All support hybrid (dense + keyword) search.
 
 ```python
 from ai4rag.core.experiment.experiment import AI4RAGExperiment

ai4rag/core/experiment/experiment.py

@@ -425,11 +425,6 @@ def run_single_evaluation(self, rag_params: RAGParamsType) -> float:
         number_of_chunks = retrieval_params[AI4RAGParamNames.NUMBER_OF_CHUNKS]
 
         search_mode = retrieval_params.get(AI4RAGParamNames.SEARCH_MODE, "vector")
-        if search_mode != "vector" and self.vector_store_config.provider == "chroma":
-            raise RAGExperimentError(
-                f"Search mode '{search_mode}' is not supported with chroma vector store. "
-                "Only 'vector' mode is supported for chroma."
-            )
 
         context_template_text = foundation_model.context_template_text
         system_message_text = foundation_model.system_message_text
@@ -500,11 +495,6 @@ def run_single_evaluation(self, rag_params: RAGParamsType) -> float:
 
         collection_name = vector_store.collection_name
 
-        # The store's connection/client is only needed for indexing and retrieval,
-        # both of which finish before scoring; closing it deterministically here
-        # (rather than waiting on garbage collection) keeps a long HPO search from
-        # accumulating one open connection per evaluated pattern, including on
-        # trials that fail and get caught by search()'s objective_function.
         with vector_store:
             if not self._collection_exists(collection_name=collection_name):
                 chunking_method = chunking_params.get(AI4RAGParamNames.CHUNKING_METHOD)
@@ -903,9 +893,10 @@ def _evaluate_response(
     def _collection_exists(self, collection_name: str) -> bool:
         """
         This method checks if a collection with a given name already exists.
-        The trick comes with chromadb. We always need to assume that collection
-        does not exist, as we create new instance of chroma in memory per each
-        run.
+        Existence is tracked by this run's own bookkeeping
+        (``self.results.collection_names``) rather than by querying the backend,
+        so the check is backend-agnostic and reflects only collections this
+        experiment created and can safely reuse.
 
         Parameters
         ----------

ai4rag/core/experiment/mps.py

@@ -26,7 +26,7 @@
 from ai4rag.rag.retrieval.retriever import Retriever
 from ai4rag.rag.template.simple_rag_template import SimpleRAG
 from ai4rag.rag.vector_store.base_vector_store import BaseVectorStore
-from ai4rag.rag.vector_store.chroma import ChromaVectorStore
+from ai4rag.rag.vector_store.local_store import temporary_milvus_lite_store
 from ai4rag.utils.constants import AI4RAGParamNames, PreSelectorConstants
 
 __all__ = ["PreSelectorError", "ModelsPreSelector"]
@@ -140,8 +140,9 @@ def evaluate_patterns(self):
         If knowledge base references were provided, the retriever is created once and reused.
         Otherwise, a separate vector store is built for each embedding model, and the best-performing ones are selected.
 
-        For evaluation only sample of the documents is used, embedded and added
-        to chroma vector store.
+        For evaluation only a sample of the documents is used, embedded and added
+        to a throwaway, file-backed Milvus Lite vector store that is torn down as
+        soon as the embedding model's evaluation finishes.
 
         This method does not return anything, but in the end changes attributes
         of the instance: self.evaluation_results is a mapping holding results
@@ -156,17 +157,16 @@ def evaluate_patterns(self):
         chunked_documents = self._chunk_documents(documents)
 
         for i, embedding_model in enumerate(self.embedding_models):
+            collection_name = f"ai4rag_mps_collection_{i}"
             try:
-                collection_name = f"ai4rag_mps_collection_{i}"
-                try:
-                    vector_store = self._create_vector_store(
-                        embedding_model, chunked_documents, collection_name=collection_name
-                    )
-                except Exception as exc:
-                    raise IndexingError(exc, collection_name, embedding_model.model_id) from exc
+                with temporary_milvus_lite_store(embedding_model, collection_name=collection_name) as vector_store:
+                    try:
+                        self._index_documents(vector_store, chunked_documents)
+                    except Exception as exc:
+                        raise IndexingError(exc, collection_name, embedding_model.model_id) from exc
 
-                retriever = Retriever(vector_store, **self.retrieval_params)
-                self._evaluate_foundation_models(retriever=retriever, embedding_model=embedding_model)
+                    retriever = Retriever(vector_store, **self.retrieval_params)
+                    self._evaluate_foundation_models(retriever=retriever, embedding_model=embedding_model)
 
             except IndexingError as exc:
                 self._exception_handler.handle_exception(exc)
@@ -221,56 +221,35 @@ def _evaluate_foundation_models(self, retriever: Retriever, embedding_model: Bas
                 continue
 
     @staticmethod
-    def _create_vector_store(
-        embedding_model: BaseEmbeddingModel,
-        chunked_documents: list[AI4RAGChunk],
-        collection_name: str,
-    ) -> BaseVectorStore:
-        """
-        Create instance of vector store with given chunked documents and embedding model.
+    def _index_documents(vector_store: BaseVectorStore, chunked_documents: list[AI4RAGChunk]) -> None:
+        """Embed and add the chunked documents to *vector_store*, retrying once.
+
+        A single retry absorbs transient embedding-service hiccups. If the retry
+        also fails, the exception propagates so the caller can wrap it in an
+        :class:`IndexingError` and skip this embedding model.
 
         Parameters
         ----------
-        embedding_model : BaseEmbeddingModel
-            Embedding model used for collection creation.
+        vector_store : BaseVectorStore
+            Store the chunks are embedded into.
 
         chunked_documents : list[AI4RAGChunk]
             Chunked documents for the embedding process.
 
-        collection_name : str
-            Name of the collection in the chroma vector database.
-
-        Returns
-        -------
-        VectorStore
-            Instance for communication with properly created index in the
-            vector database.
-
         Raises
         ------
-        PreSelectorError
-            When 2 attempts of embedding documents are failing
+        Exception
+            Propagates the second failure when both the initial embedding attempt
+            and the retry fail.
         """
-        logger.info("Building index for pre-evaluation using embedding model: '%s'.", embedding_model.model_id)
-
-        vector_store = ChromaVectorStore(
-            embedding_model=embedding_model,
-            collection_name=collection_name,
-        )
-
         logger.debug("MPS: Embedding documents ...")
         try:
             vector_store.add_documents(chunked_documents)
         except Exception as err:  # pylint: disable=broad-exception-caught
-            logger.warning("Failed to create in-memory vector index due to: %s.", repr(err), exc_info=True)
-            try:
-                vector_store.add_documents(chunked_documents)
-            except Exception as exc:  # pylint: disable=broad-exception-caught
-                raise PreSelectorError(f"Failed to create in-memory vector index due to: {repr(exc)}.") from exc
+            logger.warning("Failed to build the vector index due to: %s. Retrying once.", repr(err), exc_info=True)
+            vector_store.add_documents(chunked_documents)
         logger.debug("MPS: Embedding documents finished!")
 
-        return vector_store
-
     def _evaluate_single_pattern(
         self, foundation_model: BaseFoundationModel, retriever: Retriever
     ) -> EvaluationMetricsResult:

ai4rag/core/experiment/utils.py

@@ -31,7 +31,7 @@
 
 _semantic_chunker_cache = {}
 
-VectorStoreType: TypeAlias = Literal["milvus", "chroma"]
+VectorStoreType: TypeAlias = Literal["milvus", "milvus_lite", "pgvector"]
 
 
 class RAGExperimentError(Exception):

ai4rag/evaluator/judge_selection.py

@@ -194,16 +194,16 @@ def _run_reference_rag(
     list
         List of :class:`EvaluationData` instances ready for judge scoring.
     """
-    from ai4rag.rag.vector_store.chroma import ChromaVectorStore
+    from ai4rag.rag.vector_store.local_store import temporary_milvus_lite_store
 
     chunker = LangChainChunker(chunk_size=512, method="recursive", chunk_overlap=128)
     chunks = chunker.split_documents(documents)
-    vector_store = ChromaVectorStore(embedding_model=embedding_model, collection_name="ai4rag_judge_calibration")
-    vector_store.add_documents(chunks)
-    retriever = Retriever(vector_store=vector_store, number_of_chunks=3, method="simple", search_mode="vector")
-    rag = SimpleRAG(foundation_model=foundation_model, retriever=retriever)
-    inference_response = query_rag(rag=rag, questions=list(benchmark_data.questions), max_threads=max_threads)
-    return build_evaluation_data(benchmark_data=benchmark_data, inference_response=inference_response)
+    with temporary_milvus_lite_store(embedding_model, collection_name="ai4rag_judge_calibration") as vector_store:
+        vector_store.add_documents(chunks)
+        retriever = Retriever(vector_store=vector_store, number_of_chunks=3, method="simple", search_mode="vector")
+        rag = SimpleRAG(foundation_model=foundation_model, retriever=retriever)
+        inference_response = query_rag(rag=rag, questions=list(benchmark_data.questions), max_threads=max_threads)
+        return build_evaluation_data(benchmark_data=benchmark_data, inference_response=inference_response)
 
 
 def _ordered_question_scores(evaluation_result: EvaluationMetricsResult, metric_name: str) -> list[float | None]:

ai4rag/rag/vector_store/__init__.py

@@ -4,8 +4,8 @@
 # -----------------------------------------------------------------------------
 from ai4rag.rag.vector_store.base_vector_store import BaseVectorStore
 from ai4rag.rag.vector_store.config import (
-    ChromaConfig,
     MilvusConfig,
+    MilvusLiteConfig,
     PGVectorConfig,
     get_vector_store_config,
     get_vector_store_env_vars,
@@ -14,8 +14,8 @@
 
 __all__ = [
     "BaseVectorStore",
-    "ChromaConfig",
     "MilvusConfig",
+    "MilvusLiteConfig",
     "PGVectorConfig",
     "get_vector_store",
     "get_vector_store_config",