Use queryish start/stop properties

Author: tomusher

src/django_ai_core/contrib/index/query.py

@@ -59,8 +59,8 @@ class SourceResultMixin(BaseResultMixin):
     Mixin for source search - returns unique source objects with over-fetching.
 
     This mixin overrides run_query() to implement over-fetching
-    and deduplication. By the time run_query() is called, self.limit and
-    self.offset have been set by slicing operations.
+    and deduplication. By the time run_query() is called, self.start and
+    self.stop have been set by slicing operations.
     """
 
     overfetch_multiplier: int = 3
@@ -80,8 +80,8 @@ def run_query(self):
         """
         Execute query with over-fetching and deduplication.
         """
-        requested_limit = getattr(self, "limit", None) or 20
-        requested_offset = getattr(self, "offset", None) or 0
+        requested_limit = self.limit
+        requested_offset = self.offset or 0
 
         yield from self._overfetch(requested_offset, requested_limit)
 
@@ -126,7 +126,7 @@ def _fetch_batch(self, limit: int):
         This method calls the parent class's (storage provider's) run_query()
         with specific offset and limit.
         """
-        temp_qs = self.clone(limit=limit)  # type: ignore
+        temp_qs = self.clone(start=0, stop=limit)  # type: ignore
 
         # Skip SourceResultMixin in the MRO of temp_qs to resolve
         # run_query to the method on the storage provided queryish object

src/django_ai_core/contrib/index/storage/base.py

@@ -21,7 +21,7 @@ class BaseStorageQuerySet(Queryish, Generic[StorageProviderType]):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.limit = 20
+        self.stop = 20
 
     def run_query(self) -> Iterator["BaseStorageDocument"]:
         """Execute the query and return the results."""

src/django_ai_core/contrib/index/storage/inmemory.py

@@ -11,8 +11,6 @@ def run_query(self):
         storage_provider = self.storage_provider
         filter_map = {filter[0]: filter[1] for filter in self.filters}
 
-        limit = self.limit or 10
-
         embedding = filter_map.pop("embedding", None)
         if embedding is None:
             raise ValueError("embedding filter is required")
@@ -31,7 +29,7 @@ def run_query(self):
             similarities, key=lambda pair: pair[0], reverse=True
         )
         for document in [pair[1] for pair in sorted_similarities][
-            self.offset : self.offset + limit
+            self.start : self.stop
         ]:
             yield document
 

src/django_ai_core/contrib/index/storage/pgvector/provider.py

@@ -48,7 +48,7 @@ def run_query(self) -> Generator[BaseStorageDocument, None, None]:
         for key, value in filter_map.items():
             queryset = queryset.filter(**{f"metadata__{key}": value})
 
-        queryset = queryset[self.offset : self.limit]
+        queryset = queryset[self.start : self.stop]
 
         for instance in queryset:
             yield self.get_instance(instance)

tests/unit/contrib/index/test_query_handler.py

@@ -101,7 +101,7 @@ def __init__(self, documents=None):
 
     def run_query(self):
         """Mock run_query that returns documents with respect to limit."""
-        limit = self.limit or len(self._documents)
+        limit = self.stop or len(self._documents)
         yield from self._documents[:limit]