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

Author: cpyle0819

aws_doc_sdk_examples_tools/agent/README.md

@@ -1,12 +1,13 @@
 # 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 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 back into structured JSON (`parse_json_files.py`).
+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`).
 
 ---
 
@@ -33,15 +34,11 @@ python make_prompts.py \
 - `--system-prompts`: One or more system prompts, either as strings or file paths.
 - `--out`: (Optional) Output directory. Defaults to `.ailly_prompts`.
 
-This will:
-- Create Markdown files for each snippet.
-- Create a `.aillyrc` file with the provided system prompts.
-
 ---
 
 ## Step 2: Run Ailly CLI
 
-Run Ailly on the generated prompts. From inside the .ailly_prompts directory:
+Run Ailly on the generated prompts:
 
 ```bash
 npx @ailly/cli --root .ailly_prompts
@@ -53,15 +50,33 @@ This will create `{file}.ailly.md` output files in the `.ailly_prompts` director
 
 ## Step 3: Parse Ailly output
 
-After Ailly has generated response files (`*.ailly.md`), parse them into JSON objects using `parse_json_files.py`:
+Parse the `{name}.ailly.md` files into JSON using `parse_json_files.py`:
+
+```bash
+python parse_json_files.py .ailly_prompts/*.ailly.md --out example_updates.json
+```
+
+**Arguments:**
+- Positional: List of files to parse (e.g., `*.ailly.md`).
+- `--out`: (Optional) Path for the JSON output. Defaults to `out.json`.
+
+---
+
+## Step 4: Update DocGen with Ailly results
+
+Use `update_doc_gen.py` to load the parsed data back into your `DocGen` project:
 
 ```bash
-python parse_json_files.py .ailly_prompts/*.ailly.md --out output.json
+python update_doc_gen.py \
+  --doc-gen-root /path/to/your/docgen/project \
+  --updates-path example_updates.json
 ```
 
 **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`.
+- `--doc-gen-root`: Path to the root of your DocGen project.
+- `--updates-path`: JSON file generated in Step 3 (default: `example_updates.json`).
+
+This will update the `title`, `title_abbrev`, and `synopsis` fields in the corresponding DocGen examples.
 
 ---
 
@@ -78,5 +93,10 @@ python make_prompts.py \
 npx @ailly/cli --root .ailly_prompts
 
 # Step 3: Parse results
-python parse_json_files.py .ailly_prompts/*.ailly.md --out results.json
+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
 ```

aws_doc_sdk_examples_tools/agent/make_prompts.py

@@ -20,12 +20,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:
@@ -54,15 +65,20 @@ def parse_prompts_arg(values: List[str]) -> List[str]:
     return prompts
 
 
+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: str = ".ailly_prompts"
 ) -> None:
     """Generate prompts and configuration files for Ailly."""
     out_path = Path(out)
     setup_ailly(system_prompts, out_path)
-
+    validate_root_path(doc_gen_root)
     doc_gen = make_doc_gen(doc_gen_root)
-    write_prompts(doc_gen.snippets, out_path)
+    write_prompts(doc_gen, out_path)
 
 
 if __name__ == "__main__":

aws_doc_sdk_examples_tools/agent/parse_json_files.py

@@ -41,7 +41,7 @@ def process_files(file_paths: List[str]) -> Dict[str, Dict]:
                 content = f.read()
             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}")
         except Exception as e:

aws_doc_sdk_examples_tools/agent/update_doc_gen.py

@@ -0,0 +1,73 @@
+import argparse
+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, updates_path: Path) -> None:
+    doc_gen = DocGen.from_root(doc_gen_root)
+    examples = examples_from_updates(updates_path)
+    update_examples(doc_gen, examples)
+    print(
+        [
+            {
+                "title": ex.title,
+                "title_abbrev": ex.title_abbrev,
+                "synopsis": ex.synopsis,
+            }
+            for id, 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)