From ad9bbab17cf3c8f25cde0814ee864fa60d70b6a5 Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sun, 13 Sep 2026 13:47:07 +0200 Subject: [PATCH 1/8] refactor: expose provider batch_size on the EmbeddingProvider protocol The indexing pipeline needs to know how many texts a provider accepts per request so it can accumulate symbols into full batches. The value already existed as a private `_BATCH_SIZE` in each provider module but was not reachable through the protocol. The protocol property carries a concrete default body rather than `...`: every bundled provider subclasses EmbeddingProvider, so an ellipsis would silently return None for any provider that forgot to override it. Co-Authored-By: Claude Opus 5 (1M context) --- server/embeddings/base.py | 14 +++++++++++++ server/embeddings/jina.py | 4 ++++ server/embeddings/jina_api.py | 4 ++++ server/embeddings/ollama.py | 4 ++++ server/embeddings/openai.py | 4 ++++ server/embeddings/voyage.py | 4 ++++ tests/embeddings/test_batch_size.py | 31 +++++++++++++++++++++++++++++ uv.lock | 2 +- 8 files changed, 66 insertions(+), 1 deletion(-) create mode 100644 tests/embeddings/test_batch_size.py diff --git a/server/embeddings/base.py b/server/embeddings/base.py index 345b448..9458887 100644 --- a/server/embeddings/base.py +++ b/server/embeddings/base.py @@ -2,12 +2,26 @@ from typing import Protocol, runtime_checkable +# Conservative ceiling for providers that do not declare their own. Matches the +# smallest batch size any bundled provider uses (jina TEI, ollama). +_DEFAULT_BATCH_SIZE = 32 + @runtime_checkable class EmbeddingProvider(Protocol): @property def dimensions(self) -> int: ... + @property + def batch_size(self) -> int: + """Max texts the provider accepts per request. + + Concrete default rather than `...` — every bundled provider subclasses + this Protocol, so an ellipsis body would silently return None for any + provider that forgot to override it. + """ + return _DEFAULT_BATCH_SIZE + async def embed_batch(self, texts: list[str]) -> list[list[float]]: ... async def embed_query(self, text: str) -> list[float]: ... diff --git a/server/embeddings/jina.py b/server/embeddings/jina.py index 4ea0667..e5ccb1f 100644 --- a/server/embeddings/jina.py +++ b/server/embeddings/jina.py @@ -23,6 +23,10 @@ def __init__(self) -> None: def dimensions(self) -> int: return self._dims + @property + def batch_size(self) -> int: + return _BATCH_SIZE + @staticmethod def _extract(data) -> list[list[float]]: # TEI returns a list of vectors directly diff --git a/server/embeddings/jina_api.py b/server/embeddings/jina_api.py index 1f3654b..fa7b9dd 100644 --- a/server/embeddings/jina_api.py +++ b/server/embeddings/jina_api.py @@ -76,6 +76,10 @@ def __init__(self) -> None: def dimensions(self) -> int: return self._dims + @property + def batch_size(self) -> int: + return _BATCH_SIZE + async def embed_batch(self, texts: list[str]) -> list[list[float]]: return await self._embed(texts, task="retrieval.passage") diff --git a/server/embeddings/ollama.py b/server/embeddings/ollama.py index 2d67c18..05adf41 100644 --- a/server/embeddings/ollama.py +++ b/server/embeddings/ollama.py @@ -41,6 +41,10 @@ def __init__(self) -> None: def dimensions(self) -> int: return self._dims + @property + def batch_size(self) -> int: + return _BATCH_SIZE + async def embed_batch(self, texts: list[str]) -> list[list[float]]: return await embed_in_batches( texts, diff --git a/server/embeddings/openai.py b/server/embeddings/openai.py index e77c72d..e59d112 100644 --- a/server/embeddings/openai.py +++ b/server/embeddings/openai.py @@ -51,6 +51,10 @@ def __init__(self) -> None: def dimensions(self) -> int: return self._dims + @property + def batch_size(self) -> int: + return _BATCH_SIZE + def _make_body(self, inputs: list[str]) -> dict: body: dict = {"model": self._model, "input": inputs} if self._dims_override is not None: diff --git a/server/embeddings/voyage.py b/server/embeddings/voyage.py index e4685a4..f01b7fc 100644 --- a/server/embeddings/voyage.py +++ b/server/embeddings/voyage.py @@ -57,6 +57,10 @@ def __init__(self) -> None: def dimensions(self) -> int: return self._dims + @property + def batch_size(self) -> int: + return _BATCH_SIZE + async def embed_batch(self, texts: list[str]) -> list[list[float]]: return await self._embed(texts, input_type="document") diff --git a/tests/embeddings/test_batch_size.py b/tests/embeddings/test_batch_size.py new file mode 100644 index 0000000..55a2397 --- /dev/null +++ b/tests/embeddings/test_batch_size.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import pytest + +import server.embeddings.jina as jina_module +import server.embeddings.jina_api as jina_api_module +import server.embeddings.ollama as ollama_module +import server.embeddings.openai as openai_module +import server.embeddings.voyage as voyage_module +from server.embeddings.base import _DEFAULT_BATCH_SIZE, EmbeddingProvider + +# The pipeline batches symbols up to `batch_size`, so the property must stay in +# lockstep with the `_BATCH_SIZE` each provider passes to `embed_in_batches` — +# otherwise the pipeline packs batches the provider immediately re-splits. +_PROVIDERS = [ + (voyage_module.VoyageEmbeddingProvider, voyage_module, 128), + (openai_module.OpenAIEmbeddingProvider, openai_module, 128), + (jina_api_module.JinaApiEmbeddingProvider, jina_api_module, 128), + (jina_module.JinaEmbeddingProvider, jina_module, 32), + (ollama_module.OllamaEmbeddingProvider, ollama_module, 32), +] + + +@pytest.mark.parametrize(("provider_cls", "module", "expected"), _PROVIDERS) +def test_batch_size_matches_module_constant(provider_cls, module, expected) -> None: + # Read through the descriptor — constructing a provider needs API keys. + assert provider_cls.batch_size.fget(None) == module._BATCH_SIZE == expected + + +def test_protocol_default_is_the_conservative_size() -> None: + assert EmbeddingProvider.batch_size.fget(None) == _DEFAULT_BATCH_SIZE == 32 diff --git a/uv.lock b/uv.lock index 084d493..00b28a5 100644 --- a/uv.lock +++ b/uv.lock @@ -1367,7 +1367,7 @@ wheels = [ [[package]] name = "semcode" -version = "1.3.1" +version = "1.3.2" source = { editable = "." } dependencies = [ { name = "cachetools" }, From 899b10e3c17ed5075d5c42e1ed18f6d6a48ae0cc Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sun, 13 Sep 2026 13:49:29 +0200 Subject: [PATCH 2/8] perf: batch symbols across files and fetch blobs concurrently MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The indexing loop processed one file at a time: blob fetch, then dense embed, then sparse embed, then upsert. Network round-trips never overlapped, and a single file yields only a handful of symbols, so the pipeline sent dozens-to-hundreds of tiny embedding requests to providers that accept 128 inputs each. Restructured into three stages inside index_service: - a producer task fetches and parses with a bounded semaphore, following the shape already used by fetch_commits_with_diffs - a batcher accumulates symbols across files up to the provider's batch_size and runs the dense and sparse embeds under one gather — dense is network-bound and sparse runs in a thread executor, so they overlap for free - a writer keeps upsert-before-delete per file The unit of store mutation stays one file, so per-file error isolation is unchanged. A batch-level embedding failure falls back to embedding that batch file by file, so one bad file costs an extra round-trip rather than dropping its whole batch. parse_file deliberately stays on the event loop: registry.py shares one parser instance per language and tree-sitter parsers are not safe for concurrent use. Progress `current` is now a monotone count of files resolved rather than a loop index, which reaches `total` instead of jumping over skipped files. Co-Authored-By: Claude Opus 5 (1M context) --- server/indexer/pipeline.py | 338 +++++++++++++++++++++++++++++-------- 1 file changed, 268 insertions(+), 70 deletions(-) diff --git a/server/indexer/pipeline.py b/server/indexer/pipeline.py index 25ebc3a..86f3a0e 100644 --- a/server/indexer/pipeline.py +++ b/server/indexer/pipeline.py @@ -1,22 +1,28 @@ from __future__ import annotations +import asyncio import logging import re import textwrap -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass from datetime import UTC, datetime from typing import Any import httpx +from qdrant_client.models import SparseVector -from server.config import settings +from server.config import ServiceConfig, settings from server.embeddings import get_embedding_provider from server.embeddings.base import EmbeddingProvider from server.embeddings.bm25 import BM25SparseProvider, get_sparse_embedding_provider from server.embeddings.code_tokenizer import symbol_name_tokens from server.indexer.cleanup import prune_orphaned_services -from server.indexer.github_source import fetch_blob_content, list_github_files +from server.indexer.github_source import ( + GitHubFile, + fetch_blob_content, + list_github_files, +) from server.parser.base import CodeSymbol, ParseError from server.parser.registry import parse_file from server.state import get_reindex_lock, get_service_registry @@ -25,6 +31,31 @@ logger = logging.getLogger(__name__) +# Bounded fan-out for blob downloads, matching _TREE_WALK_CONCURRENCY / +# _DIFF_CONCURRENCY in github_source.py. +_FETCH_CONCURRENCY = 10 +# Parsed files waiting to be embedded. Bounds peak memory and applies +# backpressure to the fetchers while a batch is embedding. +_PARSED_QUEUE_SIZE = 20 +# Used when the provider does not declare a batch_size (duck-typed test stubs). +_FALLBACK_EMBED_BATCH_SIZE = 32 +# Secondary cut so a batch of unusually large symbols never becomes a giant request. +_MAX_BATCH_CHARS = 400_000 + + +@dataclass(slots=True) +class _ParsedFile: + stored_path: str + blob_sha: str + symbols: list[CodeSymbol] # empty => drop stale entries, nothing to embed + + +@dataclass(slots=True) +class _EmbeddedFile: + parsed: _ParsedFile + dense: list[list[float]] + sparse: list[SparseVector] + @dataclass class ProgressEvent: @@ -157,6 +188,60 @@ def _symbol_to_payload( } +async def _drain_batches( + queue: asyncio.Queue[_ParsedFile | None], batch_size: int +) -> AsyncIterator[list[_ParsedFile]]: + """Groups parsed files until their combined symbol count fills a provider batch. + + Terminates on the ``None`` sentinel the producer always sends, flushing + whatever is still pending. A single file holding more symbols than + *batch_size* is emitted on its own; ``embed_in_batches`` re-splits it. + """ + pending: list[_ParsedFile] = [] + pending_symbols = 0 + pending_chars = 0 + + while True: + item = await queue.get() + if item is None: + break + pending.append(item) + pending_symbols += len(item.symbols) + pending_chars += sum(len(s.source or "") for s in item.symbols) + if pending_symbols >= batch_size or pending_chars >= _MAX_BATCH_CHARS: + yield pending + pending = [] + pending_symbols = 0 + pending_chars = 0 + + if pending: + yield pending + + +def _split_by_file( + files: list[_ParsedFile], + dense: list[list[float]], + sparse: list[SparseVector], +) -> list[_EmbeddedFile]: + """Cuts flat batch vectors back into per-file runs, in the order they were sent.""" + total = sum(len(f.symbols) for f in files) + if len(dense) != total or len(sparse) != total: + # Vectors are positional; a miscount would silently attach one symbol's + # vector to another symbol. + raise ValueError( + f"Embedding count mismatch: {len(dense)} dense / {len(sparse)} sparse " + f"vectors for {total} symbols" + ) + + embedded: list[_EmbeddedFile] = [] + offset = 0 + for f in files: + end = offset + len(f.symbols) + embedded.append(_EmbeddedFile(f, dense[offset:end], sparse[offset:end])) + offset = end + return embedded + + class IndexPipeline: def __init__( self, store: QdrantStore, registry: ServiceRegistry | None = None @@ -166,6 +251,122 @@ def __init__( self._sparse_embedder: BM25SparseProvider = get_sparse_embedding_provider() self._registry = registry or get_service_registry() + async def _produce_parsed_files( + self, + svc: ServiceConfig, + targets: list[tuple[GitHubFile, str]], + http_client: httpx.AsyncClient, + queue: asyncio.Queue[_ParsedFile | None], + on_file_done: Callable[[], None], + ) -> None: + """Fetches and parses *targets* concurrently, feeding results to *queue*.""" + sem = asyncio.Semaphore(_FETCH_CONCURRENCY) + + async def _fetch_and_parse(f: GitHubFile, stored_path: str) -> None: + async with sem: + try: + content = await fetch_blob_content( + settings.github_token, + svc.github_repo, + f.blob_sha, + client=http_client, + ) + except Exception as exc: # noqa: BLE001 — keep existing index entries on any fetch failure + logger.error("Failed to fetch %s: %s", stored_path, exc) + on_file_done() + return + + # parse_file stays on the event loop on purpose: registry.py shares one + # parser instance per language and tree-sitter parsers are not safe for + # concurrent use, so a thread hop here would be a data race. + try: + symbols = parse_file(content, stored_path) + except ParseError: + logger.error( + "Skipping index update for %s: parser failed, " + "existing entries preserved", + stored_path, + ) + on_file_done() + return + + # Queued outside the semaphore so a full queue holds no fetch slot. + await queue.put(_ParsedFile(stored_path, f.blob_sha, symbols)) + + try: + await asyncio.gather(*[_fetch_and_parse(f, p) for f, p in targets]) + finally: + # Sentinel, even on cancellation — the consumer must always terminate. + await queue.put(None) + + async def _embed_files( + self, files: list[_ParsedFile], service_name: str + ) -> list[_EmbeddedFile]: + """Embeds every symbol in *files* as one batch, overlapping dense and sparse.""" + if not files: + return [] + + symbols = [s for f in files for s in f.symbols] + dense_texts = [_build_embedding_text(s, service_name) for s in symbols] + sparse_texts = [_build_bm25_text(s) for s in symbols] + try: + dense, sparse = await asyncio.gather( + self._embedder.embed_batch(dense_texts), + self._sparse_embedder.embed_batch(sparse_texts), + ) + except Exception as exc: # noqa: BLE001 — one bad file must not drop the whole batch + logger.warning( + "Batch embedding failed for %d files (%s) — retrying file by file", + len(files), + exc, + ) + return await self._embed_files_individually(files, service_name) + + return _split_by_file(files, dense, sparse) + + async def _embed_files_individually( + self, files: list[_ParsedFile], service_name: str + ) -> list[_EmbeddedFile]: + """Per-file retry after a batch failure, so only the bad file is dropped.""" + embedded: list[_EmbeddedFile] = [] + for f in files: + try: + dense, sparse = await asyncio.gather( + self._embedder.embed_batch( + [_build_embedding_text(s, service_name) for s in f.symbols] + ), + self._sparse_embedder.embed_batch( + [_build_bm25_text(s) for s in f.symbols] + ), + ) + except Exception as exc: # noqa: BLE001 — keep existing index entries until embedding succeeds + logger.error("Embedding failed for %s: %s", f.stored_path, exc) + continue + embedded.append(_EmbeddedFile(f, dense, sparse)) + return embedded + + async def _write_embedded_file( + self, service_name: str, embedded: _EmbeddedFile + ) -> int: + """Upserts one file's symbols and prunes the ids it no longer covers.""" + parsed = embedded.parsed + payloads = [ + _symbol_to_payload(s, service_name, parsed.blob_sha) for s in parsed.symbols + ] + # Upsert new/changed symbols before deleting stale ones, so there's + # never a window where the file has zero indexed symbols. + previous_ids = await self._store.get_point_ids_by_file( + service_name, parsed.stored_path + ) + new_ids = await self._store.upsert_chunks( + payloads, embedded.dense, embedded.sparse + ) + stale_ids = previous_ids - set(new_ids) + await self._store.delete_by_ids(list(stale_ids)) + + logger.info("Indexed %s: %d symbols", parsed.stored_path, len(parsed.symbols)) + return len(parsed.symbols) + async def index_service( self, service_name: str, @@ -203,88 +404,85 @@ async def index_service( existing_hashes = await self._store.get_indexed_file_hashes(svc.name) - indexed_files = 0 - total_chunks = 0 - skipped = 0 total_files = len(github_files) - - for i, f in enumerate(github_files): - # "{service_name}/{path_in_repo}" — consistent path format across all tools - stored_path = f"{svc.name}/{f.rel_path}" - - # blob_sha IS the content fingerprint — no download needed to detect unchanged files - if not force and existing_hashes.get(stored_path) == f.blob_sha: - skipped += 1 - continue - - try: - content = await fetch_blob_content( - settings.github_token, - svc.github_repo, - f.blob_sha, - client=http_client, - ) - except Exception as exc: # noqa: BLE001 — keep existing index entries on any fetch failure - logger.error("Failed to fetch %s: %s", stored_path, exc) - continue - - try: - symbols = parse_file(content, stored_path) - except ParseError: - logger.error( - "Skipping index update for %s: parser failed, " - "existing entries preserved", - stored_path, - ) - continue - - if not symbols: - # File has no indexable symbols; clean up any stale entries. - await self._store.delete_by_file(svc.name, stored_path) - continue - - texts_dense = [_build_embedding_text(s, svc.name) for s in symbols] - texts_sparse = [_build_bm25_text(s) for s in symbols] - try: - dense_vectors = await self._embedder.embed_batch(texts_dense) - sparse_vectors = await self._sparse_embedder.embed_batch( - texts_sparse - ) - except Exception as exc: # noqa: BLE001 — keep existing index entries until embedding succeeds - logger.error("Embedding failed for %s: %s", stored_path, exc) - continue - - payloads = [ - _symbol_to_payload(s, svc.name, f.blob_sha) for s in symbols - ] - # Upsert new/changed symbols before deleting stale ones, so there's - # never a window where the file has zero indexed symbols. - previous_ids = await self._store.get_point_ids_by_file( - svc.name, stored_path - ) - new_ids = await self._store.upsert_chunks( - payloads, dense_vectors, sparse_vectors + # "{service_name}/{path_in_repo}" — consistent path format across all tools. + # blob_sha IS the content fingerprint — no download needed to detect + # unchanged files, so they never become fetch targets. + targets = [ + (f, path) + for f, path in ( + (f, f"{svc.name}/{f.rel_path}") for f in github_files ) - stale_ids = previous_ids - set(new_ids) - await self._store.delete_by_ids(list(stale_ids)) + if force or existing_hashes.get(path) != f.blob_sha + ] + skipped = total_files - len(targets) + + indexed_files = 0 + total_chunks = 0 + # Every file reaches a terminal state exactly once, so this only grows. + processed = skipped - indexed_files += 1 - total_chunks += len(symbols) - logger.info("Indexed %s: %d symbols", stored_path, len(symbols)) + def _mark_done() -> None: + nonlocal processed + processed += 1 + async def _emit_progress() -> None: if progress_callback: await progress_callback( ProgressEvent( phase="upserting", - current=i + 1, + current=processed, total=total_files, percentage=round( - (i + 1) / max(total_files, 1) * 100, 1 + processed / max(total_files, 1) * 100, 1 ), service=service_name, ) ) + batch_size = getattr( + self._embedder, "batch_size", _FALLBACK_EMBED_BATCH_SIZE + ) + queue: asyncio.Queue[_ParsedFile | None] = asyncio.Queue( + maxsize=_PARSED_QUEUE_SIZE + ) + producer = asyncio.create_task( + self._produce_parsed_files( + svc, targets, http_client, queue, _mark_done + ) + ) + try: + async for batch in _drain_batches(queue, batch_size): + to_embed = [] + for parsed in batch: + if parsed.symbols: + to_embed.append(parsed) + else: + # No indexable symbols; clean up any stale entries. + await self._store.delete_by_file( + svc.name, parsed.stored_path + ) + _mark_done() + + for embedded in await self._embed_files(to_embed, svc.name): + total_chunks += await self._write_embedded_file( + svc.name, embedded + ) + indexed_files += 1 + # Files dropped by an embedding failure are resolved too. + processed += len(to_embed) + + await _emit_progress() + + await producer + finally: + # Without this an exception in the consumer orphans the producer + # inside the http client's context manager. + producer.cancel() + await asyncio.gather(producer, return_exceptions=True) + + await _emit_progress() + all_stored_paths = {f"{svc.name}/{f.rel_path}" for f in github_files} stale_paths = [p for p in existing_hashes if p not in all_stored_paths] for stale_path in stale_paths: From 2eff86f23ed4f7cb064da3c6363f8d61ea922f1d Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sun, 13 Sep 2026 13:51:46 +0200 Subject: [PATCH 3/8] test: cover cross-file batching, fetch concurrency and embed fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each test was checked against a mutated pipeline to confirm it fails for the defect it targets: reusing one slice for every file breaks the vector split test, a semaphore of 1 breaks the concurrency test, and removing the per-file retry breaks the batch-containment test. _symbols_for names symbols uniquely per file on purpose — _build_embedding_text does not include the file path, so same-named symbols in different files produce identical embedding texts and make a text-to-vector lookup ambiguous. Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_pipeline.py | 363 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 363 insertions(+) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 1015d03..9210a09 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import logging from unittest.mock import AsyncMock, patch @@ -583,3 +584,365 @@ async def _list_github_files(*a, **k): ) assert max_active == 1 + + +class _RecordingEmbedder: + """Returns each text's position in its batch, so vectors are traceable to texts.""" + + dimensions = 1 + + def __init__(self, batch_size: int | None = None) -> None: + self.calls: list[list[str]] = [] + if batch_size is not None: + self.batch_size = batch_size + + async def embed_batch(self, texts: list[str]) -> list[list[float]]: + self.calls.append(list(texts)) + return [[float(i)] for i in range(len(texts))] + + +def _symbols_for(stored_path: str, count: int) -> list[CodeSymbol]: + # Names are unique across files: _build_embedding_text does not include the + # file path, so same-named symbols would produce identical embedding texts. + stem = stored_path.rsplit("/", 1)[-1].removesuffix(".py") + return [ + CodeSymbol( + name=f"{stem}_fn{i}", + symbol_type="function", + language="python", + source=f"def {stem}_fn{i}(): pass", + file_path=stored_path, + start_line=i + 1, + end_line=i + 1, + ) + for i in range(count) + ] + + +def _multi_file_patches(svc, files, symbols_by_path, fetch=None): + """The four module-level patches every index_service test needs.""" + return ( + patch.object( + type(pipeline_module.settings), "load_services", return_value=[svc] + ), + patch.object( + pipeline_module, "list_github_files", AsyncMock(return_value=files) + ), + patch.object( + pipeline_module, + "fetch_blob_content", + fetch or AsyncMock(return_value=b"source"), + ), + patch.object( + pipeline_module, + "parse_file", + lambda content, stored_path: symbols_by_path[stored_path], + ), + ) + + +async def test_symbols_from_multiple_files_share_one_embedding_call() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(3)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 2) for i in range(3) + } + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + embedder = _RecordingEmbedder() + pipeline._embedder = embedder + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path) + with p1, p2, p3, p4: + result = await pipeline.index_service("svc", force=True) + + # Six symbols across three files, embedded as a single batch... + assert len(embedder.calls) == 1 + assert len(embedder.calls[0]) == 6 + # ...but still written one file at a time. + assert store.upsert_chunks.call_count == 3 + assert all(len(call.args[0]) == 2 for call in store.upsert_chunks.call_args_list) + assert result == {"files": 3, "chunks": 6, "skipped": 0} + + +async def test_vectors_are_split_back_to_the_correct_file() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(3)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 2) for i in range(3) + } + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + embedder = _RecordingEmbedder() + pipeline._embedder = embedder + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path) + with p1, p2, p3, p4: + await pipeline.index_service("svc", force=True) + + batch_texts = embedder.calls[0] + # Every upserted symbol must carry the vector for its own text, whatever + # order the files completed in. + for call in store.upsert_chunks.call_args_list: + payloads, dense_vectors, _sparse = call.args + for payload, vector in zip(payloads, dense_vectors): + symbol = next( + s + for s in symbols_by_path[payload["file_path"]] + if s.name == payload["symbol_name"] + ) + expected = batch_texts.index(_build_embedding_text(symbol, "svc")) + assert vector == [float(expected)] + + +async def test_batch_is_cut_at_provider_batch_size() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(3)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 1) for i in range(3) + } + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + embedder = _RecordingEmbedder(batch_size=2) + pipeline._embedder = embedder + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path) + with p1, p2, p3, p4: + result = await pipeline.index_service("svc", force=True) + + assert [len(texts) for texts in embedder.calls] == [2, 1] + assert result["files"] == 3 + + +async def test_one_bad_file_does_not_drop_the_rest_of_the_batch() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(3)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 1) for i in range(3) + } + # Only m1's symbol is poisoned. + symbols_by_path["svc/m1.py"][0].source = "POISON" + + class _PoisonEmbedder(_RecordingEmbedder): + async def embed_batch(self, texts: list[str]) -> list[list[float]]: + if any("POISON" in t for t in texts): + raise RuntimeError("provider rejected input") + return await super().embed_batch(texts) + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + pipeline._embedder = _PoisonEmbedder() + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path) + with p1, p2, p3, p4: + result = await pipeline.index_service("svc", force=True) + + # The batch failed, the per-file retry salvaged the two healthy files. + written = { + call.args[0][0]["file_path"] for call in store.upsert_chunks.call_args_list + } + assert written == {"svc/m0.py", "svc/m2.py"} + assert result == {"files": 2, "chunks": 2, "skipped": 0} + # The poisoned file keeps whatever was already indexed for it. + store.delete_by_file.assert_not_called() + + +async def test_blob_fetches_run_concurrently() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(8)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 1) for i in range(8) + } + + active = 0 + max_active = 0 + + async def _fetch(*a, **k): + nonlocal active, max_active + active += 1 + max_active = max(max_active, active) + await asyncio.sleep(0.01) + active -= 1 + return b"source" + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path, fetch=_fetch) + with p1, p2, p3, p4: + result = await pipeline.index_service("svc", force=True) + + assert 1 < max_active <= pipeline_module._FETCH_CONCURRENCY + assert result["files"] == 8 + + +async def test_fetch_failure_for_one_file_does_not_block_others() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(2)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 1) for i in range(2) + } + + async def _fetch(token, repo, blob_sha, client=None): + if blob_sha == "sha0": + raise RuntimeError("boom") + return b"source" + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {"svc/m0.py": "old"} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path, fetch=_fetch) + with p1, p2, p3, p4: + result = await pipeline.index_service("svc", force=True) + + assert result == {"files": 1, "chunks": 1, "skipped": 0} + store.upsert_chunks.assert_called_once() + assert store.upsert_chunks.call_args.args[0][0]["file_path"] == "svc/m1.py" + # The unfetchable file keeps its existing entries. + store.delete_by_file.assert_not_called() + + +async def test_empty_symbol_file_in_a_mixed_batch_still_deletes_stale() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [ + GitHubFile(rel_path="empty.py", blob_sha="sha0"), + GitHubFile(rel_path="full.py", blob_sha="sha1"), + ] + symbols_by_path = { + "svc/empty.py": [], + "svc/full.py": _symbols_for("svc/full.py", 1), + } + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path) + with p1, p2, p3, p4: + result = await pipeline.index_service("svc", force=True) + + store.delete_by_file.assert_called_once_with("svc", "svc/empty.py") + store.upsert_chunks.assert_called_once() + assert result == {"files": 1, "chunks": 1, "skipped": 0} + + +async def test_progress_current_is_monotonic_and_reaches_total() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(8)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 1) for i in range(8) + } + symbols_by_path["svc/m5.py"] = [] # no indexable symbols + + def _parse(content, stored_path): + if stored_path == "svc/m6.py": + raise ParseError(stored_path) + return symbols_by_path[stored_path] + + store = AsyncMock() + # Two files are unchanged and get skipped without a fetch. + store.get_indexed_file_hashes.return_value = { + "svc/m0.py": "sha0", + "svc/m1.py": "sha1", + } + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + pipeline._embedder = _RecordingEmbedder(batch_size=2) + + events = [] + + async def _callback(event) -> None: + events.append(event) + + with ( + patch.object( + type(pipeline_module.settings), "load_services", return_value=[svc] + ), + patch.object( + pipeline_module, "list_github_files", AsyncMock(return_value=files) + ), + patch.object( + pipeline_module, "fetch_blob_content", AsyncMock(return_value=b"source") + ), + patch.object(pipeline_module, "parse_file", _parse), + ): + result = await pipeline.index_service( + "svc", force=False, progress_callback=_callback + ) + + upserting = [e.current for e in events if e.phase == "upserting"] + assert upserting == sorted(upserting), "progress must never go backwards" + assert all(c <= 8 for c in upserting) + # Every file is accounted for: 4 indexed + 1 empty + 1 parse failure + 2 skipped. + assert upserting[-1] == 8 + assert result == {"files": 4, "chunks": 4, "skipped": 2} + + +async def test_dense_and_sparse_embeds_overlap() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path="m.py", blob_sha="sha")] + symbols_by_path = {"svc/m.py": _symbols_for("svc/m.py", 1)} + + dense_entered = asyncio.Event() + sparse_entered = asyncio.Event() + + class _GatedDense(_StubEmbedder): + async def embed_batch(self, texts): + dense_entered.set() + await asyncio.wait_for(sparse_entered.wait(), timeout=1) + return [[0.0]] * len(texts) + + class _GatedSparse(_StubEmbedder): + async def embed_batch(self, texts): + sparse_entered.set() + await asyncio.wait_for(dense_entered.wait(), timeout=1) + return [[0.0]] * len(texts) + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.return_value = [] + + pipeline = _make_pipeline(store) + pipeline._embedder = _GatedDense() + pipeline._sparse_embedder = _GatedSparse() + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path) + with p1, p2, p3, p4: + # Each side waits for the other to start, so this only completes if the + # two embeds are in flight together. + result = await pipeline.index_service("svc", force=True) + + assert result["files"] == 1 From 95c8da556b3e2b4da6e1b395260ec979114b6bac Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sun, 13 Sep 2026 13:52:56 +0200 Subject: [PATCH 4/8] docs: describe concurrent fetching and cross-file embedding batches Also corrects two entries that were already stale before this change: the upsert section still described delete-before-upsert, which the pipeline stopped doing in 2cdf49c, and the matching "delete-before-upsert gap" observation described a window that no longer exists. Replaced the latter with the rate-limit exposure that concurrent fetching does introduce. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 4 ++++ docs/docs/ingestion.md | 33 +++++++++++++++++++++------------ server/routes/reindex.py | 2 ++ 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 2548ed2..469d940 100644 --- a/README.md +++ b/README.md @@ -345,6 +345,10 @@ Frame shapes: {"type": "error", "message": "..."} ``` +For `/reindex` the `upserting` frames carry `current` as a monotonically increasing count of files +resolved — indexed, skipped as unchanged, or dropped by a fetch/parse failure — so it always ends at +`total`. Files are indexed concurrently, so frames are emitted per batch rather than per file. + For `/reindex-history` the `phase` value is `discovery|embedding|upserting` and the `done` result is `{"new": int, "skipped": int, "diff_updated": int}`. diff --git a/docs/docs/ingestion.md b/docs/docs/ingestion.md index 921835d..f341be3 100644 --- a/docs/docs/ingestion.md +++ b/docs/docs/ingestion.md @@ -16,7 +16,7 @@ Ingestion is managed by `IndexPipeline` (`server/indexer/pipeline.py`). For each 1. Discovers all indexable files from GitHub 2. Skips files whose content hasn't changed since the last index run -3. Downloads changed file content, parses it into `CodeSymbol` entries, generates dense and sparse embeddings, and upserts them into Qdrant +3. Downloads changed file content concurrently, parses it into `CodeSymbol` entries, generates dense and sparse embeddings in batches that span multiple files, and upserts them into Qdrant 4. Removes index entries for files that have been deleted from the repository The pipeline is triggered via the `/reindex` HTTP endpoint (streaming NDJSON progress) or the `index_all` MCP admin tool. @@ -76,10 +76,14 @@ content = await fetch_blob_content( Fetching by blob SHA is more efficient than path-based fetching during indexing: the SHA is already known from the tree response, and the blob API is a direct content lookup with no ref resolution overhead. +Fetches run concurrently, bounded by a semaphore (`_FETCH_CONCURRENCY`, default 10 — the same bound `github_source.py` uses for tree walks and commit diffs). Each fetched file is parsed and handed to a bounded queue (`_PARSED_QUEUE_SIZE`), so downloading continues while the previous group of symbols is being embedded. A fetch failure logs an error and drops only that file; its existing index entries are preserved. + ### 4. Parsing `parse_file(content, stored_path)` dispatches to the language-specific parser via the registry. The result is a `list[CodeSymbol]` — one entry per indexable symbol (class, method, function, interface, etc.). +Parsing runs on the event loop rather than in a thread pool, deliberately: `registry.py` builds one parser instance per language and shares it across every file, and tree-sitter `Parser` objects are not safe for concurrent use. Moving `parse_file` to a worker thread would be a data race. + If a file produces no symbols (empty, unsupported format, or parse failure), any existing index entries for that file are cleaned up and the file is skipped. ### 5. Embedding Text Construction @@ -121,25 +125,32 @@ This text is then pre-processed by `split_code_identifiers` (see [sparse-vectors ### 6. Embedding -Both embedding calls are made sequentially per file batch: +Symbols are accumulated **across files** until they fill the provider's batch size (`EmbeddingProvider.batch_size` — 128 for Voyage/OpenAI/Jina's hosted API, 32 for self-hosted Jina TEI and Ollama), or until the batch reaches `_MAX_BATCH_CHARS`. A single file usually yields only a handful of symbols, so batching across files is what keeps requests full instead of sending one tiny request per file. + +The dense and sparse embeds for a batch run concurrently: ```python -dense_vectors = await self._embedder.embed_batch(texts_dense) -sparse_vectors = await self._sparse_embedder.embed_batch(texts_sparse) +dense, sparse = await asyncio.gather( + self._embedder.embed_batch(dense_texts), + self._sparse_embedder.embed_batch(sparse_texts), +) ``` -If either call raises an exception, the file is skipped and existing index entries are preserved until the next successful run. +Dense is network-bound and sparse runs in a thread executor, so the two overlap for free. + +If a batch call raises, the pipeline retries that batch **file by file**, so a single unembeddable file costs one extra round-trip rather than dropping every file batched alongside it. A file that still fails is skipped, and its existing index entries are preserved until the next successful run. ### 7. Upsert -Before inserting new vectors, all existing entries for the file are removed: +Writes stay scoped to one file at a time, even though embedding is batched. New vectors are upserted *before* stale ones are deleted, so the file never has zero indexed symbols: ```python -await self._store.delete_by_file(svc.name, stored_path) -await self._store.upsert_chunks(payloads, dense_vectors, sparse_vectors) +previous_ids = await self._store.get_point_ids_by_file(service_name, stored_path) +new_ids = await self._store.upsert_chunks(payloads, dense, sparse) +await self._store.delete_by_ids(list(previous_ids - set(new_ids))) ``` -This ensures clean replacement when symbols are added, removed, or renamed within a file. Each point's ID is a deterministic `uuid5` derived from `service:file_path:symbol_name:start_line`, so symbols moving to a new line produce new IDs (handled correctly by the delete-first approach). +Each point's ID is a deterministic `uuid5` derived from `service:file_path:symbol_name:start_line`, so symbols moving to a new line produce new IDs and the old ones fall out as stale. Because the ID includes `file_path`, two different files can never produce colliding IDs — which is what makes batching across files safe. Each point carries a payload with 20+ fields (see Data Model below). @@ -197,13 +208,11 @@ All `CodeSymbol` fields are stored verbatim, plus: ## Observations -**Sequential embedding calls** — `embed_batch` for dense and `embed_batch` for sparse are awaited sequentially. They are independent operations targeting different providers; wrapping them in `asyncio.gather` would reduce per-file embedding latency by ~50%. - **No embedding retry** — a transient API error on either embedding call causes the file to be silently skipped, leaving its existing index stale indefinitely. There is no exponential backoff or retry queue. Reindexing requires either a force reindex or waiting for the file's content to change. **BM25 text still omits some dense-only metadata** — `_build_bm25_text` folds in name, package, annotations, and HTTP method/route, but the dense preamble's service name, language, and symbol-type phrasing (e.g. "Java method") are still dense-only. A BM25 query for "Python method" will not match unless the word "Python" or "method" appears elsewhere in the folded-in fields or the source code itself. -**delete-before-upsert gap** — The pipeline deletes all entries for a file before upserting the new ones. If the process is interrupted between delete and upsert, the file has no index entries. The next incremental run will redownload and reindex the file correctly — but until then, queries miss the file entirely. +**GitHub rate limits under concurrent fetching** — `fetch_blob_content` has no 429/backoff handling of its own (unlike the embedding path, which retries via `post_with_retry`). With `_FETCH_CONCURRENCY` downloads in flight, a secondary rate limit surfaces as per-file fetch errors, leaving those files un-reindexed until the next run. **GitHub Trees truncation** — Very large repositories may have their tree response silently truncated by the GitHub API. The pipeline logs a warning but does not retry or paginate to recover the missing entries. diff --git a/server/routes/reindex.py b/server/routes/reindex.py index ad04a69..ffb5a0a 100644 --- a/server/routes/reindex.py +++ b/server/routes/reindex.py @@ -26,6 +26,8 @@ async def reindex(request: Request) -> StreamingResponse: Emits progress frames while indexing, followed by a final summary frame: {"type": "progress", "phase": "discovery"|"upserting"|"cleanup", "current": int, "total": int, "percentage": float, "service": str} + `current` counts files resolved (indexed, skipped, or failed) and only ever + increases; indexing is concurrent, so frames arrive per batch, not per file. {"type": "done", "result": {"files": int, "chunks": int, "skipped": int}} Body (optional JSON): From 1aa873aa1ebf45c6140c0c768d6d32c01007a94c Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sun, 13 Sep 2026 13:53:42 +0200 Subject: [PATCH 5/8] perf: skip duplicate progress frames at an unchanged file count The final emit after the drain loop repeated the last batch's frame verbatim. Progress frames are now strictly increasing. Co-Authored-By: Claude Opus 5 (1M context) --- server/indexer/pipeline.py | 6 +++++- tests/test_pipeline.py | 3 ++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/server/indexer/pipeline.py b/server/indexer/pipeline.py index 86f3a0e..92f7e91 100644 --- a/server/indexer/pipeline.py +++ b/server/indexer/pipeline.py @@ -426,8 +426,12 @@ def _mark_done() -> None: nonlocal processed processed += 1 + last_emitted = -1 + async def _emit_progress() -> None: - if progress_callback: + nonlocal last_emitted + if progress_callback and processed != last_emitted: + last_emitted = processed await progress_callback( ProgressEvent( phase="upserting", diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 9210a09..69f7dac 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -903,7 +903,8 @@ async def _callback(event) -> None: ) upserting = [e.current for e in events if e.phase == "upserting"] - assert upserting == sorted(upserting), "progress must never go backwards" + # Strictly increasing: never backwards, and no repeated frame at the same count. + assert upserting == sorted(set(upserting)) assert all(c <= 8 for c in upserting) # Every file is accounted for: 4 indexed + 1 empty + 1 parse failure + 2 skipped. assert upserting[-1] == 8 From eb1c394af56a43cbaa4707f80a99e6f291f9ca12 Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sun, 13 Sep 2026 13:54:19 +0200 Subject: [PATCH 6/8] perf: skip the per-file retry when a failed batch holds one file There is nothing to isolate, so the retry only repeats the call that just failed. Keeps the single-file failure path to one attempt, as before. Co-Authored-By: Claude Opus 5 (1M context) --- server/indexer/pipeline.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/server/indexer/pipeline.py b/server/indexer/pipeline.py index 92f7e91..0901a2d 100644 --- a/server/indexer/pipeline.py +++ b/server/indexer/pipeline.py @@ -315,6 +315,10 @@ async def _embed_files( self._sparse_embedder.embed_batch(sparse_texts), ) except Exception as exc: # noqa: BLE001 — one bad file must not drop the whole batch + if len(files) == 1: + # Nothing to isolate; retrying would just repeat the same call. + logger.error("Embedding failed for %s: %s", files[0].stored_path, exc) + return [] logger.warning( "Batch embedding failed for %d files (%s) — retrying file by file", len(files), From a17621760c7f814913499d17a5e6d32086fb474a Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sun, 13 Sep 2026 13:58:58 +0200 Subject: [PATCH 7/8] fix: retry GitHub 5xx and transport errors, not just rate limits _gh_get already retried 403/429 but let transport errors, timeouts and 5xx fall straight through. Concurrent blob fetching makes exactly those more likely, and the embedding path has handled all three since post_with_retry landed. Rate limits keep waiting for the window GitHub names; 5xx and transport errors use a short fixed backoff, since they carry no such hint and usually clear immediately. Also corrects the ingestion doc note added in 95c8da5, which claimed fetch_blob_content had no rate-limit handling at all. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs/ingestion.md | 4 +- server/indexer/github_source.py | 86 +++++++++++++++++---- tests/test_github_source_retry.py | 121 ++++++++++++++++++++++++++++++ 3 files changed, 193 insertions(+), 18 deletions(-) create mode 100644 tests/test_github_source_retry.py diff --git a/docs/docs/ingestion.md b/docs/docs/ingestion.md index f341be3..49f2460 100644 --- a/docs/docs/ingestion.md +++ b/docs/docs/ingestion.md @@ -76,7 +76,7 @@ content = await fetch_blob_content( Fetching by blob SHA is more efficient than path-based fetching during indexing: the SHA is already known from the tree response, and the blob API is a direct content lookup with no ref resolution overhead. -Fetches run concurrently, bounded by a semaphore (`_FETCH_CONCURRENCY`, default 10 — the same bound `github_source.py` uses for tree walks and commit diffs). Each fetched file is parsed and handed to a bounded queue (`_PARSED_QUEUE_SIZE`), so downloading continues while the previous group of symbols is being embedded. A fetch failure logs an error and drops only that file; its existing index entries are preserved. +Fetches run concurrently, bounded by a semaphore (`_FETCH_CONCURRENCY`, default 10 — the same bound `github_source.py` uses for tree walks and commit diffs). Each fetched file is parsed and handed to a bounded queue (`_PARSED_QUEUE_SIZE`), so downloading continues while the previous group of symbols is being embedded. Every request goes through `_gh_get`, which retries rate limits, 5xx, and transport errors. A fetch that still fails logs an error and drops only that file; its existing index entries are preserved. ### 4. Parsing @@ -212,7 +212,7 @@ All `CodeSymbol` fields are stored verbatim, plus: **BM25 text still omits some dense-only metadata** — `_build_bm25_text` folds in name, package, annotations, and HTTP method/route, but the dense preamble's service name, language, and symbol-type phrasing (e.g. "Java method") are still dense-only. A BM25 query for "Python method" will not match unless the word "Python" or "method" appears elsewhere in the folded-in fields or the source code itself. -**GitHub rate limits under concurrent fetching** — `fetch_blob_content` has no 429/backoff handling of its own (unlike the embedding path, which retries via `post_with_retry`). With `_FETCH_CONCURRENCY` downloads in flight, a secondary rate limit surfaces as per-file fetch errors, leaving those files un-reindexed until the next run. +**GitHub retries are bounded** — `_gh_get` retries rate limits (403/429, waiting for the window named by `Retry-After` / `X-RateLimit-Reset`, capped at 120s), 5xx, and transport errors, for `_GH_ATTEMPTS` attempts total. A failure that outlives those attempts surfaces as a per-file fetch error, leaving that file un-reindexed until the next run rather than failing the whole service. **GitHub Trees truncation** — Very large repositories may have their tree response silently truncated by the GitHub API. The pipeline logs a warning but does not retry or paginate to recover the missing entries. diff --git a/server/indexer/github_source.py b/server/indexer/github_source.py index 1cceb82..58302b2 100644 --- a/server/indexer/github_source.py +++ b/server/indexer/github_source.py @@ -19,6 +19,15 @@ _DIFF_CONCURRENCY = 10 _TREE_WALK_CONCURRENCY = 10 +# Attempts per GitHub GET, shared by rate limits, 5xx and transport errors. +_GH_ATTEMPTS = 3 +# Backoff for 5xx and transport errors. Rate limits ignore this and wait for the +# window named by Retry-After / X-RateLimit-Reset instead. +_GH_BACKOFF_DELAYS = (1.0, 5.0) +# Transient network failures worth another attempt. httpx.TimeoutException is a +# subclass of TransportError, but naming it keeps the intent obvious. +_RETRYABLE_EXCEPTIONS = (httpx.TransportError, httpx.TimeoutException) + logger = logging.getLogger(__name__) @@ -83,23 +92,68 @@ async def _gh_get( params: dict | None = None, timeout: float = 30.0, ) -> Any: - """GET a GitHub API URL, retrying up to 3 times on rate-limit responses (403/429).""" + """GET a GitHub API URL, retrying rate limits (403/429), 5xx and transport errors. + + Rate limits wait for the window GitHub names via ``Retry-After`` / + ``X-RateLimit-Reset`` (capped at 120s); 5xx and transport errors use a short + fixed backoff, since they carry no such hint and usually clear immediately. + """ headers = _auth_headers(token) - for _ in range(3): - r = await client.get(url, headers=headers, params=params, timeout=timeout) - if r.status_code not in (403, 429): - r.raise_for_status() - return r.json() - reset_ts = float(r.headers.get("X-RateLimit-Reset", 0)) - retry_after = float(r.headers.get("Retry-After", 60)) - now = time.time() - wait = min(max(retry_after, reset_ts - now if reset_ts > now else 0.0), 120.0) - logger.warning( - "GitHub rate-limited (HTTP %d) — retrying in %.0fs", r.status_code, wait - ) - await asyncio.sleep(wait) - r.raise_for_status() # raise final rate-limit error after exhausting retries - return r.json() # unreachable + r: httpx.Response | None = None + + for attempt in range(_GH_ATTEMPTS): + last_attempt = attempt == _GH_ATTEMPTS - 1 + try: + r = await client.get(url, headers=headers, params=params, timeout=timeout) + except _RETRYABLE_EXCEPTIONS as exc: + if last_attempt: + raise + wait = _GH_BACKOFF_DELAYS[min(attempt, len(_GH_BACKOFF_DELAYS) - 1)] + logger.warning( + "GitHub request failed (%s: %s) — retrying in %.0fs (attempt %d/%d)", + type(exc).__name__, + exc, + wait, + attempt + 1, + _GH_ATTEMPTS, + ) + await asyncio.sleep(wait) + continue + + if r.status_code in (403, 429): + if last_attempt: + break + reset_ts = float(r.headers.get("X-RateLimit-Reset", 0)) + retry_after = float(r.headers.get("Retry-After", 60)) + now = time.time() + wait = min( + max(retry_after, reset_ts - now if reset_ts > now else 0.0), 120.0 + ) + logger.warning( + "GitHub rate-limited (HTTP %d) — retrying in %.0fs", r.status_code, wait + ) + await asyncio.sleep(wait) + continue + + if r.status_code >= 500: + if last_attempt: + break + wait = _GH_BACKOFF_DELAYS[min(attempt, len(_GH_BACKOFF_DELAYS) - 1)] + logger.warning( + "GitHub server error (%d) — retrying in %.0fs (attempt %d/%d)", + r.status_code, + wait, + attempt + 1, + _GH_ATTEMPTS, + ) + await asyncio.sleep(wait) + continue + + break + + assert r is not None # a transport error on the last attempt re-raises + r.raise_for_status() # surfaces the final rate-limit / 5xx error + return r.json() def _filter_tree_blobs( diff --git a/tests/test_github_source_retry.py b/tests/test_github_source_retry.py new file mode 100644 index 0000000..a0c93ab --- /dev/null +++ b/tests/test_github_source_retry.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import httpx +import pytest +import respx + +import server.indexer.github_source as gh +from server.indexer.github_source import fetch_blob_content + +_API = "https://api.github.com" +_REPO = "owner/repo" +_BLOB_URL = f"{_API}/repos/{_REPO}/git/blobs/deadbeef" + + +@pytest.fixture +def no_sleep(monkeypatch): + """Record backoff waits instead of serving them.""" + waits: list[float] = [] + + async def _sleep(seconds: float) -> None: + waits.append(seconds) + + monkeypatch.setattr(gh.asyncio, "sleep", _sleep) + return waits + + +def _blob(content: str) -> httpx.Response: + import base64 + + encoded = base64.b64encode(content.encode()).decode() + return httpx.Response(200, json={"content": encoded}) + + +@respx.mock +async def test_transport_error_is_retried(no_sleep) -> None: + route = respx.get(_BLOB_URL).mock( + side_effect=[ + httpx.ConnectError("connection reset"), + _blob("recovered"), + ] + ) + + content = await fetch_blob_content("tok", _REPO, "deadbeef") + + assert content == b"recovered" + assert route.call_count == 2 + assert no_sleep == [gh._GH_BACKOFF_DELAYS[0]] + + +@respx.mock +async def test_server_error_is_retried(no_sleep) -> None: + route = respx.get(_BLOB_URL).mock( + side_effect=[httpx.Response(502), _blob("recovered")] + ) + + content = await fetch_blob_content("tok", _REPO, "deadbeef") + + assert content == b"recovered" + assert route.call_count == 2 + + +@respx.mock +async def test_transport_error_raises_after_exhausting_attempts(no_sleep) -> None: + route = respx.get(_BLOB_URL).mock(side_effect=httpx.ConnectError("down")) + + with pytest.raises(httpx.ConnectError): + await fetch_blob_content("tok", _REPO, "deadbeef") + + assert route.call_count == gh._GH_ATTEMPTS + + +@respx.mock +async def test_server_error_raises_after_exhausting_attempts(no_sleep) -> None: + route = respx.get(_BLOB_URL).mock(return_value=httpx.Response(503)) + + with pytest.raises(httpx.HTTPStatusError): + await fetch_blob_content("tok", _REPO, "deadbeef") + + assert route.call_count == gh._GH_ATTEMPTS + + +@respx.mock +async def test_rate_limit_waits_for_the_reset_window(no_sleep) -> None: + route = respx.get(_BLOB_URL).mock( + side_effect=[ + httpx.Response(429, headers={"Retry-After": "30"}), + _blob("recovered"), + ] + ) + + content = await fetch_blob_content("tok", _REPO, "deadbeef") + + assert content == b"recovered" + assert route.call_count == 2 + # Rate limits honour the server's window, not the short 5xx backoff. + assert no_sleep == [30.0] + + +@respx.mock +async def test_rate_limit_wait_is_capped(no_sleep) -> None: + respx.get(_BLOB_URL).mock( + side_effect=[ + httpx.Response(403, headers={"Retry-After": "9999"}), + _blob("recovered"), + ] + ) + + await fetch_blob_content("tok", _REPO, "deadbeef") + + assert no_sleep == [120.0] + + +@respx.mock +async def test_client_error_is_not_retried(no_sleep) -> None: + route = respx.get(_BLOB_URL).mock(return_value=httpx.Response(404)) + + with pytest.raises(httpx.HTTPStatusError): + await fetch_blob_content("tok", _REPO, "deadbeef") + + assert route.call_count == 1 + assert no_sleep == [] From 34faefbe5b9e7ef61da487e32fac8459312e1838 Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sun, 13 Sep 2026 14:00:02 +0200 Subject: [PATCH 8/8] perf: overlap Qdrant writes across files within a batch Each file's own upsert-then-delete stays ordered, so the index never has a window with zero symbols for a file. Only different files overlap, and _symbol_point_id includes the file path, so two files can never contend for the same point. Co-Authored-By: Claude Opus 5 (1M context) --- server/indexer/pipeline.py | 30 +++++++++++++--- tests/test_pipeline.py | 72 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 5 deletions(-) diff --git a/server/indexer/pipeline.py b/server/indexer/pipeline.py index 0901a2d..cf04221 100644 --- a/server/indexer/pipeline.py +++ b/server/indexer/pipeline.py @@ -37,6 +37,9 @@ # Parsed files waiting to be embedded. Bounds peak memory and applies # backpressure to the fetchers while a batch is embedding. _PARSED_QUEUE_SIZE = 20 +# Concurrent Qdrant write round-trips. Each file's own upsert-then-delete stays +# ordered; only different files overlap, and their point ids never collide. +_WRITE_CONCURRENCY = 4 # Used when the provider does not declare a batch_size (duck-typed test stubs). _FALLBACK_EMBED_BATCH_SIZE = 32 # Secondary cut so a batch of unusually large symbols never becomes a giant request. @@ -371,6 +374,23 @@ async def _write_embedded_file( logger.info("Indexed %s: %d symbols", parsed.stored_path, len(parsed.symbols)) return len(parsed.symbols) + async def _write_batch( + self, service_name: str, embedded_files: list[_EmbeddedFile] + ) -> int: + """Writes a batch's files concurrently, returning the symbols written. + + Safe to overlap because each file's own upsert-then-delete stays ordered + and `_symbol_point_id` includes the file path, so two files can never + touch the same point. + """ + sem = asyncio.Semaphore(_WRITE_CONCURRENCY) + + async def _write_one(embedded: _EmbeddedFile) -> int: + async with sem: + return await self._write_embedded_file(service_name, embedded) + + return sum(await asyncio.gather(*[_write_one(e) for e in embedded_files])) + async def index_service( self, service_name: str, @@ -472,11 +492,11 @@ async def _emit_progress() -> None: ) _mark_done() - for embedded in await self._embed_files(to_embed, svc.name): - total_chunks += await self._write_embedded_file( - svc.name, embedded - ) - indexed_files += 1 + embedded_files = await self._embed_files(to_embed, svc.name) + total_chunks += await self._write_batch( + svc.name, embedded_files + ) + indexed_files += len(embedded_files) # Files dropped by an embedding failure are resolved too. processed += len(to_embed) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 69f7dac..2a31d8d 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -947,3 +947,75 @@ async def embed_batch(self, texts): result = await pipeline.index_service("svc", force=True) assert result["files"] == 1 + + +async def test_writes_for_different_files_overlap() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(4)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 1) for i in range(4) + } + + active = 0 + max_active = 0 + + async def _upsert(payloads, dense, sparse): + nonlocal active, max_active + active += 1 + max_active = max(max_active, active) + await asyncio.sleep(0.01) + active -= 1 + return [] + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.return_value = set() + store.upsert_chunks.side_effect = _upsert + + pipeline = _make_pipeline(store) + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path) + with p1, p2, p3, p4: + result = await pipeline.index_service("svc", force=True) + + assert 1 < max_active <= pipeline_module._WRITE_CONCURRENCY + assert result == {"files": 4, "chunks": 4, "skipped": 0} + + +async def test_concurrent_writes_keep_upsert_before_delete_per_file() -> None: + svc = ServiceConfig(name="svc", github_repo="org/repo", exclude=[]) + files = [GitHubFile(rel_path=f"m{i}.py", blob_sha=f"sha{i}") for i in range(4)] + symbols_by_path = { + f"svc/m{i}.py": _symbols_for(f"svc/m{i}.py", 1) for i in range(4) + } + + events: list[tuple[str, str]] = [] + + async def _upsert(payloads, dense, sparse): + path = payloads[0]["file_path"] + events.append(("upsert", path)) + await asyncio.sleep(0.01) + return [f"new-{path}"] + + async def _delete_by_ids(ids): + # Stale ids are namespaced by file, so the path is recoverable. + for i in ids: + events.append(("delete", i.removeprefix("old-"))) + + store = AsyncMock() + store.get_indexed_file_hashes.return_value = {} + store.get_point_ids_by_file.side_effect = lambda service, path: {f"old-{path}"} + store.upsert_chunks.side_effect = _upsert + store.delete_by_ids.side_effect = _delete_by_ids + + pipeline = _make_pipeline(store) + + p1, p2, p3, p4 = _multi_file_patches(svc, files, symbols_by_path) + with p1, p2, p3, p4: + await pipeline.index_service("svc", force=True) + + # Files may interleave, but within one file the upsert always precedes the + # delete — the index never has a window with zero symbols for that file. + for path in symbols_by_path: + ordered = [action for action, p in events if p == path] + assert ordered == ["upsert", "delete"], path