Batch Ailly calls from Lliam

Rather than make one Ailly call, this batches them into groups of 150 policies and calls Ailly on each batch.

Also change the default prefix from [DEFAULT] to DEFAULT. The former was causing build failures on the TCA side.

Author: cpyle0819

.gitignore

@@ -7,3 +7,4 @@ __pycache__
 build/
 dist/
 .ailly_iam_policy
+*.log

aws_doc_sdk_examples_tools/agent/bin/main.py

@@ -1,21 +1,35 @@
 from pathlib import Path
 from subprocess import run
 from typing import List
+import time
+from datetime import timedelta, datetime
 
+import logging
 import typer
 
 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")
+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:
+    """Format duration in seconds to hours:minutes:seconds format."""
+    td = timedelta(seconds=seconds)
+    hours, remainder = divmod(td.total_seconds(), 3600)
+    minutes, seconds = divmod(remainder, 60)
+    return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d}"
+
+
 @app.command()
 def update(
     iam_tributary_root: str,
@@ -34,8 +48,51 @@ def update(
             out_dir=AILLY_DIR_PATH,
             language="IAMPolicyGrammar",
         )
-        run(["npx @ailly/[email protected]", "--root", AILLY_DIR])
 
+        # Find all batch directories
+        batch_dirs = [
+            d.name
+            for d in AILLY_DIR_PATH.iterdir()
+            if d.is_dir() and d.name.startswith("batch_")
+        ]
+        if batch_dirs:
+            # Run ailly on each batch directory
+            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")
+        else:
+            # Fallback to running on the main directory if no batches found
+            run(
+                [
+                    "/Users/corepyle/Documents/repos/ailly/cli/index.js",
+                    "--root",
+                    AILLY_DIR,
+                ]
+            )
+
+    logger.info("Processing generated content")
     process_ailly_files(
         input_dir=str(AILLY_DIR_PATH), output_file=str(IAM_UPDATES_PATH)
     )

aws_doc_sdk_examples_tools/agent/make_prompts.py

@@ -2,17 +2,14 @@
 
 import logging
 import os
+import yaml
 from pathlib import Path
 from typing import List
-import yaml
-
-from aws_doc_sdk_examples_tools.doc_gen import DocGen, Snippet
 
-DEFAULT_METADATA_PREFIX = "[DEFAULT]"
+from aws_doc_sdk_examples_tools.doc_gen import DocGen
 
+DEFAULT_METADATA_PREFIX = "DEFAULT"
 
-# Setup logging
-logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
@@ -26,6 +23,8 @@ def make_doc_gen(root: Path) -> DocGen:
 def write_prompts(doc_gen: DocGen, out_dir: Path, language: str) -> None:
     examples = doc_gen.examples
     snippets = doc_gen.snippets
+
+    filtered_examples = []
     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
@@ -35,30 +34,56 @@ def write_prompts(doc_gen: DocGen, out_dir: Path, language: str) -> None:
         if title.startswith(DEFAULT_METADATA_PREFIX) and title_abbrev.startswith(
             DEFAULT_METADATA_PREFIX
         ):
-            prompt_path = out_dir / f"{example_id}.md"
-            snippet_key = (
-                example.languages[language]
-                .versions[0]
-                .excerpts[0]
-                .snippet_files[0]
-                .replace("/", ".")
-            )
-            snippet = snippets[snippet_key]
-            prompt_path.write_text(snippet.code, encoding="utf-8")
+            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}"
+    )
+
+    for batch_num in range(num_batches):
+        batch_dir = out_dir / f"batch_{(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]
+            prompt_path = batch_dir / f"{example_id}.md"
+
+            try:
+                snippet_key = (
+                    example.languages[language]
+                    .versions[0]
+                    .excerpts[0]
+                    .snippet_files[0]
+                    .replace("/", ".")
+                )
+                snippet = snippets[snippet_key]
+                prompt_path.write_text(snippet.code, encoding="utf-8")
+            except (KeyError, IndexError, AttributeError) as e:
+                logger.warning(f"Error processing example {example_id}: {e}")
 
 
 def setup_ailly(system_prompts: List[str], out_dir: Path) -> None:
     """Create the .aillyrc configuration file."""
     fence = "---"
     options = {
         "isolated": "true",
-        "mcp": {
-            "awslabs.aws-documentation-mcp-server": {
-                "type": "stdio",
-                "command": "uvx",
-                "args": ["awslabs.aws-documentation-mcp-server@latest"],
-            }
-        },
+        "overwrite": "true"
+        # MCP assistance did not produce noticeably different results, but it was
+        # slowing things down by 10x. Disabled for now.
+        # "mcp": {
+        #     "awslabs.aws-documentation-mcp-server": {
+        #         "type": "stdio",
+        #         "command": "uvx",
+        #         "args": ["awslabs.aws-documentation-mcp-server@latest"],
+        #     }
+        # },
     }
     options_block = yaml.dump(options).strip()
     prompts_block = "\n".join(system_prompts)

aws_doc_sdk_examples_tools/agent/process_ailly_files.py

@@ -106,7 +106,7 @@ def process_ailly_files(
     input_path = Path(input_dir)
 
     try:
-        for file_path in input_path.glob(file_pattern):
+        for file_path in input_path.rglob(file_pattern):
             logger.info(f"Processing file: {file_path}")
             parsed_data = parse_ailly_file(str(file_path))
             if parsed_data:

setup.py

@@ -19,5 +19,6 @@
         "pathspec==0.11.2",
         "PyYAML==6.0.1",
         "yamale==4.0.4",
+        "typer==0.16.0",
     ],
 )