Enhance LLIAM tooling with AWS Documentation MCP server integration

.gitignore

@@ -6,3 +6,4 @@ __pycache__
 .idea/
 build/
 dist/
+.ailly_iam_policy

aws_doc_sdk_examples_tools/agent/README.md

@@ -36,7 +36,11 @@ python -m aws_doc_sdk_examples_tools.agent.bin.main \
 
 ### 🔧 Arguments
 
-Run `python -m aws_doc_sdk_examples_tools.agent.bin.main --help` for more info.
+- `iam_tributary_root`: Path to the root directory of your IAM policy tributary
+- `--system-prompts`: List of system prompt files or strings to include in the Ailly configuration
+- `--skip-generation`: Skip the prompt generation and Ailly execution steps (useful for reprocessing existing outputs)
+
+Run `python -m aws_doc_sdk_examples_tools.agent.bin.main update --help` for more info.
 
 ---
 

aws_doc_sdk_examples_tools/agent/bin/main.py

@@ -17,18 +17,24 @@
 
 
 @app.command()
-def update(iam_tributary_root: str, system_prompts: List[str] = []) -> None:
+def update(
+    iam_tributary_root: str,
+    system_prompts: List[str] = [],
+    skip_generation: bool = False,
+) -> None:
     """
     Generate new IAM policy metadata for a tributary.
     """
     doc_gen_root = Path(iam_tributary_root)
-    make_prompts(
-        doc_gen_root=doc_gen_root,
-        system_prompts=system_prompts,
-        out_dir=AILLY_DIR_PATH,
-        language="IAMPolicyGrammar",
-    )
-    run(["npx", "@ailly/cli", "--root", AILLY_DIR])
+
+    if not skip_generation:
+        make_prompts(
+            doc_gen_root=doc_gen_root,
+            system_prompts=system_prompts,
+            out_dir=AILLY_DIR_PATH,
+            language="IAMPolicyGrammar",
+        )
+        run(["npx @ailly/[email protected]", "--root", AILLY_DIR])
 
     process_ailly_files(
         input_dir=str(AILLY_DIR_PATH), output_file=str(IAM_UPDATES_PATH)

aws_doc_sdk_examples_tools/agent/make_prompts.py

@@ -4,6 +4,7 @@
 import os
 from pathlib import Path
 from typing import List
+import yaml
 
 from aws_doc_sdk_examples_tools.doc_gen import DocGen, Snippet
 
@@ -26,8 +27,9 @@ def write_prompts(doc_gen: DocGen, out_dir: Path, language: str) -> None:
     examples = doc_gen.examples
     snippets = doc_gen.snippets
     for example_id, example in examples.items():
-        # "Title" and "Abbrev" are the defaults. If they're not there, it suggests we've already
-        # added new titles.
+        # TCXContentAnalyzer prefixes new metadata title/title_abbrev entries with
+        # the DEFAULT_METADATA_PREFIX. Checking this here to make sure we're only
+        # running the LLM tool on new extractions.
         title = example.title or ""
         title_abbrev = example.title_abbrev or ""
         if title.startswith(DEFAULT_METADATA_PREFIX) and title_abbrev.startswith(
@@ -48,8 +50,17 @@ def write_prompts(doc_gen: DocGen, out_dir: Path, language: str) -> None:
 def setup_ailly(system_prompts: List[str], out_dir: Path) -> None:
     """Create the .aillyrc configuration file."""
     fence = "---"
-    options = {"isolated": "true"}
-    options_block = "\n".join(f"{key}: {value}" for key, value in options.items())
+    options = {
+        "isolated": "true",
+        "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)
 
     content = f"{fence}\n{options_block}\n{fence}\n{prompts_block}"

aws_doc_sdk_examples_tools/agent/process_ailly_files.py

@@ -15,10 +15,7 @@
 logger = logging.getLogger(__name__)
 
 EXPECTED_KEYS: Set[str] = set(["title", "title_abbrev"])
-VALUE_PREFIXES: Dict[str, str] = {
-    "title": "Example policy: ",
-    "title_abbrev": "Example: ",
-}
+VALUE_PREFIXES: Dict[str, str] = {"title": "", "title_abbrev": "", "synopsis": ""}
 
 
 class MissingExpectedKeys(Exception):
@@ -84,6 +81,7 @@ def parse_ailly_file(
             if key in result:
                 result[key] = f"{prefix}{result[key]}"
 
+        result["title_abbrev"] = result["title"]
         result["id"] = Path(file_path).name.split(".md.ailly.md")[0]
         result["_source_file"] = file_path
 

requirements.txt

@@ -9,3 +9,4 @@ requests==2.32.0
 types-PyYAML==6.0.12.12
 yamale==4.0.4
 typer==0.15.3
+uv==0.7.10