# Concurrent Ingestion Pipeline — Implementation Plan

## Goal

Speed up ingestion by introducing concurrency to the I/O-bound stages of the
pipeline (embedding generation, graph writes, file reads). The primary
bottleneck is sequential embedding API calls (~100ms each, ~100s total for a
1000-embedding run). Secondary bottleneck is sequential Neo4j writes.

## Design Principles

1. **Async-first for I/O.** Use `asyncio` with `async for` / `async generators`
   for the embedding and graph-write hot paths. Threading is only used where
   async cannot reach (tree-sitter C calls, synchronous third-party libraries).
2. **Explicit pool lifecycle.** Every pool/semaphore is created in a known scope
   and torn down deterministically — no background threads outliving the
   pipeline.
3. **Conservative concurrency limits.** Semaphores cap concurrent embedding
   calls (default 8) and graph writes (default 16). Values are configurable,
   not hard-coded.
4. **Stages stay sequential.** The 5-stage pipeline order is preserved. Concurrency
   is *within* stages, never across them.
5. **Backward-compatible entry point.** `ingest_repo()` remains a sync function.
   It calls `asyncio.run()` once at the top level.

---

## Implementation Stages

This is a three-stage rollout. Each stage is a standalone PR that can be
reviewed, tested, and merged independently. If a stage introduces regressions
it can be reverted without affecting prior work.

---

### Stage 1 — Async Embedding Provider + Batched Vector Store

**Files changed:**
- `packages/clew-ingestion/src/embeddings.py`
- `packages/clew-ingestion/src/vector.py`
- `packages/clew-ingestion/tests/test_vector.py`
- `packages/clew-ingestion/tests/test_vector_batch.py`

**What changes:**

1. **Fix `OpenAIEmbeddings.embed()` to be truly async.**
   - Replace the sync `openai.OpenAI` client with `openai.AsyncOpenAI`.
   - The method signature is already `async def embed()` — the implementation
     just needs to match.

2. **Add `embed_batch()` to both embedding providers.**
   ```python
   async def embed_batch(self, texts: list[str]) -> list[list[float]]:
   ```
   - OpenAI: use a single `client.embeddings.create(input=texts)` call
     (the API already supports batch input, up to 2048 items).
   - Ollama: use `asyncio.gather()` with a `Semaphore(8)` to fan out
     individual requests.
   - This is the highest-impact change: replaces N serial API calls with
     1 batch call (OpenAI) or N concurrent calls capped by semaphore (Ollama).

3. **Add `VectorStore.add_batch()` — async batch method.**
   ```python
   async def add_batch(
       self,
       items: list[tuple[str, dict]],  # (text, metadata) pairs
   ) -> list[str]:
   ```
   - Calls `embed_batch()` once for all texts.
   - Extends internal `_buffer` with all resulting points.
   - Flushes if buffer exceeds `_batch_size`.
   - Returns list of generated point IDs.

4. **Keep `VectorStore.add()` working as-is** for backward compat with
   `sync.py` and existing tests. Internally rewire it to call `add_batch`
   with a single item using `asyncio.run()` (same as today, no behavior
   change).

5. **Update `EmbeddingProvider` protocol** in `vector.py` to add
   `embed_batch` with a default that falls back to sequential `embed()`:
   ```python
   async def embed_batch(self, texts: list[str]) -> list[list[float]]:
       return [await self.embed(t) for t in texts]
   ```

**Tests:**
- Unit test `embed_batch` with mocked HTTP responses.
- Unit test `add_batch` with a mock embedding provider.
- Verify `add()` still works identically (regression).

**Risk:** Low. No pipeline logic changes. Providers are behind a protocol.

---

### Stage 2 — Async Parsing Stage with Concurrent File Processing

**Files changed:**
- `packages/clew-ingestion/src/pipeline/base.py`
- `packages/clew-ingestion/src/pipeline/orchestrator.py`
- `packages/clew-ingestion/src/pipeline/stages/parsing.py`
- `packages/clew-ingestion/src/pipeline/context.py`
- `packages/clew-ingestion/src/ingest.py`
- `packages/clew-ingestion/tests/test_ingest.py`

