Add search_sources and search_documents methods to VectorIndex

Author: tomusher

docs/modules/index/index.md

@@ -64,4 +64,70 @@ The Index Module provides a vector indexing system for Django applications. It e
     ```
 
 5. Build your indexes with `manage.py rebuild_indexes`
-6. Query your index with `MyIndex().search("query")`
+6. Query your index with `MyIndex().search_sources("query")`
+
+## Querying Indexes
+
+When indexes are built, source objects are often chunked in to many separate Documents before they are embedded and inserted in to the index.
+
+This means that a query to the underlying index can return multiple Documents from the same source object. For example; if you have a `Book` Django model with a big summary to be embedded, searching the index might return many Documents from the same Book.
+
+This can be fine in some cases, in RAG applications the most relevant chunks are usually what you want, even if they all come from the same source.
+
+In other cases, such as finding similar content, this behaviour can be a hindrance.
+
+To solve this, Vector Indexes provide two query methods depending on your needs:
+
+### Document Search
+
+```python
+
+MyVectorIndex().search_documents("Similar to this")
+```
+
+`search_documents` returns a queryset-like interface over Document objects. If the underlying vector provider returns multiple Documents from the same source object, these will all be returned.
+
+This is useful for RAG-like applications where the most relevant chunks are important.
+
+### Source Search
+
+```python
+
+MyVectorIndex().search_sources("Similar to this")
+```
+
+When using the `search_sources` method, a Vector Index will attempt to map results from the index back to original source objects, i.e. in the `Book` example, when using `ModelSource(model=Book)`, this method will return a queryset-like interface over `Book` models.
+
+As the underlying storage provider is likely to return multiple Documents for the same source object, this method overfetches Documents to attempt to ensure enough source objects are returned for your query.
+
+This overfetching behaviour can be customised:
+
+```python
+
+MyVectorIndex().search_sources(
+    "Similar to this",
+    overfetch_multiplier=4,
+    max_overfetch_iterations=3
+)
+```
+
+Where:
+
+-   `overfetch_multiplier` defines how many multiples of the requested limit will be retrieved from the source, e.g. if you request 5 results and provide an `overfetch_multiplier` of 4, 20 Documents will be retrieved from the index internally. The top 5 unique sources from these will then be returned.
+-   `max_overfetch_iterations` defines the maximum number of times the underlying search will be repeated to get all-unique source objects, e.g. if the initial search doesn't return enough unique objects, it will be repeated with an increasing number of items up to `max_overfetch_iterations` times.
+
+### Converting Between Result Types
+
+You can convert between result types on an existing queryset:
+
+```python
+# Start with document search, convert to sources
+docs = MyVectorIndex().search_documents("query")
+sources = docs.as_sources()
+
+# Start with source search, convert to documents
+sources = MyVectorIndex().search_sources("query")
+docs = sources.as_documents()
+```
+
+This can be useful in RAG applications where you want to use `Documents` for building context, but then present source objects to users as the 'Sources referenced'.

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "django-ai-core"
-version = "0.1.0"
+version = "0.1.1"
 description = "Django AI Core provides developer-focused AI features for implementing AI tooling in to Django sites."
 readme = "README.md"
 license = "BSD-3-Clause"

src/django_ai_core/contrib/index/base.py

@@ -7,10 +7,10 @@
 
 if TYPE_CHECKING:
     from .embedding import EmbeddingTransformer
-    from .query import QueryHandler, ResultQuerySetMixin
+    from .query import QueryHandler
     from .schema import Document
     from .source import Source
-    from .storage.base import StorageProvider
+    from .storage.base import BaseStorageQuerySet, StorageProvider
 
 
 logger = logging.getLogger(__name__)
@@ -85,16 +85,37 @@ def update(self, documents: Iterable["Document"]):
             if isinstance(source, HasPostIndexUpdateHook):
                 source.post_index_update(self)
 
-    def search(self, query: str) -> "ResultQuerySetMixin":
-        """Search the index and return a queryish object of results.
+    def search_sources(
+        self,
+        query: str,
+        *,
+        overfetch_multiplier: int | None = None,
+        max_overfetch_iterations: int | None = None,
+    ) -> "BaseStorageQuerySet":
+        """Search the index and return a queryish object of results
+        mapped back to original source objects.
         Args:
             query: The search query string
         Returns:
             ResultQuerySet instance for the search results
         """
-        return self.query_handler.search(query)
+        return self.query_handler.search_sources(
+            query,
+            overfetch_multiplier=overfetch_multiplier,
+            max_overfetch_iterations=max_overfetch_iterations,
+        )
+
+    def search_documents(self, query: str) -> "BaseStorageQuerySet":
+        """Search the index and return a queryish object of Documents
+        as stored in the underlying index.
+        Args:
+            query: The search query string
+        Returns:
+            ResultQuerySet instance for the search results
+        """
+        return self.query_handler.search_documents(query)
 
-    def find_similar(self, obj: object) -> "ResultQuerySetMixin":
+    def find_similar(self, obj: object) -> "BaseStorageQuerySet":
         """Find objects similar to the given object.
         Args:
             obj: The object to find similar objects to