Skip to content

Commit 193837d

Browse files
authored
Enhance LLIAM tooling with AWS Documentation MCP server integration (#173)
* Enhance LLIAM tooling with AWS Documentation MCP server integration - Add AWS Documentation MCP server to LLIAM configuration - Add option to skip LLM generation step with --skip-generation flag - Remove title prefixes from generated metadata - Pin ailly CLI version to 1.7.0-rc1 - Add uv dependency to requirements.txt - Add .ailly_iam_policy to gitignore - Remove unused example policy file - Code formatting and comment improvements * Update README for ailly agent.
1 parent bc84e6e commit 193837d

File tree

7 files changed

+38
-37
lines changed

7 files changed

+38
-37
lines changed

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,4 @@ __pycache__
66
.idea/
77
build/
88
dist/
9+
.ailly_iam_policy

aws_doc_sdk_examples_tools/agent/README.md

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,11 @@ python -m aws_doc_sdk_examples_tools.agent.bin.main \
3636

3737
### 🔧 Arguments
3838

39-
Run `python -m aws_doc_sdk_examples_tools.agent.bin.main --help` for more info.
39+
- `iam_tributary_root`: Path to the root directory of your IAM policy tributary
40+
- `--system-prompts`: List of system prompt files or strings to include in the Ailly configuration
41+
- `--skip-generation`: Skip the prompt generation and Ailly execution steps (useful for reprocessing existing outputs)
42+
43+
Run `python -m aws_doc_sdk_examples_tools.agent.bin.main update --help` for more info.
4044

4145
---
4246

aws_doc_sdk_examples_tools/agent/bin/main.py

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,18 +17,24 @@
1717

1818

1919
@app.command()
20-
def update(iam_tributary_root: str, system_prompts: List[str] = []) -> None:
20+
def update(
21+
iam_tributary_root: str,
22+
system_prompts: List[str] = [],
23+
skip_generation: bool = False,
24+
) -> None:
2125
"""
2226
Generate new IAM policy metadata for a tributary.
2327
"""
2428
doc_gen_root = Path(iam_tributary_root)
25-
make_prompts(
26-
doc_gen_root=doc_gen_root,
27-
system_prompts=system_prompts,
28-
out_dir=AILLY_DIR_PATH,
29-
language="IAMPolicyGrammar",
30-
)
31-
run(["npx", "@ailly/cli", "--root", AILLY_DIR])
29+
30+
if not skip_generation:
31+
make_prompts(
32+
doc_gen_root=doc_gen_root,
33+
system_prompts=system_prompts,
34+
out_dir=AILLY_DIR_PATH,
35+
language="IAMPolicyGrammar",
36+
)
37+
run(["npx @ailly/[email protected]", "--root", AILLY_DIR])
3238

3339
process_ailly_files(
3440
input_dir=str(AILLY_DIR_PATH), output_file=str(IAM_UPDATES_PATH)

aws_doc_sdk_examples_tools/agent/make_prompts.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import os
55
from pathlib import Path
66
from typing import List
7+
import yaml
78

89
from aws_doc_sdk_examples_tools.doc_gen import DocGen, Snippet
910

@@ -26,8 +27,9 @@ def write_prompts(doc_gen: DocGen, out_dir: Path, language: str) -> None:
2627
examples = doc_gen.examples
2728
snippets = doc_gen.snippets
2829
for example_id, example in examples.items():
29-
# "Title" and "Abbrev" are the defaults. If they're not there, it suggests we've already
30-
# added new titles.
30+
# TCXContentAnalyzer prefixes new metadata title/title_abbrev entries with
31+
# the DEFAULT_METADATA_PREFIX. Checking this here to make sure we're only
32+
# running the LLM tool on new extractions.
3133
title = example.title or ""
3234
title_abbrev = example.title_abbrev or ""
3335
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:
4850
def setup_ailly(system_prompts: List[str], out_dir: Path) -> None:
4951
"""Create the .aillyrc configuration file."""
5052
fence = "---"
51-
options = {"isolated": "true"}
52-
options_block = "\n".join(f"{key}: {value}" for key, value in options.items())
53+
options = {
54+
"isolated": "true",
55+
"mcp": {
56+
"awslabs.aws-documentation-mcp-server": {
57+
"type": "stdio",
58+
"command": "uvx",
59+
"args": ["awslabs.aws-documentation-mcp-server@latest"],
60+
}
61+
},
62+
}
63+
options_block = yaml.dump(options).strip()
5364
prompts_block = "\n".join(system_prompts)
5465

5566
content = f"{fence}\n{options_block}\n{fence}\n{prompts_block}"

aws_doc_sdk_examples_tools/agent/policy.json.example.md

Lines changed: 0 additions & 20 deletions
This file was deleted.

aws_doc_sdk_examples_tools/agent/process_ailly_files.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,7 @@
1515
logger = logging.getLogger(__name__)
1616

1717
EXPECTED_KEYS: Set[str] = set(["title", "title_abbrev"])
18-
VALUE_PREFIXES: Dict[str, str] = {
19-
"title": "Example policy: ",
20-
"title_abbrev": "Example: ",
21-
}
18+
VALUE_PREFIXES: Dict[str, str] = {"title": "", "title_abbrev": "", "synopsis": ""}
2219

2320

2421
class MissingExpectedKeys(Exception):
@@ -84,6 +81,7 @@ def parse_ailly_file(
8481
if key in result:
8582
result[key] = f"{prefix}{result[key]}"
8683

84+
result["title_abbrev"] = result["title"]
8785
result["id"] = Path(file_path).name.split(".md.ailly.md")[0]
8886
result["_source_file"] = file_path
8987

requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,4 @@ requests==2.32.0
99
types-PyYAML==6.0.12.12
1010
yamale==4.0.4
1111
typer==0.15.3
12+
uv==0.7.10

0 commit comments

Comments
 (0)