Fix update json getting Ailly'd

On subsequent runs, update json files were being processed by Ailly
because of a pattern match on the update files. This changes the update
json naming pattern and makes some other refactors.

Author: cpyle0819

aws_doc_sdk_examples_tools/agent/bin/main.py

@@ -1,27 +1,31 @@
 from pathlib import Path
 from subprocess import run
 from typing import List
-import time
+from typing_extensions import Annotated
 from datetime import timedelta, datetime
 
+import json
 import logging
+import time
 import typer
 
+from aws_doc_sdk_examples_tools.agent.shared_constants import BATCH_PREFIX
 from aws_doc_sdk_examples_tools.agent.make_prompts import make_prompts
 from aws_doc_sdk_examples_tools.agent.process_ailly_files import process_ailly_files
 from aws_doc_sdk_examples_tools.agent.update_doc_gen import update_doc_gen
 from aws_doc_sdk_examples_tools.yaml_writer import prepare_write, write_many
 
 logging.basicConfig(
-    level=logging.INFO, filename=f"lliam-run-{datetime.now()}.log", filemode="w"
+    level=logging.INFO,
+    filename=f"lliam-run-{datetime.now().strftime('%Y%m%d_%H%M%S')}.log",
+    filemode="w",
 )
 logger = logging.getLogger(__name__)
 
 app = typer.Typer()
 
 AILLY_DIR = ".ailly_iam_policy"
 AILLY_DIR_PATH = Path(AILLY_DIR)
-IAM_UPDATES_PATH = AILLY_DIR_PATH / "iam_updates.json"
 
 
 def format_duration(seconds: float) -> str:
@@ -30,72 +34,160 @@ def format_duration(seconds: float) -> str:
 
 
 @app.command()