**What changes:**

1. **Add `AsyncPipelineStage` protocol alongside existing `PipelineStage`.**
   ```python
   @runtime_checkable
   class AsyncPipelineStage(Protocol):
       name: str
       async def execute(self, context: IngestionContext) -> ProcessingResult: ...
   ```
   The orchestrator checks `isinstance(stage, AsyncPipelineStage)` and awaits
   if so; otherwise calls synchronously. This lets us convert stages one at a
   time without a flag day.

2. **Convert `IngestionPipeline.run()` to `async def run()`.**
   - The `for stage in self.stages` loop becomes:
     ```python
     for stage in self.stages:
         if isinstance(stage, AsyncPipelineStage):
             result = await stage.execute(context)
         else:
             result = stage.execute(context)
     ```
   - `ingest_repo()` wraps the call: `asyncio.run(pipeline.run(...))`.

3. **Convert `ParsingStage` to `AsyncPipelineStage`.**
   - The `execute()` method becomes `async def execute()`.
   - File processing uses an **async generator** that yields `(file_item, content)`
     pairs. Files are read from disk in a thread pool (`asyncio.to_thread`)
     to avoid blocking the event loop:
     ```python
     async def _file_reader(self, files: list[FileItem]):
         """Async generator that yields files with loaded content."""
         for file_item in files:
             content = await asyncio.to_thread(file_item.load_content)
             yield file_item, content
     ```
   - A **worker coroutine** processes one file: parse AST (via `to_thread`
     since tree-sitter is CPU/C), then collect the text+metadata pairs for
     that file's embeddings and graph writes.
   - The main loop gathers file results in **bounded batches**:
     ```python
     batch: list[tuple[str, dict]] = []  # (text, metadata) for embeddings
     graph_ops: list[Callable] = []       # deferred graph writes

     async for file_item, content in self._file_reader(context.files):
         # Parse AST (CPU, run in thread)
         definitions = await asyncio.to_thread(
             context.parser.parse_file, file_item.path, content.encode("utf-8")
         )

         # Collect file-level embed item
         batch.append((content[:2000], {file-level metadata}))
         # Collect graph op
         graph_ops.append(lambda: context.graph_store.create_file_node(...))

         # Collect definition-level items into batch + nodes list
         for defn in definitions:
             context.nodes.append(ParsedNode(...))

         # Flush batch when large enough
         if len(batch) >= BATCH_SIZE:
             point_ids = await context.vector_store.add_batch(batch)
             # Execute graph writes concurrently
             await asyncio.gather(*(
                 asyncio.to_thread(op) for op in graph_ops
             ))
             batch.clear()
             graph_ops.clear()

     # Final flush for remainder
     if batch:
         await context.vector_store.add_batch(batch)
         await asyncio.gather(*(asyncio.to_thread(op) for op in graph_ops))
     ```

4. **Thread-safety for `context.nodes`.**
   - Parsing is the *only* stage that appends to `context.nodes`.
   - Within this stage the appends happen in the `async for` loop body which
     runs on a single coroutine (no concurrent mutation). No lock needed.
   - If we later parallelize the inner loop with `TaskGroup`, we add a
     `threading.Lock` at that point. Not now.

5. **Pool management: `asyncio.to_thread` uses the default executor.**
   - We set a bounded `ThreadPoolExecutor(max_workers=4)` as the default
     executor at pipeline start:
     ```python
     loop = asyncio.get_running_loop()
     pool = ThreadPoolExecutor(max_workers=4, thread_name_prefix="clew-parse")
     loop.set_default_executor(pool)
     ```
   - The pool is shut down in `FinalizationStage` or via a `try/finally` in
     the orchestrator's `run()`.

**The key design choice:** we do NOT parallelize file processing across files
yet. The `async for` loop processes files one at a time, but the I/O within
each iteration (embedding, graph writes) is batched and concurrent. This is
safer and already captures most of the win (batch embeddings cut ~100s to ~2s).

