diff --git a/src/core/config.py b/src/core/config.py index 67e0e20..2cb02ea 100644 --- a/src/core/config.py +++ b/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)] diff --git a/src/rag/code_retriever.py b/src/rag/code_retriever.py index 9fcb29e..60c997c 100644 --- a/src/rag/code_retriever.py +++ b/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: diff --git a/src/rag/indexing.py b/src/rag/indexing.py new file mode 100644 index 0000000..80c91e3 --- /dev/null +++ b/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) diff --git a/src/rag/standards_retriever.py b/src/rag/standards_retriever.py index 0d91091..4a845bc 100644 --- a/src/rag/standards_retriever.py +++ b/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() diff --git a/tests/test_indexing.py b/tests/test_indexing.py new file mode 100644 index 0000000..98bfb3a --- /dev/null +++ b/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()