Add runner for iam update agent (#163)

Author: cpyle0819

aws_doc_sdk_examples_tools/agent/README.md

@@ -1,102 +1,74 @@
 # Ailly Prompt Workflow
 
-This project provides three scripts to **generate** Ailly prompts, **run** them, **parse** the results, and **update** a DocGen instance.
-
-## Overview
-
-1. **Generate** Ailly prompts from DocGen snippets (`make_prompts.py`).
-2. **Run** the Ailly CLI on the generated prompts (`npx @ailly/cli`).
-3. **Parse** the Ailly outputs into structured JSON (`parse_json_files.py`).
-4. **Update** your DocGen project with the parsed results (`update_doc_gen.py`).
-
----
-
-## Prerequisites
-
-- Python 3.8+
-- Node.js & npm (for `npx`)
+This project automates the process of generating, running, parsing, and applying [Ailly](https://www.npmjs.com/package/@ailly/cli) prompt outputs to an AWS DocGen project. It combines all steps into one streamlined command using a single Python script.
 
 ---
 
-## Step 1: Generate Ailly prompts
-
-Use `make_prompts.py` to create a directory of Markdown files and a `.aillyrc` configuration file.
+## 📦 Overview
 
-```bash
-python make_prompts.py \
-  --doc-gen-root /path/to/your/docgen/project \
-  --system-prompts "You are a helpful assistant..." \
-  --out .ailly_prompts
-```
+This tool:
+1. **Generates** Ailly prompts from DocGen snippets.
+2. **Runs** Ailly CLI to get enhanced metadata.
+3. **Parses** Ailly responses into structured JSON.
+4. **Updates** your DocGen examples with the new metadata.
 
-**Arguments:**
-- `--doc-gen-root`: Path to your DocGen project root.
-- `--system-prompts`: One or more system prompts, either as strings or file paths.
-- `--out`: (Optional) Output directory. Defaults to `.ailly_prompts`.
+All of this is done with one command.
 
 ---
 
-## Step 2: Run Ailly CLI
+## ✅ Prerequisites
 
-Run Ailly on the generated prompts:
-
-```bash
-npx @ailly/cli --root .ailly_prompts
-```
-
-This will create `{file}.ailly.md` output files in the `.ailly_prompts` directory (or whatever output directory you specified).
+- Python 3.8+
+- Node.js and npm (for `npx`)
+- A DocGen project directory
 
 ---
 
-## Step 3: Parse Ailly output
+## 🚀 Usage
 
-Parse the `{name}.ailly.md` files into JSON using `parse_json_files.py`:
+From your project root, run:
 
 ```bash
-python parse_json_files.py .ailly_prompts/*.ailly.md --out example_updates.json
+python -m aws_doc_sdk_examples_tools.agent.bin.main \
+  /path/to/your/docgen/project \
+  --system-prompts path/to/system_prompt.txt
 ```
 
-**Arguments:**
-- Positional: List of files to parse (e.g., `*.ailly.md`).
-- `--out`: (Optional) Path for the JSON output. Defaults to `out.json`.
+### 🔧 Arguments
+
+Run `python -m aws_doc_sdk_examples_tools.agent.bin.main --help` for more info.
 
 ---
 
-## Step 4: Update DocGen with Ailly results
+## 🗂 What This Does
 
-Use `update_doc_gen.py` to load the parsed data back into your `DocGen` project:
+Under the hood, this script:
 
-```bash
-python update_doc_gen.py \
-  --doc-gen-root /path/to/your/docgen/project \
-  --updates-path example_updates.json
-```
+1. Creates a directory `.ailly_iam_policy` containing:
+   - One Markdown file per snippet.
+   - A `.aillyrc` configuration file.
 
-**Arguments:**
-- `--doc-gen-root`: Path to the root of your DocGen project.
-- `--updates-path`: JSON file generated in Step 3 (default: `example_updates.json`).
+2. Runs `npx @ailly/cli` to generate `.ailly.md` outputs.
 
-This will update the `title`, `title_abbrev`, and `synopsis` fields in the corresponding DocGen examples.
+3. Parses the Ailly `.ailly.md` files into a single `iam_updates.json` file.
+
+4. Updates each matching `Example` in the DocGen instance with:
+   - `title`
+   - `title_abbrev`
+   - `synopsis`
 
 ---
 
-## Example Full Workflow
+## 💡 Example
 
 ```bash
-# Step 1: Generate prompts
-python make_prompts.py \
-  --doc-gen-root ~/projects/aws-docgen-root \
-  --system-prompts system_prompt.txt \
-  --out .ailly_prompts
-
-# Step 2: Run Ailly
-npx @ailly/cli --root .ailly_prompts
-
-# Step 3: Parse results
-python parse_json_files.py .ailly_prompts/*.ailly.md --out example_updates.json
-
-# Step 4: Update DocGen
-python update_doc_gen.py \
-  --doc-gen-root ~/projects/aws-docgen-root \
-  --updates-path example_updates.json
+python -m aws_doc_sdk_examples_tools.agent.bin.main \
+  ~/projects/AWSIAMPolicyExampleReservoir \
+  --system-prompts prompts/system_prompt.txt
 ```
+
+This will:
+- Write prompts and config to `.ailly_iam_policy/`
+- Run Ailly and capture results
+- Parse and save output as `.ailly_iam_policy/iam_updates.json`
+- Apply updates to your DocGen examples

aws_doc_sdk_examples_tools/agent/bin/main.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+from subprocess import run
+from typing import List
+
+import typer
+
+from aws_doc_sdk_examples_tools.agent.make_prompts import main as make_prompts
+from aws_doc_sdk_examples_tools.agent.parse_json_files import main as parse_json_files
+from aws_doc_sdk_examples_tools.agent.update_doc_gen import main as update_doc_gen
+
+app = typer.Typer()
+
+AILLY_DIR = ".ailly_iam_policy"
+AILLY_DIR_PATH = Path(AILLY_DIR)
+IAM_UPDATES_PATH = AILLY_DIR_PATH / "iam_updates.json"
+
+
+def get_ailly_files(dir: Path) -> List[Path]:
+    return [
+        file
+        for file in dir.iterdir()
+        if file.is_file() and file.name.endswith(".ailly.md")
+    ]
+
+
+@app.command()
+def update(iam_tributary_root: str, system_prompts: List[str] = []) -> None:
+    doc_gen_root = Path(iam_tributary_root)
+    make_prompts(
+        doc_gen_root=doc_gen_root, system_prompts=system_prompts, out=AILLY_DIR_PATH
+    )
+    run(["npx", "@ailly/cli", "--root", AILLY_DIR])
+    file_paths = get_ailly_files(AILLY_DIR_PATH)
+    parse_json_files(file_paths=file_paths, out=IAM_UPDATES_PATH)
+    update_doc_gen(doc_gen_root=doc_gen_root, iam_updates_path=IAM_UPDATES_PATH)
+
+
+if __name__ == "__main__":
+    app()

aws_doc_sdk_examples_tools/agent/make_prompts.py

@@ -1,10 +1,9 @@
 #!/usr/bin/env python
 
-import argparse
 import logging
 import os
 from pathlib import Path
-from typing import Dict, List
+from typing import List
 
 from aws_doc_sdk_examples_tools.doc_gen import DocGen, Snippet
 
@@ -53,7 +52,7 @@ def setup_ailly(system_prompts: List[str], out: Path) -> None:
     aillyrc_path.write_text(content, encoding="utf-8")
 
 
-def parse_prompts_arg(values: List[str]) -> List[str]:
+def read_system_prompts(values: List[str]) -> List[str]:
     """Parse system prompts from a list of strings or file paths."""
     prompts = []
     for value in values:
@@ -70,38 +69,10 @@ def validate_root_path(doc_gen_root: Path):
     assert doc_gen_root.is_dir()
 
 
-def main(
-    doc_gen_root: Path, system_prompts: List[str], out: str = ".ailly_prompts"
-) -> None:
+def main(doc_gen_root: Path, system_prompts: List[str], out: Path) -> None:
     """Generate prompts and configuration files for Ailly."""
-    out_path = Path(out)
-    setup_ailly(system_prompts, out_path)
+    system_prompts = read_system_prompts(system_prompts)
+    setup_ailly(system_prompts, out)
     validate_root_path(doc_gen_root)
     doc_gen = make_doc_gen(doc_gen_root)
-    write_prompts(doc_gen, out_path)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Write Ailly prompts for DocGen snippets and parse the results."
-    )
-    parser.add_argument(
-        "--doc-gen-root", required=True, help="Path to a DocGen ready project."
-    )
-    parser.add_argument(
-        "--system-prompts",
-        nargs="+",
-        required=True,
-        help="List of prompt strings or file paths to store in a .aillyrc file.",
-    )
-    parser.add_argument(
-        "--out",
-        default=".ailly_prompts",
-        help="Directory where Ailly prompt files will be written. Defaults to '.ailly_prompts'.",
-    )
-
-    args = parser.parse_args()
-
-    doc_gen_root = Path(args.doc_gen_root)
-    system_prompts = parse_prompts_arg(args.system_prompts)
-    main(doc_gen_root, system_prompts, out=args.out)
+    write_prompts(doc_gen, out)

aws_doc_sdk_examples_tools/agent/parse_json_files.py

@@ -1,4 +1,3 @@
-import argparse
 import json
 import logging
 from os.path import basename
@@ -33,17 +32,16 @@ def extract_json_from_text(text: str) -> Optional[Dict]:
     return None
 
 
-def process_files(file_paths: List[str]) -> Dict[str, Dict]:
+def process_files(file_paths: List[Path]) -> Dict[str, Dict]:
     results = {}
     for path in file_paths:
         try:
-            with open(path, "r", encoding="utf-8") as f:
-                content = f.read()
+            content = path.read_text()
             json_data = extract_json_from_text(content)
             if json_data is not None:
                 results[basename(path).replace(".md.ailly.md", "")] = json_data
             else:
-                logger.warning(f"No valid JSON object found in file: {f.name}")
+                logger.warning(f"No valid JSON object found in file: {path.name}")
         except Exception as e:
             logger.warning(f"Error processing file {path}: {e}")
     return results
@@ -53,25 +51,6 @@ def write_objects(object: Dict, out: Path):
     out.write_text(json.dumps(object, indent=2))
 
 
-def main():
-    parser = argparse.ArgumentParser(
-        description="Extract JSON objects from one or more text files."
-    )
-
-    parser.add_argument("files", nargs="+", help="List of files to process")
-
-    parser.add_argument(
-        "--out",
-        default="out.json",
-        help="File path where the resultant json will be written.",
-    )
-
-    args = parser.parse_args()
-
-    json_objects = process_files(args.files)
-    out = Path(args.out)
+def main(file_paths: List[Path], out: Path):
+    json_objects = process_files(file_paths)
     write_objects(json_objects, out)
-
-
-if __name__ == "__main__":
-    main()

aws_doc_sdk_examples_tools/agent/update_doc_gen.py

@@ -1,4 +1,3 @@
-import argparse
 import json
 import logging
 from pathlib import Path
@@ -37,9 +36,9 @@ def update_examples(doc_gen: DocGen, examples: Iterable[Example]) -> None:
             logger.warning(f"Could not find example with id: {example.id}")
 
 
-def main(doc_gen_root: Path, updates_path: Path) -> None:
+def main(doc_gen_root: Path, iam_updates_path: Path) -> None:
     doc_gen = DocGen.from_root(doc_gen_root)
-    examples = examples_from_updates(updates_path)
+    examples = examples_from_updates(iam_updates_path)
     update_examples(doc_gen, examples)
     print(
         [
@@ -48,26 +47,6 @@ def main(doc_gen_root: Path, updates_path: Path) -> None:
                 "title_abbrev": ex.title_abbrev,
                 "synopsis": ex.synopsis,
             }
-            for id, ex in doc_gen.examples.items()
+            for _, ex in doc_gen.examples.items()
         ]
     )
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Apply a list of example updates to a doc gen instance"
-    )
-    parser.add_argument(
-        "--doc-gen-root", required=True, help="Path to a DocGen ready project."
-    )
-    parser.add_argument(
-        "--updates-path",
-        default="example_updates.json",
-        help="JSON file containing a list of example updates.",
-    )
-
-    args = parser.parse_args()
-
-    doc_gen_root = Path(args.doc_gen_root)
-    updates_path = Path(args.updates_path)
-    main(doc_gen_root, updates_path)

requirements.txt

@@ -8,3 +8,4 @@ PyYAML==6.0.1
 requests==2.32.0
 types-PyYAML==6.0.12.12
 yamale==4.0.4
+typer==0.15.3