Add packages option to run-ailly

Author: cpyle0819

aws_doc_sdk_examples_tools/lliam/domain/commands.py

@@ -17,6 +17,7 @@ class CreatePrompts(Command):
 @dataclass
 class RunAilly(Command):
     batches: List[str]
+    packages: List[str]
 
 
 @dataclass

aws_doc_sdk_examples_tools/lliam/entry_points/lliam_app.py

@@ -5,10 +5,7 @@
 import logging
 import typer
 
-from aws_doc_sdk_examples_tools.lliam.config import (
-    AILLY_DIR,
-    BATCH_PREFIX
-)
+from aws_doc_sdk_examples_tools.lliam.config import AILLY_DIR, BATCH_PREFIX
 from aws_doc_sdk_examples_tools.lliam.domain import commands
 from aws_doc_sdk_examples_tools.lliam.service_layer import messagebus, unit_of_work
 
@@ -38,18 +35,21 @@ def create_prompts(iam_tributary_root: str, system_prompts: List[str] = []):
 def run_ailly(
     batches: Annotated[
         Optional[str],
-        typer.Option(
-            help="Batch names to process (comma-separated list)"
-        ),
+        typer.Option(help="Batch names to process (comma-separated list)"),
+    ] = None,
+    packages: Annotated[
+        Optional[str], typer.Option(help="Comma delimited list of packages to update")
     ] = None,
 ) -> None:
     """
     Run ailly to generate IAM policy content and process the results.
     If batches is specified, only those batches will be processed.
     If batches is omitted, all batches will be processed.
+    If packages is specified, only those packages will be processed.
     """
     requested_batches = parse_batch_names(batches)
-    cmd = commands.RunAilly(batches=requested_batches)
+    package_names = parse_package_names(packages)
+    cmd = commands.RunAilly(batches=requested_batches, packages=package_names)
     messagebus.handle(cmd)
 
 
@@ -58,9 +58,7 @@ def update_reservoir(
     iam_tributary_root: str,
     batches: Annotated[
         Optional[str],
-        typer.Option(
-            help="Batch names to process (comma-separated list)"
-        ),
+        typer.Option(help="Batch names to process (comma-separated list)"),
     ] = None,
     packages: Annotated[
         Optional[str], typer.Option(help="Comma delimited list of packages to update")
@@ -86,7 +84,7 @@ def parse_batch_names(batch_names_str: Optional[str]) -> List[str]:
     """
     if not batch_names_str:
         return []
-    
+
     batch_names = []
 
     for name in batch_names_str.split(","):

aws_doc_sdk_examples_tools/lliam/service_layer/messagebus.py

@@ -12,9 +12,7 @@
 Message = commands.Command
 
 
-def handle(
-    message: commands.Command, uow: Optional[unit_of_work.FsUnitOfWork] = None
-):
+def handle(message: commands.Command, uow: Optional[unit_of_work.FsUnitOfWork] = None):
     queue = [message]
 
     while queue:
@@ -25,9 +23,7 @@ def handle(
             raise Exception(f"{message} was not a Command")
 
 
-def handle_command(
-    command: commands.Command, uow: Optional[unit_of_work.FsUnitOfWork]
-):
+def handle_command(command: commands.Command, uow: Optional[unit_of_work.FsUnitOfWork]):
     handler = COMMAND_HANDLERS[type(command)]
     handler(command, uow)
 

aws_doc_sdk_examples_tools/lliam/service_layer/run_ailly.py

@@ -13,6 +13,14 @@
     BATCH_PREFIX,
 )
 
+AILLY_CMD_BASE = [
+    "ailly",
+    "--max-depth",
+    "10",
+    "--root",
+    str(AILLY_DIR_PATH),
+]
+
 logger = logging.getLogger(__file__)
 
 
@@ -23,7 +31,7 @@ def handle_run_ailly(cmd: RunAilly, uow: None):
         total_start_time = time.time()
 
         for batch in resolved_batches:
-            run_ailly_single_batch(batch)
+            run_ailly_single_batch(batch, cmd.packages)
 
         total_end_time = time.time()
         total_duration = total_end_time - total_start_time
@@ -56,19 +64,27 @@ def resolve_requested_batches(batch_names: List[str]) -> List[Path]:
     return batch_paths
 
 
-def run_ailly_single_batch(batch: Path) -> None:
+def run_ailly_single_batch(batch: Path, packages: List[str] = None) -> 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,
-    ]
+    if packages:
+        paths = []
+        for package in packages:
+            package_files = [
+                f"{batch.name}/{p.name}" for p in batch.glob(f"*{package}*.md")
+            ]
+            paths.extend(package_files)
+
+        if not paths:
+            logger.warning(f"No matching files found for packages: {packages}")
+            return
+
+        cmd = AILLY_CMD_BASE + paths
+    else:
+        cmd = AILLY_CMD_BASE + [batch.name]
+
     logger.info(f"Running {cmd}")
     run(cmd)
 
@@ -79,7 +95,9 @@ def run_ailly_single_batch(batch: Path) -> None:
     )
 
     logger.info(f"Processing generated content for {batch.name}")
-    process_ailly_files(input_dir=batch, output_file=iam_updates_path)
+    process_ailly_files(
+        input_dir=batch, output_file=iam_updates_path, packages=packages
+    )
 
 
 EXPECTED_KEYS: Set[str] = set(["title", "title_abbrev"])
@@ -177,7 +195,10 @@ def parse_package_name(policy_update: Dict[str, str]) -> Optional[str]:
 
 
 def process_ailly_files(
-    input_dir: Path, output_file: Path, file_pattern: str = "*.md.ailly.md"
+    input_dir: Path,
+    output_file: Path,
+    file_pattern: str = "*.md.ailly.md",
+    packages: List[str] = None,
 ) -> None:
     """
     Process all .md.ailly.md files in the input directory and write the results as JSON to the output file.
@@ -186,6 +207,7 @@ def process_ailly_files(
         input_dir: Directory containing .md.ailly.md files
         output_file: Path to the output JSON file
         file_pattern: Pattern to match files (default: "*.md.ailly.md")
+        packages: Optional list of packages to filter by
     """
     results = defaultdict(list)
 
@@ -197,6 +219,13 @@ def process_ailly_files(
                 package_name = parse_package_name(policy_update)
                 if not package_name:
                     raise TypeError(f"Could not get package name from policy update.")
+
+                if packages and package_name not in packages:
+                    logger.info(
+                        f"Skipping package {package_name} (not in requested packages)"
+                    )
+                    continue
+
                 results[package_name].append(policy_update)
 
         with open(output_file, "w", encoding="utf-8") as out_file: