feat: add metadata based filtering to VectorStoreOptions (#564)

Author: mhordynski

docs/how-to/document_search/search-documents.md

@@ -73,6 +73,35 @@ Searching for elements is performed using a vector store. [`DocumentSearch`][rag
 
     To learn more about using Hybrid Search, refer to [How to Perform Hybrid Search with Multiple Vector Stores](../vector_stores/hybrid.md).
 
+## Limit results with metadata-based filtering
+
+You can filter search results based on document metadata using the `where` clause in `VectorStoreOptions`. This allows you to narrow down results to specific document types, sources, or any other metadata fields you've defined.
+
+```python
+from ragbits.core.vector_stores.base import VectorStoreOptions
+from ragbits.document_search import DocumentSearch, DocumentSearchOptions
+
+# Create vector store options with metadata filtering
+vector_store_options = VectorStoreOptions(
+    k=2,  # Number of results to return
+    score_threshold=0.6,  # Minimum similarity score
+    where={"document_meta": {"document_type": "txt"}}  # Filter by document type
+)
+
+# Create document search options with the vector store options
+options = DocumentSearchOptions(vector_store_options=vector_store_options)
+
+# Search with the filtering options
+results = await document_search.search("Your search query", options=options)
+```
+
+The `where` clause supports various filtering conditions. For example, you can filter by:
+- Document type
+- Source
+- Custom metadata fields
+
+This filtering happens at the vector store level, making the search more efficient by reducing the number of documents that need to be processed.
+
 ## Rephrase query
 
 By default, the input query is provided directly to the embedding model. However, there is an option to add an additional step before vector search. Ragbits offers several common rephrasing techniques that can be utilized to refine the query and generate better embeddings for retrieval.

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Allow to limit VectorStore results by metadata (#564)
 - Switch from imghdr to filetype for image file type check (#563)
 - Remove prompt lab (#549)
 - Add batched() helper method to utils (#555)

packages/ragbits-core/pyproject.toml

@@ -49,7 +49,7 @@ dependencies = [
 
 [project.optional-dependencies]
 chroma = [
-    "chromadb>=0.6.3,<1.0.0",
+    "chromadb>=1.0.0,<2.0.0",
 ]
 local = [
     "sentence-transformers>=4.0.2,<5.0.0",

packages/ragbits-core/src/ragbits/core/vector_stores/base.py

@@ -13,7 +13,7 @@
 from ragbits.core.utils.config_handling import ConfigurableComponent, ObjectConstructionConfig
 from ragbits.core.utils.pydantic import SerializableBytes
 
-WhereQuery = dict[str, str | int | float | bool]
+WhereQuery = dict[str, str | int | float | bool | dict]
 
 
 class VectorStoreEntry(BaseModel):
@@ -69,10 +69,13 @@ class VectorStoreOptions(Options):
             Note that this is based on score, which may be different from the raw
             similarity metric used by the vector store (see `VectorStoreResult`
             for more details).
+        where: The filter dictionary - the keys are the field names and the values are the values to filter by.
+            Not specifying the key means no filtering.
     """
 
     k: int = 5
     score_threshold: float | None = None
+    where: WhereQuery | None = None
 
 
 VectorStoreOptionsT = TypeVar("VectorStoreOptionsT", bound=VectorStoreOptions)

packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py

@@ -3,7 +3,8 @@
 from uuid import UUID
 
 import chromadb
-from chromadb.api import ClientAPI, types
+from chromadb.api import ClientAPI
+from chromadb.api.types import IncludeMetadataDocuments, IncludeMetadataDocumentsEmbeddingsDistances
 from typing_extensions import Self
 
 from ragbits.core.audit.traces import trace
@@ -193,15 +194,13 @@ async def retrieve(
             query_vector = (await self._embedder.embed_text([text]))[0]
             query_vector = cast(list[float], query_vector)
 
+            where_dict = self._create_chroma_filter(merged_options.where)
+
             results = self._collection.query(
                 query_embeddings=query_vector,
                 n_results=merged_options.k,
-                include=[
-                    types.IncludeEnum.metadatas,
-                    types.IncludeEnum.embeddings,
-                    types.IncludeEnum.distances,
-                    types.IncludeEnum.documents,
-                ],
+                include=IncludeMetadataDocumentsEmbeddingsDistances,
+                where=where_dict,
             )
 
             ids = [id for batch in results.get("ids", []) for id in batch]
@@ -266,14 +265,13 @@ async def list(
         with trace(
             where=where, collection=self._collection, index_name=self._index_name, limit=limit, offset=offset
         ) as outputs:
-            # Cast `where` to chromadb's Where type
-            where_chroma: chromadb.Where | None = dict(where) if where else None
+            where_chroma = self._create_chroma_filter(where)
 
             results = self._collection.get(
                 where=where_chroma,
                 limit=limit,
                 offset=offset,
-                include=[types.IncludeEnum.metadatas, types.IncludeEnum.documents],
+                include=IncludeMetadataDocuments,
             )
 
             ids = results.get("ids") or []
@@ -301,3 +299,22 @@ async def list(
     def _flatten_metadata(metadata: dict) -> dict:
         """Flattens the metadata dictionary. Removes any None values as they are not supported by ChromaDB."""
         return {k: v for k, v in flatten_dict(metadata).items() if v is not None}
+
+    @staticmethod
+    def _create_chroma_filter(where: WhereQuery | None) -> chromadb.Where | None:
+        """
+        Creates a ChromaDB filter from a WhereQuery.
+
+        Args:
+            where: The filter dictionary - the keys are the field names and the values are the values to filter by.
+
+        Returns:
+            The ChromaDB filter.
+        """
+        if not where:
+            return None
+
+        # If there are multiple filters, combine them with $and
+        if len(where) > 1:
+            return cast(chromadb.Where, {"$and": [{k: v} for k, v in flatten_dict(where).items()]})
+        return cast(chromadb.Where, where)

packages/ragbits-core/src/ragbits/core/vector_stores/in_memory.py

@@ -90,6 +90,14 @@ async def retrieve(
             results: list[VectorStoreResult] = []
 
             for entry_id, vector in self._embeddings.items():
+                entry = self._entries[entry_id]
+
+                # Apply metadata filtering
+                if merged_options.where and not all(
+                    entry.metadata.get(key) == value for key, value in merged_options.where.items()
+                ):
+                    continue
+
                 # Calculate score based on vector type
                 if isinstance(query_vector, SparseVector) and isinstance(vector, SparseVector):
                     # For sparse vectors, use dot product between query and document vectors

packages/ragbits-core/src/ragbits/core/vector_stores/pgvector.py

@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Any, NamedTuple, cast
+from typing import Any, NamedTuple
 from uuid import UUID
 
 import asyncpg
@@ -173,13 +173,19 @@ def _create_retrieve_query(
         # _table_name has been validated in the class constructor, and it is a valid table name.
         query = f"SELECT *, vector {distance_operator} $1 as distance, {score_formula} as score FROM {self._table_name}"  # noqa S608
 
-        values: list[Any] = [
-            self._vector_to_string(vector),
-        ]
+        values: list[Any] = [self._vector_to_string(vector)]
+        where_clauses = []
 
         if query_options.score_threshold is not None:
-            query += " WHERE score >= $2"
-            values.extend([query_options.score_threshold])
+            where_clauses.append("score >= $" + str(len(values) + 1))
+            values.append(query_options.score_threshold)
+
+        if query_options.where:
+            where_clauses.append(f"metadata @> ${len(values) + 1}")
+            values.append(json.dumps(query_options.where))
+
+        if where_clauses:
+            query += " WHERE " + " AND ".join(where_clauses)
 
         query += " ORDER BY distance"
 
@@ -351,25 +357,23 @@ async def retrieve(
         Returns:
             The retrieved entries.
         """
-        query_options = (self.default_options | options) if options else self.default_options
+        merged_options = (self.default_options | options) if options else self.default_options
+
         with trace(
             text=text,
+            options=merged_options.dict(),
             table_name=self._table_name,
-            query_options=query_options,
             vector_size=self._vector_size,
             distance_method=self._distance_method,
             embedder=repr(self._embedder),
             embedding_type=self._embedding_type,
         ) as outputs:
-            vector = (await self._embedder.embed_text([text]))[0]
-            vector = cast(list[float], vector)
-
-            query_options = (self.default_options | options) if options else self.default_options
-            retrieve_query, values = self._create_retrieve_query(vector, query_options)
+            query_vector = (await self._embedder.embed_text([text]))[0]
+            query, values = self._create_retrieve_query(query_vector, merged_options)
 
             try:
                 async with self._client.acquire() as conn:
-                    results = await conn.fetch(retrieve_query, *values)
+                    results = await conn.fetch(query, *values)
 
                 outputs.results = [
                     VectorStoreResult(

packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py

@@ -24,7 +24,6 @@
     EmbeddingType,
     VectorStoreEntry,
     VectorStoreOptions,
-    VectorStoreOptionsT,
     VectorStoreResult,
     VectorStoreWithEmbedder,
     WhereQuery,
@@ -214,7 +213,11 @@ async def store(self, entries: list[VectorStoreEntry]) -> None:
                 wait=True,
             )
 
-    async def retrieve(self, text: str, options: VectorStoreOptionsT | None = None) -> list[VectorStoreResult]:
+    async def retrieve(
+        self,
+        text: str,
+        options: VectorStoreOptions | None = None,
+    ) -> list[VectorStoreResult]:
         """
         Retrieves entries from the Qdrant collection based on vector similarity.
 
@@ -236,7 +239,7 @@ async def retrieve(self, text: str, options: VectorStoreOptionsT | None = None)
         )
         with trace(
             text=text,
-            options=merged_options,
+            options=merged_options.dict(),
             index_name=self._index_name,
             distance_method=self._distance_method,
             embedder=repr(self._embedder),
@@ -252,6 +255,7 @@ async def retrieve(self, text: str, options: VectorStoreOptionsT | None = None)
                 score_threshold=score_threshold,
                 with_payload=True,
                 with_vectors=True,
+                query_filter=self._create_qdrant_filter(merged_options.where),
             )
 
             outputs.results = []
@@ -290,16 +294,19 @@ async def remove(self, ids: list[UUID]) -> None:
             )
 
     @staticmethod
-    def _create_qdrant_filter(where: WhereQuery) -> Filter:
+    def _create_qdrant_filter(where: WhereQuery | None) -> Filter:
         """
         Creates the QdrantFilter from the given WhereQuery.
 
         Args:
-            where: The WhereQuery to filter.
+            where: The WhereQuery to filter. If None, returns an empty filter.
 
         Returns:
             The created filter.
         """
+        if where is None:
+            return Filter(must=[])
+
         where = flatten_dict(where)  # type: ignore
 
         return Filter(

packages/ragbits-core/tests/integration/vector_stores/test_vector_store.py

@@ -250,3 +250,38 @@ async def test_handling_document_ingestion_with_different_content_and_verifying_
     assert document_1_content in document_contents
     assert document_2_new_content in document_contents
     assert document_2_content not in document_contents
+
+
+async def test_vector_store_retrieve_with_where_clause(
+    text_vector_store: VectorStoreWithDenseEmbedder,
+    vector_store_entries: list[VectorStoreEntry],
+) -> None:
+    await text_vector_store.store(vector_store_entries)
+
+    # Test with a simple where clause
+    results = await text_vector_store.retrieve(
+        text="foo",
+        options=VectorStoreOptions(
+            where={
+                "foo": "bar",
+                "nested_foo": {"nested_bar": "nested_baz"},
+            }
+        ),
+    )
+
+    # Should only return the first entry which matches both conditions
+    assert len(results) == 1
+    assert results[0].entry.id == vector_store_entries[0].id
+    assert results[0].entry.metadata["foo"] == "bar"
+    assert results[0].entry.metadata["nested_foo"]["nested_bar"] == "nested_baz"
+
+    # Test with a where clause that matches no entries
+    results = await text_vector_store.retrieve(
+        text="foo",
+        options=VectorStoreOptions(
+            where={
+                "foo": "nonexistent",
+            }
+        ),
+    )
+    assert len(results) == 0