**Tests:**
- Existing `test_ingest.py` is updated: `ingest_repo()` still returns an
  int exit code; the test doesn't need to know about async.
- New test: verify `_file_reader` async generator yields all files.
- New test: verify batch embedding is called with correct groupings.

**Risk:** Medium. The orchestrator changes are the most sensitive part. The
`isinstance` check for async/sync stages is the escape hatch — if anything
breaks, a stage can revert to sync.

---

### Stage 3 — Async Processing Stage + Batched Graph Writes

**Files changed:**
- `packages/clew-ingestion/src/graph.py`
- `packages/clew-ingestion/src/pipeline/stages/processing.py`
- `packages/clew-ingestion/src/pipeline/processors/definition.py`
- `packages/clew-ingestion/src/pipeline/processors/import_processor.py`
- `packages/clew-ingestion/src/pipeline/processors/call.py`
- `packages/clew-ingestion/src/pipeline/processors/registry.py`

**What changes:**

1. **Add `GraphStore.execute_batch()` for batched Cypher.**
   ```python
   def execute_batch(self, operations: list[tuple[str, dict]]) -> None:
       """Execute multiple Cypher queries in a single transaction."""
       with self.driver.session() as session:
           with session.begin_transaction() as tx:
               for query, params in operations:
                   tx.run(query, **params)
               tx.commit()
   ```
   This replaces N sessions with 1 transaction. The individual
   `create_*` methods remain for backward compat.

2. **Convert `ProcessingStage` to `AsyncPipelineStage`.**
   - Nodes are processed in batches (e.g. 50 at a time).
   - Definition nodes: collect their text for `add_batch()`, collect
     their Cypher ops for `execute_batch()`.
   - Import/Call nodes: collect Cypher ops only, flush via `execute_batch()`.
   - Flow:
     ```python
     async def execute(self, context):
         embed_batch = []
         graph_batch = []

         for node in context.nodes:
             if node.type == "definition":
                 embed_batch.append((node.content, {metadata}))
                 graph_batch.append((CREATE_CODE_NODE_CYPHER, {params}))
             elif node.type == "import":
                 graph_batch.append((CREATE_IMPORT_CYPHER, {params}))
             elif node.type == "call":
                 graph_batch.append((CREATE_CALL_CYPHER, {params}))

             if len(embed_batch) >= BATCH_SIZE:
                 point_ids = await context.vector_store.add_batch(embed_batch)
                 # Assign qdrant_ids back to the corresponding graph ops
                 self._patch_qdrant_ids(graph_batch, point_ids)
                 await asyncio.to_thread(
                     context.graph_store.execute_batch, graph_batch
                 )
                 embed_batch.clear()
                 graph_batch.clear()

         # Flush remainder
         ...
     ```

3. **The `NodeProcessorRegistry` and individual processors are NOT removed.**
   They remain available for the `sync.py` incremental path, which processes
   nodes one at a time. The batched path in `ProcessingStage` is a parallel
   code path chosen when running the full pipeline.

4. **Qdrant ID plumbing.** The tricky part: `DefinitionProcessor` today calls
   `vector_store.add()` and passes the returned `qdrant_id` to
   `graph_store.create_code_node()`. In the batched path, we need to:
   - Call `add_batch()` for all definition texts.
   - Get back a list of `point_ids` in order.
   - Pair each `point_id` with its corresponding graph Cypher params.
   - Then `execute_batch()` with the complete params.

   This is handled by keeping embed items and graph ops in parallel lists
   and zipping the IDs back after the batch call.

**Tests:**
- Unit test `execute_batch` with a mock Neo4j driver.
- Unit test the batched processing stage end-to-end with mock stores.
- Verify the `qdrant_id` values flow correctly from `add_batch` to graph ops.

**Risk:** Medium-high. The `qdrant_id` plumbing requires careful ordering.
Keeping the per-node processor path as a fallback means we can disable
batching with a config flag if needed.

---

## Pool Management & Cleanup (Cross-cutting)

This is called out as a priority, so here is the explicit strategy:

