Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ authors = [
requires-python = ">=3.11"
dependencies = [
"any-llm-sdk==0.20.3",
"queryish>=0.2.0",
"queryish>=0.3.0",
]
classifiers = [
"Development Status :: 3 - Alpha",
Expand Down
10 changes: 5 additions & 5 deletions src/django_ai_core/contrib/index/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/django_ai_core/contrib/index/storage/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
4 changes: 1 addition & 3 deletions src/django_ai_core/contrib/index/storage/inmemory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/contrib/index/test_query_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]


Expand Down
1,319 changes: 726 additions & 593 deletions uv.lock

Large diffs are not rendered by default.