Add agent script to update a DocGen instance with json output. (#162)

Author: cpyle0819

aws_doc_sdk_examples_tools/agent/README.md

@@ -1,82 +1,74 @@
 # Ailly Prompt Workflow
 
-This project provides two scripts to **generate**, **run**, and **parse** [Ailly](https://www.npmjs.com/package/@ailly/cli) prompts based on your AWS DocGen snippets.
+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.
 
-## Overview
+---
+
+## 📦 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 back into structured JSON (`parse_json_files.py`).
+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.
+
+All of this is done with one command.
 
 ---
 
-## Prerequisites
+## ✅ Prerequisites
 
 - Python 3.8+
-- Node.js & npm (for `npx`)
+- Node.js and npm (for `npx`)
+- A DocGen project directory
 
 ---
 
-## Step 1: Generate Ailly prompts
+## 🚀 Usage
 
-Use `make_prompts.py` to create a directory of Markdown files and a `.aillyrc` configuration file.
+From your project root, run:
 
 ```bash
-python make_prompts.py \
-  --doc-gen-root /path/to/your/docgen/project \
-  --system-prompts "You are a helpful assistant..." \
-  --out .ailly_prompts
+python -m aws_doc_sdk_examples_tools.agent.bin.main \
+  /path/to/your/docgen/project \
+  --system-prompts path/to/system_prompt.txt
 ```
 
-**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`.
+### 🔧 Arguments
 
-This will:
-- Create Markdown files for each snippet.
-- Create a `.aillyrc` file with the provided system prompts.
+Run `python -m aws_doc_sdk_examples_tools.agent.bin.main --help` for more info.
 
 ---
 
-## Step 2: Run Ailly CLI
-
-Run Ailly on the generated prompts. From inside the .ailly_prompts directory:
-
-```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).
+## 🗂 What This Does
 
----
+Under the hood, this script:
 
-## Step 3: Parse Ailly output
+1. Creates a directory `.ailly_iam_policy` containing:
+   - One Markdown file per snippet.
+   - A `.aillyrc` configuration file.
 
-After Ailly has generated response files (`*.ailly.md`), parse them into JSON objects using `parse_json_files.py`:
+2. Runs `npx @ailly/cli` to generate `.ailly.md` outputs.
 
-```bash
-python parse_json_files.py .ailly_prompts/*.ailly.md --out output.json
-```
+3. Parses the Ailly `.ailly.md` files into a single `iam_updates.json` file.
 
-**Arguments:**
-- Positional: List of files to parse (you can use a glob like `*.ailly.md`).
-- `--out`: (Optional) Output file path for the resulting JSON array. Defaults to `out.json`.
+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 results.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):
+    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] = []):
+    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
 
@@ -20,12 +19,23 @@ def make_doc_gen(root: Path) -> DocGen:
     return doc_gen
 
 
-def write_prompts(snippets: Dict[str, Snippet], out: Path) -> None:
-    """Write each snippet's code into a separate Markdown file."""
+def write_prompts(doc_gen: DocGen, out: Path) -> None:
     out.mkdir(parents=True, exist_ok=True)
-    for snippet_id, snippet in snippets.items():
-        snippet_path = out / f"{snippet_id}.md"
-        snippet_path.write_text(snippet.code, encoding="utf-8")
+    examples = doc_gen.examples
+    snippets = doc_gen.snippets
+    for example_id, example in examples.items():
+        # Postfix with `.md` so Ailly will pick it up.
+        prompt_path = out / f"{example_id}.md"
+        # This assumes we're running DocGen specifically on AWSIAMPolicyExampleReservoir.
+        snippet_key = (
+            example.languages["IAMPolicyGrammar"]
+            .versions[0]
+            .excerpts[0]
+            .snippet_files[0]
+            .replace("/", ".")
+        )
+        snippet = snippets[snippet_key]
+        prompt_path.write_text(snippet.code, encoding="utf-8")
 
 
 def setup_ailly(system_prompts: List[str], out: Path) -> None:
@@ -42,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:
@@ -54,38 +64,15 @@ def parse_prompts_arg(values: List[str]) -> List[str]:
     return prompts
 
 
-def main(
-    doc_gen_root: Path, system_prompts: List[str], out: str = ".ailly_prompts"
-) -> None:
-    """Generate prompts and configuration files for Ailly."""
-    out_path = Path(out)
-    setup_ailly(system_prompts, out_path)
+def validate_root_path(doc_gen_root: Path):
+    assert "AWSIAMPolicyExampleReservoir" in str(doc_gen_root)
+    assert doc_gen_root.is_dir()
+
 
+def main(doc_gen_root: Path, system_prompts: List[str], out: Path) -> None:
+    """Generate prompts and configuration files for Ailly."""
+    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.snippets, 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)] = json_data
+                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

@@ -0,0 +1,52 @@
+import json
+import logging
+from pathlib import Path
+from typing import Iterable
+
+from aws_doc_sdk_examples_tools.doc_gen import DocGen, Example
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def examples_from_updates(updates_path: Path) -> Iterable[Example]:
+    updates = json.loads(updates_path.read_text())
+    examples = [
+        Example(
+            id=id,
+            file=None,
+            languages={},
+            title=update.get("title"),
+            title_abbrev=update.get("title_abbrev"),
+            synopsis=update.get("synopsis"),
+        )
+        for id, update in updates.items()
+    ]
+    return examples
+
+
+def update_examples(doc_gen: DocGen, examples: Iterable[Example]) -> None:
+    for example in examples:
+        if doc_gen_example := doc_gen.examples.get(example.id):
+            doc_gen_example.title = example.title
+            doc_gen_example.title_abbrev = example.title_abbrev
+            doc_gen_example.synopsis = example.synopsis
+        else:
+            logger.warning(f"Could not find example with id: {example.id}")
+
+
+def main(doc_gen_root: Path, iam_updates_path: Path) -> None:
+    doc_gen = DocGen.from_root(doc_gen_root)
+    examples = examples_from_updates(iam_updates_path)
+    update_examples(doc_gen, examples)
+    print(
+        [
+            {
+                "title": ex.title,
+                "title_abbrev": ex.title_abbrev,
+                "synopsis": ex.synopsis,
+            }
+            for _, ex in doc_gen.examples.items()
+        ]
+    )

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