|
14 | 14 |
|
15 | 15 | import ast |
16 | 16 | import logging |
| 17 | +from concurrent.futures import ThreadPoolExecutor, as_completed |
17 | 18 | from dataclasses import dataclass, field |
18 | 19 | from datetime import datetime, timezone |
19 | 20 | from pathlib import Path |
|
25 | 26 |
|
26 | 27 | logger = logging.getLogger(__name__) |
27 | 28 |
|
| 29 | +#: Cap on concurrent LLM calls during the parallel polish phase. |
| 30 | +#: Sized to comfortably fit under Anthropic's per-minute rate |
| 31 | +#: limits while still saturating the LLM-bound wall time of a |
| 32 | +#: typical ``regenerate --all-kinds`` run. |
| 33 | +_POLISH_MAX_WORKERS = 4 |
| 34 | + |
| 35 | + |
| 36 | +def _parallel_polish( |
| 37 | + pending: list[tuple[str, str, Path]], |
| 38 | + feature: object, |
| 39 | + source_info: object, |
| 40 | + use_rag: bool, |
| 41 | +) -> dict[str, tuple[str, Path]]: |
| 42 | + """Polish a batch of rendered templates concurrently. |
| 43 | +
|
| 44 | + Args: |
| 45 | + pending: List of (depth, rendered_content, out_path) tuples. |
| 46 | + feature: Feature being documented (read-only, thread-safe). |
| 47 | + source_info: Extracted source info (read-only, thread-safe). |
| 48 | + use_rag: Whether to use RAG grounding during polish. |
| 49 | +
|
| 50 | + Returns: |
| 51 | + Mapping of depth -> (polished_content, out_path). Raises |
| 52 | + the first exception encountered (propagated from the future). |
| 53 | + """ |
| 54 | + |
| 55 | + def _task(depth: str, content: str, out_path: Path) -> tuple[str, str, Path]: |
| 56 | + polished = _maybe_polish( |
| 57 | + content, |
| 58 | + feature, # type: ignore[arg-type] |
| 59 | + source_info, # type: ignore[arg-type] |
| 60 | + template_type=depth, |
| 61 | + use_rag=use_rag, |
| 62 | + ) |
| 63 | + return depth, polished, out_path |
| 64 | + |
| 65 | + results: dict[str, tuple[str, Path]] = {} |
| 66 | + workers = min(len(pending), _POLISH_MAX_WORKERS) |
| 67 | + with ThreadPoolExecutor(max_workers=workers) as executor: |
| 68 | + futures = { |
| 69 | + executor.submit(_task, depth, content, out_path): depth |
| 70 | + for depth, content, out_path in pending |
| 71 | + } |
| 72 | + for future in as_completed(futures): |
| 73 | + depth, polished, out_path = future.result() |
| 74 | + results[depth] = (polished, out_path) |
| 75 | + return results |
| 76 | + |
| 77 | + |
28 | 78 | #: Core progressive-depth template kinds. These form the |
29 | 79 | #: progressive disclosure path that attune-help renders: |
30 | 80 | #: concept → task → reference. They are generated by |
@@ -234,6 +284,9 @@ def generate_feature_templates( |
234 | 284 | ", ".join(feature.doc_paths[1:]), |
235 | 285 | ) |
236 | 286 |
|
| 287 | + # Phase 1: render all templates (fast Jinja2, sequential). |
| 288 | + # Determines which depths are active and builds the rendered skeleton. |
| 289 | + pending: list[tuple[str, str, Path]] = [] |
237 | 290 | for depth in target_depths: |
238 | 291 | if depth not in _ALL_TEMPLATE_NAMES: |
239 | 292 | logger.warning("Unknown template kind '%s', skipping", depth) |
@@ -278,17 +331,15 @@ def generate_feature_templates( |
278 | 331 | source_hash=source_hash, |
279 | 332 | source_info=source_info, |
280 | 333 | ) |
| 334 | + pending.append((depth, content, out_path)) |
281 | 335 |
|
282 | | - # LLM polish pass — improves writing quality |
283 | | - content = _maybe_polish( |
284 | | - content, |
285 | | - feature, |
286 | | - source_info, |
287 | | - template_type=depth, |
288 | | - use_rag=use_rag, |
289 | | - ) |
| 336 | + # Phase 2: LLM polish — run all depths concurrently. |
| 337 | + polished = _parallel_polish(pending, feature, source_info, use_rag) |
290 | 338 |
|
291 | | - out_path.write_text(content, encoding="utf-8") |
| 339 | + # Phase 3: write results in original depth order. |
| 340 | + for depth, content, out_path in pending: |
| 341 | + final_content, _ = polished[depth] |
| 342 | + out_path.write_text(final_content, encoding="utf-8") |
292 | 343 | result.templates.append( |
293 | 344 | GeneratedTemplate( |
294 | 345 | feature=feature.name, |
|
0 commit comments