feat: add client-side rate limiting

src/dynavec/config.py

@@ -91,6 +91,8 @@ class DynavecConfig:
     max_workers: int = 8
     parallel_writes: bool = True
     max_pool_connections: int | None = None
+    put_rps: float | None = None
+    query_rps: float | None = None
 
     # in-memory hot tier (optional): keep a hot working set in RAM so warmed
     # namespaces are served entirely from memory — no S3 Vectors query and no
@@ -158,3 +160,11 @@ def __post_init__(self) -> None:
         if self.dimension > 4096:
             logger = logging.getLogger(__name__)
             logger.warning("Amazon S3 Vectors currently supports a maximum embedding dimension of 4096. You have configured a dimension of %d. This may result in an API error during provisioning or writing data.", self.dimension)
+
+        if self.put_rps is not None and self.put_rps <= 0:
+            raise ValueError("put_rps must be positive")
+
+        if self.query_rps is not None and self.query_rps <= 0:
+            raise ValueError("query_rps must be positive")
+
+

src/dynavec/stores/s3vectors.py

@@ -18,7 +18,7 @@
 
 from ..config import DynavecConfig
 from ..logging import log_store_event
-from ..utils import retry
+from ..utils import TokenBucket, retry
 
 Metadata = dict[str, Any]
 
@@ -45,6 +45,12 @@ def __init__(self, config: DynavecConfig, boto_session: Any | None = None) -> No
 
         session = boto_session or boto3.Session()
         self._config = config
+        self._put_limiter = (
+            TokenBucket(config.put_rps) if config.put_rps is not None else None
+        )
+        self._query_limiter = (
+            TokenBucket(config.query_rps) if config.query_rps is not None else None
+        )
         client_kwargs: dict[str, object] = {"region_name": config.region}
         botocore_config = config.botocore_config()
         if botocore_config is not None:
@@ -61,6 +67,8 @@ def get_index(self) -> dict[str, Any]:
 
     @retry()
     def _put_batch(self, payload: list[dict[str, Any]]) -> None:
+        if self._put_limiter is not None:
+            self._put_limiter.acquire()
         self._client.put_vectors(
             vectorBucketName=self._config.vector_bucket,
             indexName=self._config.index,
@@ -184,7 +192,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 as the S3 Vectors paginator returns them.
 
         This is what powers streaming search: an agent can begin consuming the
         first page while later pages are still in flight.
@@ -198,6 +206,7 @@ def query_pages(
         """
         if top_k <= 0:
             return
+
         if top_k > _MAX_TOP_K:
             raise ValueError(
                 f"top_k ({top_k}) exceeds Amazon S3 Vectors maximum limit of {_MAX_TOP_K}."
@@ -206,45 +215,74 @@ def query_pages(
         effective_page_size = (
             page_size if page_size is not None else self._config.top_k_page_size
         )
+
         if effective_page_size is not None and effective_page_size <= 0:
             raise ValueError("page_size must be a positive integer.")
 
         kwargs = self._query_kwargs(
-            query_vector, top_k, filter, return_metadata, return_distance
+            query_vector,
+            top_k,
+            filter,
+            return_metadata,
+            return_distance,
         )
+
         paginator = self._client.get_paginator("query_vectors")
         yielded = 0
         buffer: list[dict[str, Any]] = []
 
-        for page in paginator.paginate(
-            PaginationConfig={"MaxItems": top_k},
-            **kwargs,
-        ):
+        page_iterator = iter(
+            paginator.paginate(
+                PaginationConfig={"MaxItems": top_k},
+                **kwargs,
+            )
+        )
+
+        while True:
+            if self._query_limiter is not None:
+                self._query_limiter.acquire()
+
+            try:
+                page = next(page_iterator)
+            except StopIteration:
+                break
+
             vectors = page.get("vectors", [])
-            if not vectors:
-                continue
 
-            if effective_page_size is None:
-                remaining = top_k - yielded
-                if len(vectors) > remaining:
-                    vectors = vectors[:remaining]
-                yield vectors
-                yielded += len(vectors)
-                if yielded >= top_k:
-                    return
-            else:
-                buffer.extend(vectors)
-                while len(buffer) >= effective_page_size and yielded < top_k:
-                    chunk = buffer[:effective_page_size]
-                    buffer = buffer[effective_page_size:]
+            if vectors:
+                if effective_page_size is None:
                     remaining = top_k - yielded
-                    if len(chunk) > remaining:
-                        chunk = chunk[:remaining]
-                    yield chunk
-                    yielded += len(chunk)
+
+                    if len(vectors) > remaining:
+                        vectors = vectors[:remaining]
+
+                    yield vectors
+                    yielded += len(vectors)
+
                     if yielded >= top_k:
                         return
 
+                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]
+
+                        yield chunk
+                        yielded += len(chunk)
+
+                        if yielded >= top_k:
+                            return
+
+            if page.get("NextToken") is None:
+                break
+
         if effective_page_size is not None and buffer and yielded < top_k:
             remaining = top_k - yielded
             yield buffer[:remaining]

src/dynavec/utils.py

@@ -8,6 +8,7 @@
 
 import functools
 import random
+import threading
 import time
 from collections.abc import Iterable, Iterator
 from typing import Any, Callable, TypeVar
@@ -120,3 +121,44 @@ def chunked(iterable: Iterable[T], size: int) -> Iterator[list[T]]:
             batch = []
     if batch:
         yield batch
+
+
+class TokenBucket:
+    def __init__(
+        self,
+        rate: float,
+        capacity: float | None = None,
+    ) -> None:
+        if rate <= 0:
+            raise ValueError("rate must be positive")
+
+        self.rate = rate
+        self.capacity = capacity if capacity is not None else rate
+
+        if self.capacity <= 0:
+            raise ValueError("capacity must be positive")
+
+        self.tokens = self.capacity
+        self.last_time = time.monotonic()
+        self._lock = threading.Lock()
+
+
+    def acquire(self) -> None:
+        while True:
+            with self._lock:
+                now = time.monotonic()
+                elapsed = now - self.last_time
+
+                self.tokens = min(
+                    self.capacity,
+                    self.tokens + elapsed * self.rate,
+                )
+                self.last_time = now
+
+                if self.tokens >= 1:
+                    self.tokens -= 1
+                    return
+
+                wait_time = (1 - self.tokens) / self.rate
+
+            time.sleep(wait_time)

tests/test_connection_pool.py

@@ -30,6 +30,18 @@ def test_max_pool_connections_validation():
         _config(max_pool_connections=-5)
 
 
+def test_request_rps_defaults_to_none():
+    assert _config().put_rps is None
+    assert _config().query_rps is None
+
+
+def test_request_rps_validation():
+    with pytest.raises(ValueError):
+        _config(put_rps=0)
+    with pytest.raises(ValueError):
+        _config(query_rps=-5)
+
+
 def test_botocore_config_carries_pool_size():
     cfg = _config(max_pool_connections=50)
     client_config = cfg.botocore_config()