### Thread Pool
- **One pool.** Created in `IngestionPipeline.run()` as the asyncio default executor.
- **Bounded.** `max_workers=4` (configurable via `context.config`).
- **Named threads.** `thread_name_prefix="clew-parse"` for debuggability.
- **Deterministic shutdown.** `pool.shutdown(wait=True)` in a `finally` block
  in `run()`, guaranteeing cleanup even on exceptions:
  ```python
  async def run(self, repo_id, repo_path_or_url):
      pool = ThreadPoolExecutor(max_workers=4, thread_name_prefix="clew-parse")
      loop = asyncio.get_running_loop()
      loop.set_default_executor(pool)
      try:
          # ... execute stages ...
      finally:
          pool.shutdown(wait=True)
  ```

### Async Semaphores
- **Embedding semaphore.** `asyncio.Semaphore(8)` inside embedding providers,
  used only in `embed_batch` for the Ollama fan-out path. OpenAI's batch
  endpoint doesn't need a semaphore (single HTTP call).
- **Graph write semaphore.** Not needed — graph writes are batched into single
  transactions, not fanned out.

### Connection Cleanup
- **Neo4j driver.** Already cleaned up in `FinalizationStage`. No change.
- **Qdrant client.** Stateless HTTP. `flush()` in `FinalizationStage` ensures
  buffered points are written. No change.
- **httpx.AsyncClient.** In `OllamaEmbeddings`, change from creating a new
  client per `embed()` call to a **shared client** with explicit lifecycle:
  ```python
  class OllamaEmbeddings:
      def __init__(self, ...):
          self._client = httpx.AsyncClient(timeout=self.timeout)

      async def close(self):
          await self._client.aclose()
  ```
  Closed in `FinalizationStage`.

### Error Handling Under Concurrency
- **`asyncio.gather(..., return_exceptions=True)`** is used wherever we fan
  out, so one failed embedding doesn't cancel the batch.
- Failed items are logged and counted in `ProcessingResult` (PARTIAL status).
- No silent swallowing — every exception is logged at ERROR level before
  being aggregated.

---

## What We Are NOT Doing

1. **No `multiprocessing`.** The GIL blocks CPU parallelism for tree-sitter,
   but tree-sitter is already fast (~10ms/file). The overhead of IPC and
   serialization would exceed the gains.
2. **No cross-stage parallelism.** Stages have data dependencies. Pipelining
   (stage N+1 starts before stage N finishes) would add complexity for
   marginal gain.
3. **No distributed workers (Celery/Redis).** Out of scope. This plan
   targets single-process concurrency.
4. **No parallel file processing in Stage 2.** We batch I/O within the
   sequential file loop. If profiling shows the `async for` loop is still
   a bottleneck, we can add `TaskGroup` parallelism in a follow-up, with
   a lock on `context.nodes`.

---

## Expected Performance Impact

| Scenario (500-file repo) | Before | After (est.) | Speedup |
|---|---|---|---|
| Embedding generation | ~100s | ~2-4s (batch) | ~25-50x |
| Graph writes | ~15s | ~1-2s (batched txn) | ~10x |
| File I/O + parsing | ~5s | ~5s (unchanged) | 1x |
| **Total ingestion** | **~120s** | **~10-15s** | **~8-12x** |

---

## Ordering & Dependencies

```
Stage 1 (embeddings + vector)
   └──> Stage 2 (parsing stage, depends on add_batch from Stage 1)
           └──> Stage 3 (processing stage, depends on add_batch + async orchestrator from Stages 1-2)
```

Each stage should be one PR. Stage 1 has no dependencies on pipeline changes
and can be developed and tested in isolation.

---

## Testing Strategy

- **Each stage adds its own unit tests** for new/changed methods.
- **`test_ingest.py` integration test is updated in Stage 2** when the
  orchestrator goes async. The test's external interface (`ingest_repo()`
  returns an int) does not change.
- **No new test infrastructure needed.** `pytest-asyncio` may be added as a
  dev dependency for testing async methods directly.
- **Mock strategy unchanged.** Embedding providers and stores are already
  injected; mocking them in tests is straightforward.
