Feature/indexing rate limit hardening

src/core/config.py

@@ -36,6 +36,28 @@ class Settings(BaseModel):
         ge=0.0,
         le=1.0,
     )
+    indexing_batch_size: int = Field(
+        default_factory=lambda: int(os.getenv("INDEXING_BATCH_SIZE", "12")),
+        ge=1,
+        le=200,
+    )
+    indexing_batch_sleep_ms: int = Field(
+        default_factory=lambda: int(os.getenv("INDEXING_BATCH_SLEEP_MS", "750")),
+        ge=0,
+        le=60000,
+    )
+    indexing_max_retries: int = Field(
+        default_factory=lambda: int(os.getenv("INDEXING_MAX_RETRIES", "5")),
+        ge=0,
+        le=20,
+    )
+    indexing_backoff_base_seconds: float = Field(
+        default_factory=lambda: float(
+            os.getenv("INDEXING_BACKOFF_BASE_SECONDS", "2.0")
+        ),
+        ge=0.0,
+        le=60.0,
+    )
 
     def require(self, *fields: str) -> None:
         missing = [field for field in fields if not getattr(self, field)]

src/rag/code_retriever.py

@@ -8,7 +8,8 @@
 from langchain_community.document_loaders.generic import GenericLoader
 from langchain_community.document_loaders.parsers import LanguageParser
 
-from src.core.config import DB_DIR
+from src.core.config import DB_DIR, load_settings
+from src.rag.indexing import batched, run_rate_limited_write, sleep_between_batches
 from src.rag.vector_store import get_embeddings, normalize_lang
 
 
@@ -18,6 +19,7 @@ def get_content_hash(content: str) -> str:
 
 class CodeRetriever:
     def __init__(self):