-def update(
-    iam_tributary_root: str,
-    system_prompts: List[str] = [],
-    skip_generation: bool = False,
+def create_prompts(iam_tributary_root: str, system_prompts: List[str] = []):
+    doc_gen_root = Path(iam_tributary_root)
+
+    make_prompts(
+        doc_gen_root=doc_gen_root,
+        system_prompts=system_prompts,
+        out_dir=AILLY_DIR_PATH,
+        language="IAMPolicyGrammar",
+    )
+
+
+def parse_batch_names(batch_nums_str: str) -> List[str]:
+    """
+    Parse batch numbers from a string.
+    """
+    if not batch_nums_str:
+        return []
+
+    result = []
+    parts = batch_nums_str.split(",")
+
+    for part in parts:
+        part = part.strip()
+        if ".." in part:
+            start, end = part.split("..")
+            start_num = int(start.strip())
+            end_num = int(end.strip())
+            result.extend(range(start_num, end_num + 1))
+        else:
+            result.append(int(part))
+
+    batch_nums = sorted(list(set(result)))
+    return map(lambda b: f"{BATCH_PREFIX}{b:03}", batch_nums)
+
+
+def resolve_requested_batches(batch_names: List[str]) -> List[Path]:
+    if not batch_names:
+        batch_paths = [
+            p
+            for p in AILLY_DIR_PATH.iterdir()
+            if p.is_dir() and p.name.startswith(BATCH_PREFIX)
+        ]
+
+        return batch_paths
+
+    batch_paths = []
+
+    for batch_name in batch_names:
+        batch_path = Path(AILLY_DIR_PATH / batch_name)
+        if not batch_path.exists():
+            raise FileNotFoundError(batch_path)
+        if not batch_path.is_dir():
+            raise NotADirectoryError(batch_path)
+        batch_paths.append(batch_path)
+
+    return batch_paths
+
+
+@app.command()
+def run_ailly(
+    batch_nums: Annotated[
+        str,
+        typer.Option(
+            help="Batch numbers to process (e.g., '33', '33..35', '33,35,37')"
+        ),
+    ] = None,
 ) -> None:
     """
-    Generate new IAM policy metadata for a tributary.
+    Run ailly to generate IAM policy content and process the results.
+    If batch_nums is specified, only those batches will be processed.
+    If batch_nums is omitted, all batches will be processed.
     """
-    doc_gen_root = Path(iam_tributary_root)
+    requested_batches = parse_batch_names(batch_nums)
+    resolved_batches = resolve_requested_batches(requested_batches)
 
-    if not skip_generation:
-        make_prompts(
-            doc_gen_root=doc_gen_root,
-            system_prompts=system_prompts,
-            out_dir=AILLY_DIR_PATH,
-            language="IAMPolicyGrammar",
+    if resolved_batches:
+        total_start_time = time.time()
+
+        for batch in resolved_batches:
+            run_ailly_single_batch(batch)
+
+        total_end_time = time.time()
+        total_duration = total_end_time - total_start_time
+        num_batches = len(resolved_batches)
+        logger.info(
+            f"[TIMECHECK] {num_batches} batches took {format_duration(total_duration)} to run"
         )
 
-        batch_dirs = [
-            d.name
-            for d in AILLY_DIR_PATH.iterdir()
-            if d.is_dir() and d.name.startswith("batch_")
-        ]
 
-        if batch_dirs:
-            total_start_time = time.time()
-
-            for batch_dir in sorted(batch_dirs):
-                batch_start_time = time.time()
-
-                cmd = [
-                    "ailly",
-                    "--max-depth",
-                    "10",
-                    "--root",
-                    AILLY_DIR,
-                    str(batch_dir),
-                ]
-                logger.info(f"Running {cmd}")
-                run(cmd)
-
-                batch_end_time = time.time()
-                batch_duration = batch_end_time - batch_start_time
-                batch_num = batch_dir.replace("batch_", "")
-                logger.info(
-                    f"[TIMECHECK] Batch {batch_num} took {format_duration(batch_duration)} to run"
-                )
-
-            total_end_time = time.time()
-            total_duration = total_end_time - total_start_time
-            num_batches = len(batch_dirs)
-            logger.info(
-                f"[TIMECHECK] {num_batches} batches took {format_duration(total_duration)} to run"
-            )
-
-    logger.info("Processing generated content")
-    process_ailly_files(
-        input_dir=str(AILLY_DIR_PATH), output_file=str(IAM_UPDATES_PATH)
+@app.command()
+def update_reservoir(
+    iam_tributary_root: str,
+    batch_nums: Annotated[
+        str,
+        typer.Option(
+            help="Batch numbers to process (e.g., '33', '33..35', '33,35,37')"
+        ),
+    ] = None,
+) -> None:
+    """
+    Update the doc_gen reservoir with processed IAM policy updates.
+    If batch_nums is specified, only those batches will be processed.
+    If batch_nums is omitted, all available update files will be processed.
+    """
+    doc_gen_root = Path(iam_tributary_root)
+    batch_names = parse_batch_names(batch_nums) if batch_nums else None
+
+    update_files = (
+        [AILLY_DIR_PATH / f"updates_{batch}.json" for batch in batch_names]
+        if batch_names
+        else list(AILLY_DIR_PATH.glob(f"updates_{BATCH_PREFIX}*.json"))
     )
 
-    doc_gen = update_doc_gen(
-        doc_gen_root=doc_gen_root, iam_updates_path=IAM_UPDATES_PATH
+    if not update_files:
+        logger.warning("No IAM update files found to process")
+        return
+
+    for update_file in sorted(update_files):
+        if update_file.exists():
+            logger.info(f"Processing updates from {update_file.name}")
+            updates = json.loads(update_file.read_text())
+
+            doc_gen = update_doc_gen(doc_gen_root=doc_gen_root, updates=updates)
+
+            writes = prepare_write(doc_gen.examples)
+            write_many(doc_gen_root, writes)
+        else:
+            logger.warning(f"Update file not found: {update_file}")
+
+
+def run_ailly_single_batch(batch: Path) -> None:
+    """Run ailly and process files for a single batch."""
+    batch_start_time = time.time()
+    iam_updates_path = AILLY_DIR_PATH / f"updates_{batch.name}.json"
+
+    cmd = [
+        "ailly",
+        "--max-depth",
+        "10",
+        "--root",
+        str(AILLY_DIR_PATH),
+        batch.name,
+    ]
+    logger.info(f"Running {cmd}")
+    run(cmd)
+
+    batch_end_time = time.time()
+    batch_duration = batch_end_time - batch_start_time
+    logger.info(
+        f"[TIMECHECK] {batch.name} took {format_duration(batch_duration)} to run"
     )
 
-    writes = prepare_write(doc_gen.examples)
-    write_many(doc_gen_root, writes)
+    logger.info(f"Processing generated content for {batch.name}")
+    process_ailly_files(input_dir=batch, output_file=iam_updates_path)
 
 
 if __name__ == "__main__":

aws_doc_sdk_examples_tools/agent/make_prompts.py

@@ -3,12 +3,15 @@
 import logging
 import os
 import yaml
+from itertools import islice
 from pathlib import Path
-from typing import List
+from typing import Any, Generator, Iterable, List, Tuple
 
-from aws_doc_sdk_examples_tools.doc_gen import DocGen
+from aws_doc_sdk_examples_tools.agent.shared_constants import BATCH_PREFIX
+from aws_doc_sdk_examples_tools.doc_gen import DocGen, Example
 
 DEFAULT_METADATA_PREFIX = "DEFAULT"
+BATCH_SIZE = 150
 
 logger = logging.getLogger(__name__)
 
@@ -20,11 +23,21 @@ def make_doc_gen(root: Path) -> DocGen:
     return doc_gen
 
 
+def batched(iterable: Iterable, n: int) -> Generator[Tuple, Any, None]:
+    "Batch data into tuples of length n. The last batch may be shorter."
+    # batched('ABCDEFG', 3) --> ABC DEF G
+    if n < 1:
+        raise ValueError("n must be at least one")
+    it = iter(iterable)
+    while batch := tuple(islice(it, n)):
+        yield batch
+
+
 def write_prompts(doc_gen: DocGen, out_dir: Path, language: str) -> None:
     examples = doc_gen.examples
     snippets = doc_gen.snippets
 
-    filtered_examples = []
+    filtered_examples: Tuple[str, Example] = []
     for example_id, example in examples.items():
         # TCXContentAnalyzer prefixes new metadata title/title_abbrev entries with
         # the DEFAULT_METADATA_PREFIX. Checking this here to make sure we're only
@@ -36,23 +49,15 @@ def write_prompts(doc_gen: DocGen, out_dir: Path, language: str) -> None:
         ):
             filtered_examples.append((example_id, example))
 
-    batch_size = 150
-    total_examples = len(filtered_examples)
-    num_batches = (total_examples + batch_size - 1) // batch_size
-
     logger.info(
-        f"Splitting {total_examples} examples into {num_batches} batches of {batch_size}"
+        f"Splitting {len(filtered_examples)} examples into batches of {BATCH_SIZE}"
     )
 
-    for batch_num in range(num_batches):
-        batch_dir = out_dir / f"batch_{(batch_num + 1):03}"
+    for batch_num, batch in enumerate(batched(filtered_examples, BATCH_SIZE)):
+        batch_dir = out_dir / f"{BATCH_PREFIX}{(batch_num + 1):03}"
         batch_dir.mkdir(exist_ok=True)
 
-        start_idx = batch_num * batch_size
-        end_idx = min((batch_num + 1) * batch_size, total_examples)
-
-        for i in range(start_idx, end_idx):
-            example_id, example = filtered_examples[i]
+        for example_id, example in batch:
             prompt_path = batch_dir / f"{example_id}.md"
 
             try:

aws_doc_sdk_examples_tools/agent/process_ailly_files.py

@@ -11,7 +11,6 @@
 from pathlib import Path
 from typing import Any, Dict, List, Set
 
-logging.basicConfig(level=logging.WARNING)
 logger = logging.getLogger(__name__)
 
 EXPECTED_KEYS: Set[str] = set(["title", "title_abbrev"])
@@ -91,7 +90,7 @@ def parse_ailly_file(
 
 
 def process_ailly_files(
-    input_dir: str, output_file: str, file_pattern: str = "*.md.ailly.md"
+    input_dir: Path, output_file: Path, file_pattern: str = "*.md.ailly.md"
 ) -> None:
     """
     Process all .md.ailly.md files in the input directory and write the results as JSON to the output file.
@@ -102,10 +101,9 @@ def process_ailly_files(
         file_pattern: Pattern to match files (default: "*.md.ailly.md")
     """
     results = []
-    input_path = Path(input_dir)
 
     try:
-        for file_path in input_path.rglob(file_pattern):
+        for file_path in input_dir.rglob(file_pattern):
             logger.info(f"Processing file: {file_path}")
             parsed_data = parse_ailly_file(str(file_path))
             if parsed_data:
@@ -115,7 +113,7 @@ def process_ailly_files(
             json.dump(results, out_file, indent=2)
 
         logger.info(
-            f"Successfully processed {len(results)} files. Output written to {output_file}"
+            f"Successfully processed {len(results)} files. Output written to {output_file.name}"
         )
 
     except Exception as e:

aws_doc_sdk_examples_tools/agent/shared_constants.py

@@ -0,0 +1 @@
+BATCH_PREFIX = "batch_"