diff --git a/src/dynavec/stores/async_s3vectors.py b/src/dynavec/stores/async_s3vectors.py index f3fd66e..db3fc62 100644 --- a/src/dynavec/stores/async_s3vectors.py +++ b/src/dynavec/stores/async_s3vectors.py @@ -257,7 +257,7 @@ async def query_pages( return_distance: bool = True, page_size: int | None = None, ) -> AsyncIterator[list[dict[str, Any]]]: - """Yield result pages from the async S3 Vectors paginator.""" + """Yield async result pages, following S3 Vectors ``nextToken``.""" if top_k <= 0: return @@ -289,22 +289,13 @@ async def query_pages( yielded = 0 buffer: list[dict[str, Any]] = [] - with translate_client_errors(): - paginator = client.get_paginator("query_vectors") - page_iterator = paginator.paginate( - PaginationConfig={"MaxItems": top_k}, - **kwargs, - ).__aiter__() - while True: if self._query_limiter is not None: await self._query_limiter.acquire_async() - try: - with translate_client_errors(): - page = await page_iterator.__anext__() - except StopAsyncIteration: - break + # Older supported SDKs expose QueryVectors without a paginator. + with translate_client_errors(): + page = await client.query_vectors(**kwargs) vectors = page.get("vectors", []) @@ -325,16 +316,12 @@ async def query_pages( buffer.extend(vectors) while ( - len(buffer) >= effective_page_size - and yielded < top_k + yielded < top_k + and len(buffer) >= min(effective_page_size, top_k - yielded) ): - chunk = buffer[:effective_page_size] - buffer = buffer[effective_page_size:] - - remaining = top_k - yielded - - if len(chunk) > remaining: - chunk = chunk[:remaining] + chunk_size = min(effective_page_size, top_k - yielded) + chunk = buffer[:chunk_size] + buffer = buffer[chunk_size:] yield chunk yielded += len(chunk) @@ -342,8 +329,10 @@ async def query_pages( if yielded >= top_k: return - if page.get("NextToken") is None: + next_token = page.get("nextToken") + if not next_token: break + kwargs["nextToken"] = next_token if ( effective_page_size is not None diff --git a/src/dynavec/stores/s3vectors.py b/src/dynavec/stores/s3vectors.py index 79fb276..7de1bdc 100644 --- a/src/dynavec/stores/s3vectors.py +++ b/src/dynavec/stores/s3vectors.py @@ -195,7 +195,7 @@ def query_pages( return_distance: bool = True, page_size: int | None = None, ) -> Iterator[list[dict[str, Any]]]: - """Yield result pages as the S3 Vectors paginator returns them. + """Yield result pages from S3 Vectors, following ``nextToken``. This is what powers streaming search: an agent can begin consuming the first page while later pages are still in flight. @@ -233,24 +233,13 @@ def query_pages( yielded = 0 buffer: list[dict[str, Any]] = [] - with translate_client_errors(): - paginator = self._client.get_paginator("query_vectors") - page_iterator = iter( - paginator.paginate( - PaginationConfig={"MaxItems": top_k}, - **kwargs, - ) - ) - while True: if self._query_limiter is not None: self._query_limiter.acquire() - try: - with translate_client_errors(): - page = next(page_iterator) - except StopIteration: - break + # Older supported SDKs expose QueryVectors without a paginator. + with translate_client_errors(): + page = self._client.query_vectors(**kwargs) vectors = page.get("vectors", []) @@ -270,14 +259,12 @@ def query_pages( else: buffer.extend(vectors) - while len(buffer) >= effective_page_size and yielded < top_k: - chunk = buffer[:effective_page_size] - buffer = buffer[effective_page_size:] - - remaining = top_k - yielded - - if len(chunk) > remaining: - chunk = chunk[:remaining] + while yielded < top_k and len(buffer) >= min( + effective_page_size, top_k - yielded + ): + chunk_size = min(effective_page_size, top_k - yielded) + chunk = buffer[:chunk_size] + buffer = buffer[chunk_size:] yield chunk yielded += len(chunk) @@ -285,8 +272,10 @@ def query_pages( if yielded >= top_k: return - if page.get("NextToken") is None: + next_token = page.get("nextToken") + if not next_token: break + kwargs["nextToken"] = next_token if effective_page_size is not None and buffer and yielded < top_k: remaining = top_k - yielded diff --git a/tests/test_async_s3vectors.py b/tests/test_async_s3vectors.py index 9a46cba..5c95174 100644 --- a/tests/test_async_s3vectors.py +++ b/tests/test_async_s3vectors.py @@ -6,24 +6,14 @@ from typing import Any from unittest.mock import AsyncMock +import aioboto3 import pytest +from aiobotocore.stub import AioStubber from dynavec.config import DynavecConfig from dynavec.stores.async_s3vectors import AsyncS3VectorsStore -class FakeAsyncPaginator: - def __init__(self, pages: list[dict[str, Any]]) -> None: - self.pages = pages - self.calls: list[dict[str, Any]] = [] - - async def paginate(self, **kwargs: Any): - self.calls.append(kwargs) - - for page in self.pages: - yield page - - def _config() -> DynavecConfig: return DynavecConfig( vector_bucket="test-bucket", @@ -71,14 +61,15 @@ def __init__( ) -> None: self.put_calls: list[dict[str, Any]] = [] self.get_calls: list[dict[str, Any]] = [] - self.paginator = FakeAsyncPaginator(pages or []) + self.pages = iter(pages or [{"vectors": []}]) + self.query_calls: list[dict[str, Any]] = [] async def put_vectors(self, **kwargs: Any) -> None: self.put_calls.append(kwargs) - def get_paginator(self, name: str) -> FakeAsyncPaginator: - assert name == "query_vectors" - return self.paginator + async def query_vectors(self, **kwargs: Any) -> dict[str, Any]: + self.query_calls.append(kwargs) + return next(self.pages) async def get_vectors(self, **kwargs: Any) -> dict[str, Any]: self.get_calls.append(kwargs) @@ -188,7 +179,7 @@ async def test_query_returns_paginated_results(): {"key": f"k{i}", "distance": i * 0.01} for i in range(100) ], - "NextToken": "token-1", + "nextToken": "token-1", }, { "vectors": [ @@ -210,6 +201,26 @@ async def test_query_returns_paginated_results(): assert len(results) == 150 assert results[0]["key"] == "k0" assert results[-1]["key"] == "k149" + assert "nextToken" not in client.query_calls[0] + assert client.query_calls[1]["nextToken"] == "token-1" + + +async def test_query_with_real_async_sdk_client(): + session = aioboto3.Session( + aws_access_key_id="testing", aws_secret_access_key="testing", region_name="us-east-1" + ) + vectors = [{"key": "doc-1", "distance": 0.25}] + + async with AsyncS3VectorsStore(_config(), session) as store: + with AioStubber(store._client) as stubber: + stubber.add_response("query_vectors", {"vectors": vectors}, { + "vectorBucketName": "test-bucket", "indexName": "test-index", "topK": 1, + "queryVector": {"float32": [1.0, 2.0, 3.0, 4.0]}, + "returnMetadata": True, "returnDistance": True, + }) + + assert await store.query([1.0, 2.0, 3.0, 4.0], top_k=1) == vectors + stubber.assert_no_pending_responses() async def test_query_truncates_at_top_k(): @@ -219,7 +230,7 @@ async def test_query_truncates_at_top_k(): {"key": f"k{i}"} for i in range(100) ], - "NextToken": "token-1", + "nextToken": "token-1", }, { "vectors": [ @@ -283,6 +294,49 @@ async def test_query_rejects_excessive_top_k(): ) +@pytest.mark.parametrize("page_size", [None, 2]) +async def test_query_pages_continues_after_empty_page_with_next_token(page_size): + client = FakeS3Client([ + {"vectors": [{"key": "k0"}], "nextToken": "token-1"}, + {"vectors": [], "nextToken": "token-2"}, + {"vectors": [{"key": "k1"}, {"key": "k2"}]}, + ]) + limiter = AsyncMock() + + async with AsyncS3VectorsStore(_config(), FakeSession(client)) as store: + store._query_limiter = limiter + pages = [ + page async for page in store.query_pages( + [0.1, 0.2, 0.3, 0.4], top_k=5, page_size=page_size + ) + ] + + assert [hit["key"] for page in pages for hit in page] == ["k0", "k1", "k2"] + assert [len(page) for page in pages] == ([1, 2] if page_size is None else [2, 1]) + assert limiter.acquire_async.await_count == 3 + + +async def test_query_pages_stops_at_top_k_with_large_page_size(): + client = FakeS3Client([ + {"vectors": [{"key": "k0"}], "nextToken": "token-1"}, + {"vectors": [{"key": "k1"}, {"key": "k2"}], "nextToken": "unused"}, + ]) + + async with AsyncS3VectorsStore(_config(), FakeSession(client)) as store: + pages = [page async for page in store.query_pages( + [1.0, 2.0, 3.0, 4.0], top_k=2, page_size=100, + filter={"kind": "test"}, return_metadata=False, return_distance=False, + )] + + assert pages == [[{"key": "k0"}, {"key": "k1"}]] + assert len(client.query_calls) == 2 + for call in client.query_calls: + assert call["topK"] == 2 + assert call["filter"] == {"kind": "test"} + assert call["returnMetadata"] is False + assert call["returnDistance"] is False + + async def test_put_batch_uses_rate_limiter(): config = DynavecConfig( vector_bucket="test-bucket", @@ -319,11 +373,11 @@ async def test_query_pages_uses_rate_limiter_per_page(): pages = [ { "vectors": [{"key": "k0"}], - "NextToken": "token-1", + "nextToken": "token-1", }, { "vectors": [{"key": "k1"}], - "NextToken": "token-2", + "nextToken": "token-2", }, { "vectors": [{"key": "k2"}], diff --git a/tests/test_logging.py b/tests/test_logging.py index 90a4311..08c3459 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -130,10 +130,9 @@ def test_dynamodb_store_structured_logging(caplog): def test_s3vectors_store_structured_logging(caplog): session = MagicMock() client = session.client.return_value - paginator = client.get_paginator.return_value - paginator.paginate.return_value = [ - {"vectors": [{"key": "ns#1", "data": {"float32": [0.1] * 8}, "distance": 0.05}]} - ] + client.query_vectors.return_value = { + "vectors": [{"key": "ns#1", "data": {"float32": [0.1] * 8}, "distance": 0.05}] + } # 1. Opt-in: When False, no logs cfg_disabled = DynavecConfig( diff --git a/tests/test_s3vectors.py b/tests/test_s3vectors.py index 32e1797..48eb538 100644 --- a/tests/test_s3vectors.py +++ b/tests/test_s3vectors.py @@ -1,7 +1,10 @@ import time +from contextlib import closing from unittest.mock import MagicMock +import boto3 import pytest +from botocore.stub import Stubber from dynavec.config import DynavecConfig from dynavec.stores.s3vectors import _MAX_TOP_K, S3VectorsStore @@ -49,34 +52,50 @@ def _store_with_pages(pages, config=None): store._put_limiter = None store._query_limiter = None - paginator = MagicMock() - paginator.paginate.return_value = pages - store._client.get_paginator.return_value = paginator - return store, paginator + store._client.query_vectors.side_effect = pages or [{"vectors": []}] + return store, store._client.query_vectors def test_query_single_page(): page = [{"key": f"k{i}", "distance": 0.01 * i} for i in range(50)] - store, paginator = _store_with_pages([{"vectors": page}]) + store, query = _store_with_pages([{"vectors": page}]) results = store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=50) assert len(results) == 50 assert results[0]["key"] == "k0" assert results[-1]["key"] == "k49" - store._client.get_paginator.assert_called_once_with("query_vectors") - assert paginator.paginate.call_args.kwargs["PaginationConfig"] == {"MaxItems": 50} - assert paginator.paginate.call_args.kwargs["topK"] == 50 + store._client.get_paginator.assert_not_called() + assert query.call_args.kwargs["topK"] == 50 + assert "nextToken" not in query.call_args.kwargs + + +def test_query_with_real_sdk_client(): + session = boto3.Session( + aws_access_key_id="testing", aws_secret_access_key="testing", region_name="us-east-1" + ) + store = S3VectorsStore(_config(), session) + vectors = [{"key": "doc-1", "distance": 0.25}] + + with closing(store._client), Stubber(store._client) as stubber: + stubber.add_response("query_vectors", {"vectors": vectors}, { + "vectorBucketName": "test-bucket", "indexName": "test-index", "topK": 1, + "queryVector": {"float32": [1.0, 2.0, 3.0, 4.0]}, + "returnMetadata": True, "returnDistance": True, + }) + + assert store.query([1.0, 2.0, 3.0, 4.0], top_k=1) == vectors + stubber.assert_no_pending_responses() def test_query_multi_page_pagination(): pages = [ { "vectors": [{"key": f"k{i}"} for i in range(100)], - "NextToken": "token-1", + "nextToken": "token-1", }, { "vectors": [{"key": f"k{i}"} for i in range(100, 200)], - "NextToken": "token-2", + "nextToken": "token-2", }, { "vectors": [{"key": f"k{i}"} for i in range(200, 250)], @@ -97,6 +116,9 @@ def test_query_multi_page_pagination(): "k248", "k249", ] + calls = store._client.query_vectors.call_args_list + assert "nextToken" not in calls[0].kwargs + assert [call.kwargs["nextToken"] for call in calls[1:]] == ["token-1", "token-2"] def test_query_fewer_results_than_top_k(): store, _ = _store_with_pages([{"vectors": [{"key": "k0"}, {"key": "k1"}]}]) @@ -111,7 +133,7 @@ def test_query_zero_or_negative_top_k(): assert store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=0) == [] assert store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=-5) == [] - store._client.get_paginator.assert_not_called() + store._client.query_vectors.assert_not_called() def test_query_exceeding_service_limit(): @@ -200,7 +222,7 @@ def test_query_truncation_at_top_k(): pages = [ { "vectors": [{"key": f"k{i}"} for i in range(100)], - "NextToken": "token-1", + "nextToken": "token-1", }, { "vectors": [{"key": f"k{i}"} for i in range(100, 200)], @@ -222,7 +244,7 @@ def test_query_pages_raw_chunks(): pages = [ { "vectors": [{"key": f"k{i}"} for i in range(100)], - "NextToken": "token-1", + "nextToken": "token-1", }, { "vectors": [{"key": f"k{i}"} for i in range(100, 150)], @@ -273,15 +295,52 @@ def test_query_pages_uses_config_top_k_page_size(): assert [len(p) for p in result_pages] == [10, 10, 5] +@pytest.mark.parametrize("page_size", [None, 2]) +def test_query_pages_continues_after_empty_page_with_next_token(page_size): + store, _ = _store_with_pages([ + {"vectors": [{"key": "k0"}], "nextToken": "token-1"}, + {"vectors": [], "nextToken": "token-2"}, + {"vectors": [{"key": "k1"}, {"key": "k2"}]}, + ]) + limiter = MagicMock() + store._query_limiter = limiter + + pages = list(store.query_pages([0.1, 0.2, 0.3, 0.4], top_k=5, page_size=page_size)) + + assert [hit["key"] for page in pages for hit in page] == ["k0", "k1", "k2"] + assert [len(page) for page in pages] == ([1, 2] if page_size is None else [2, 1]) + assert limiter.acquire.call_count == 3 + + +def test_query_pages_stops_at_top_k_with_large_page_size(): + store, query = _store_with_pages([ + {"vectors": [{"key": "k0"}], "nextToken": "token-1"}, + {"vectors": [{"key": "k1"}, {"key": "k2"}], "nextToken": "unused"}, + ]) + + pages = list(store.query_pages( + [1.0, 2.0, 3.0, 4.0], top_k=2, page_size=100, + filter={"kind": "test"}, return_metadata=False, return_distance=False, + )) + + assert pages == [[{"key": "k0"}, {"key": "k1"}]] + assert query.call_count == 2 + for call in query.call_args_list: + assert call.kwargs["topK"] == 2 + assert call.kwargs["filter"] == {"kind": "test"} + assert call.kwargs["returnMetadata"] is False + assert call.kwargs["returnDistance"] is False + + def test_query_pages_uses_rate_limiter_per_page(): pages = [ { "vectors": [{"key": "k0"}], - "NextToken": "token-1", + "nextToken": "token-1", }, { "vectors": [{"key": "k1"}], - "NextToken": "token-2", + "nextToken": "token-2", }, { "vectors": [{"key": "k2"}], @@ -324,4 +383,4 @@ def test_query_pages_does_not_rate_limit_after_paginator_exhaustion(): ) assert [len(page) for page in result_pages] == [1] - assert limiter.acquire.call_count == 1 \ No newline at end of file + assert limiter.acquire.call_count == 1 diff --git a/tests/test_s3vectors_errors.py b/tests/test_s3vectors_errors.py index 5d7a2b7..2da1d75 100644 --- a/tests/test_s3vectors_errors.py +++ b/tests/test_s3vectors_errors.py @@ -42,6 +42,7 @@ def _async_store(): client = MagicMock() client.put_vectors = AsyncMock() client.get_vectors = AsyncMock() + client.query_vectors = AsyncMock() session = MagicMock() session.client.return_value.__aenter__.return_value = client return AsyncS3VectorsStore(_config(), session), client @@ -118,6 +119,7 @@ def test_new_exceptions_are_public(error_type): ("_put_batch", ([],), "put_vectors"), ("get_vectors", (["key"],), "get_vectors"), ("delete_vectors", (["key"],), "delete_vectors"), + ("query", ([0.1, 0.2], 2), "query_vectors"), ], ) def test_sync_direct_operations_translate( @@ -138,7 +140,11 @@ def test_sync_direct_operations_translate( @pytest.mark.parametrize("code,error_type", MAPPINGS) @pytest.mark.parametrize( "operation,args,aws_operation", - [("_put_batch", ([],), "put_vectors"), ("get_vectors", (["key"],), "get_vectors")], + [ + ("_put_batch", ([],), "put_vectors"), + ("get_vectors", (["key"],), "get_vectors"), + ("query", ([0.1, 0.2], 2), "query_vectors"), + ], ) async def test_async_direct_operations_translate(code, error_type, operation, args, aws_operation): store, client = _async_store() @@ -155,24 +161,24 @@ async def test_async_direct_operations_translate(code, error_type, operation, ar def _failing_pages(original): - yield {"vectors": [{"key": "first"}], "NextToken": "next"} + yield {"vectors": [{"key": "first"}], "nextToken": "next"} raise original -async def _async_pages(original=None): - yield {"vectors": [{"key": "first"}], "NextToken": "next"} - if original is not None: - raise original - - @pytest.mark.parametrize("operation,args", [("query_pages", ([0.1, 0.2], 2)), ("list_pages", ())]) def test_sync_pagination_translates_later_page_errors(sync_store, operation, args): store, client = sync_store original = _client_error("NotFoundException") - client.get_paginator.return_value.paginate.return_value = _failing_pages(original) + if operation == "query_pages": + client.query_vectors.side_effect = [ + {"vectors": [{"key": "first"}], "nextToken": "next"}, original, + ] + else: + client.get_paginator.return_value.paginate.return_value = _failing_pages(original) pages = getattr(store, operation)(*args) client.get_paginator.assert_not_called() + client.query_vectors.assert_not_called() assert next(pages) == [{"key": "first"}] with pytest.raises(NotFoundError) as caught: next(pages) @@ -183,11 +189,13 @@ def test_sync_pagination_translates_later_page_errors(sync_store, operation, arg async def test_async_pagination_translates_later_page_errors(): store, client = _async_store() original = _client_error("NotFoundException") - client.get_paginator.return_value.paginate.return_value = _async_pages(original) + client.query_vectors.side_effect = [ + {"vectors": [{"key": "first"}], "nextToken": "next"}, original, + ] async with store: pages = store.query_pages([0.1, 0.2], 2) - client.get_paginator.assert_not_called() + client.query_vectors.assert_not_called() assert await pages.__anext__() == [{"key": "first"}] with pytest.raises(NotFoundError) as caught: await pages.__anext__() @@ -195,9 +203,8 @@ async def test_async_pagination_translates_later_page_errors(): assert caught.value.__cause__ is original -@pytest.mark.parametrize("operation,args", [("query", ([0.1, 0.2], 2)), ("list_pages", ())]) @pytest.mark.parametrize("boundary", ["get_paginator", "paginate"]) -def test_sync_paginator_setup_errors_translate(sync_store, operation, args, boundary): +def test_sync_list_paginator_setup_errors_translate(sync_store, boundary): store, client = sync_store original = _client_error("ValidationException") call = ( @@ -208,26 +215,7 @@ def test_sync_paginator_setup_errors_translate(sync_store, operation, args, boun call.side_effect = original with pytest.raises(ValidationError) as caught: - list(getattr(store, operation)(*args)) - - assert caught.value.__cause__ is original - call.assert_called_once() - - -@pytest.mark.parametrize("boundary", ["get_paginator", "paginate"]) -async def test_async_paginator_setup_errors_translate(boundary): - store, client = _async_store() - original = _client_error("ValidationException") - call = ( - client.get_paginator - if boundary == "get_paginator" - else client.get_paginator.return_value.paginate - ) - call.side_effect = original - - async with store: - with pytest.raises(ValidationError) as caught: - await store.query([0.1, 0.2], 2) + list(store.list_pages()) assert caught.value.__cause__ is original call.assert_called_once() @@ -287,32 +275,36 @@ async def test_async_batch_retries_then_succeeds(monkeypatch, code): def test_sync_query_retries_page_iteration_failure(sync_store, monkeypatch): store, client = sync_store - client.get_paginator.return_value.paginate.side_effect = [ - _failing_pages(_client_error("ThrottlingException")), - [{"vectors": [{"key": "success"}]}], + client.query_vectors.side_effect = [ + {"vectors": [{"key": "first"}], "nextToken": "next"}, + _client_error("ThrottlingException"), + {"vectors": [{"key": "success"}]}, ] sleep = MagicMock() monkeypatch.setattr("dynavec.utils.time.sleep", sleep) assert store.query([0.1, 0.2], 2) == [{"key": "success"}] - assert client.get_paginator.call_count == 2 + assert client.query_vectors.call_count == 3 + assert "nextToken" not in client.query_vectors.call_args.kwargs sleep.assert_called_once() async def test_async_query_retries_page_iteration_failure(monkeypatch): store, client = _async_store() - client.get_paginator.return_value.paginate.side_effect = [ - _async_pages(_client_error("ThrottlingException")), - _async_pages(), + client.query_vectors.side_effect = [ + {"vectors": [{"key": "first"}], "nextToken": "next"}, + _client_error("ThrottlingException"), + {"vectors": [{"key": "success"}]}, ] sleep = AsyncMock() monkeypatch.setattr("dynavec.utils.asyncio.sleep", sleep) async with store: - assert await store.query([0.1, 0.2], 2) == [{"key": "first"}] + assert await store.query([0.1, 0.2], 2) == [{"key": "success"}] - assert client.get_paginator.call_count == 2 + assert client.query_vectors.await_count == 3 + assert "nextToken" not in client.query_vectors.call_args.kwargs sleep.assert_awaited_once()