+        self.settings = load_settings()
         self.embeddings = get_embeddings()
 
         self.db = Chroma(
@@ -45,7 +47,7 @@ def index_repository(self, path: str, repo_id: str):
 
         docs = loader.load()
         seen_sources: set[str] = set()
-        added_or_updated = 0
+        docs_to_add = []
 
         for doc in docs:
             source = _normalize_source(path, doc.metadata["source"])
@@ -68,8 +70,9 @@ def index_repository(self, path: str, repo_id: str):
             doc.metadata["lang"] = normalize_lang(ext)
             doc.metadata["repo_id"] = repo_id
             doc.metadata["content_hash"] = content_hash
-            self.db.add_documents([doc])
-            added_or_updated += 1
+            docs_to_add.append(doc)
+
+        added_or_updated = self._add_documents_in_batches(docs_to_add)
 
         self._cleanup_deleted_sources(repo_id, seen_sources)
         print(f"Indexed {added_or_updated} updated/new code documents for {repo_id}")
@@ -102,6 +105,17 @@ def _cleanup_deleted_sources(self, repo_id: str, seen_sources: set[str]) -> None
         if stale_ids:
             self.db.delete(ids=stale_ids)
 
+    def _add_documents_in_batches(self, docs_to_add: list) -> int:
+        if not docs_to_add:
+            return 0
+
+        total_written = 0
+        for batch in batched(docs_to_add, self.settings.indexing_batch_size):
+            run_rate_limited_write(lambda batch=batch: self.db.add_documents(batch))
+            total_written += len(batch)
+            sleep_between_batches()
+        return total_written
+
 
 def _normalize_source(root_path: str, source: str) -> str:
     try:

src/rag/indexing.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import time
+from typing import Callable, Iterable, Sequence, TypeVar
+
+from tenacity import RetryCallState, retry, retry_if_exception, stop_after_attempt, wait_exponential
+
+from src.core.config import load_settings
+
+
+T = TypeVar("T")
+
+
+def batched(items: Sequence[T], batch_size: int) -> list[list[T]]:
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive")
+    return [list(items[index : index + batch_size]) for index in range(0, len(items), batch_size)]
+
+
+def is_rate_limit_exception(exc: BaseException) -> bool:
+    message = str(exc).lower()
+    return "429" in message or "rate limit" in message or "resource exhausted" in message
+
+
+def _before_sleep(retry_state: RetryCallState) -> None:
+    exception = retry_state.outcome.exception() if retry_state.outcome else None
+    print(
+        "Rate limited during indexing write; retrying "
+        f"(attempt {retry_state.attempt_number}) after error: {exception}"
+    )
+
+
+def run_rate_limited_write(write_operation: Callable[[], None]) -> None:
+    settings = load_settings()
+
+    @retry(
+        retry=retry_if_exception(is_rate_limit_exception),
+        stop=stop_after_attempt(settings.indexing_max_retries + 1),
+        wait=wait_exponential(multiplier=settings.indexing_backoff_base_seconds, min=settings.indexing_backoff_base_seconds),
+        before_sleep=_before_sleep,
+        reraise=True,
+    )
+    def _wrapped_write() -> None:
+        write_operation()
+
+    _wrapped_write()
+
+
+def sleep_between_batches() -> None:
+    settings = load_settings()
+    if settings.indexing_batch_sleep_ms > 0:
+        time.sleep(settings.indexing_batch_sleep_ms / 1000)

src/rag/standards_retriever.py

@@ -5,7 +5,8 @@
 
 from langchain_chroma import Chroma
 
-from src.core.config import DB_DIR
+from src.core.config import DB_DIR, load_settings
+from src.rag.indexing import batched, run_rate_limited_write, sleep_between_batches
 from src.rag.vector_store import get_embeddings, normalize_lang
 
 
@@ -15,6 +16,7 @@ def get_content_hash(content: str) -> str:
 
 class StandardsRetriever:
     def __init__(self):
+        self.settings = load_settings()
         self.embeddings = get_embeddings()
 
         self.db = Chroma(
@@ -29,7 +31,7 @@ def index_standards(self, path: str, repo_id: str):
             return
 
         seen_files = set()
-        updated_count = 0
+        docs_to_add = []
 
         print(f"Checking standards for repo: {repo_id}")
 
@@ -55,18 +57,19 @@ def index_standards(self, path: str, repo_id: str):
                 if existing_ids:
                     self.db.delete(ids=existing_ids)
 
-            self.db.add_texts(
-                texts=[content],
-                metadatas=[
-                    {
+            docs_to_add.append(
+                {
+                    "text": content,
+                    "metadata": {
                         "lang": infer_standard_language(file),
                         "repo_id": repo_id,
                         "source": file,
                         "content_hash": content_hash,
-                    }
-                ],
+                    },
+                }
             )
-            updated_count += 1
+
+        updated_count = self._add_texts_in_batches(docs_to_add)
 
         self._cleanup_deleted_files(repo_id, seen_files)
         print(f"Indexed {updated_count} updated/new standard files.")
@@ -104,6 +107,22 @@ def get_relevant_rules(self, query: str, file_ext: str, repo_id: str, k: int = 4
 
         return "\n---\n".join([doc.page_content for doc in results])
 
+    def _add_texts_in_batches(self, docs_to_add: list[dict]) -> int:
+        if not docs_to_add:
+            return 0
+
+        total_written = 0
+        for batch in batched(docs_to_add, self.settings.indexing_batch_size):
+            run_rate_limited_write(
+                lambda batch=batch: self.db.add_texts(
+                    texts=[item["text"] for item in batch],
+                    metadatas=[item["metadata"] for item in batch],
+                )
+            )
+            total_written += len(batch)
+            sleep_between_batches()
+        return total_written
+
 
 def infer_standard_language(filename: str) -> str:
     filename_lower = filename.lower()

tests/test_indexing.py

@@ -0,0 +1,26 @@
+import unittest
+
+from src.rag.indexing import batched, is_rate_limit_exception
+
+
+class IndexingHelperTests(unittest.TestCase):
+    def test_batched_splits_items_into_configurable_chunks(self):
+        items = list(range(7))
+
+        result = batched(items, 3)
+
+        self.assertEqual([[0, 1, 2], [3, 4, 5], [6]], result)
+
+    def test_batched_requires_positive_batch_size(self):
+        with self.assertRaises(ValueError):
+            batched([1, 2], 0)
+
+    def test_is_rate_limit_exception_detects_common_provider_messages(self):
+        self.assertTrue(is_rate_limit_exception(Exception("429 Too Many Requests")))
+        self.assertTrue(is_rate_limit_exception(Exception("Resource exhausted for quota")))
+        self.assertTrue(is_rate_limit_exception(Exception("Rate limit exceeded")))
+        self.assertFalse(is_rate_limit_exception(Exception("connection reset by peer")))
+
+
+if __name__ == "__main__":
+    unittest.main()