diff --git a/aws_doc_sdk_examples_tools/agent/bin/main.py b/aws_doc_sdk_examples_tools/agent/bin/main.py
index 0f99b2c..aec310b 100644
--- a/aws_doc_sdk_examples_tools/agent/bin/main.py
+++ b/aws_doc_sdk_examples_tools/agent/bin/main.py
@@ -6,7 +6,8 @@
 
 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
+from aws_doc_sdk_examples_tools.agent.update_doc_gen import update as update_doc_gen
+from aws_doc_sdk_examples_tools.yaml_writer import prepare_write, write_many
 
 app = typer.Typer()
 
@@ -32,7 +33,11 @@ def update(iam_tributary_root: str, system_prompts: List[str] = []) -> None:
     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)
+    doc_gen = update_doc_gen(
+        doc_gen_root=doc_gen_root, iam_updates_path=IAM_UPDATES_PATH
+    )
+    writes = prepare_write(doc_gen.examples)
+    write_many(doc_gen_root, writes)
 
 
 if __name__ == "__main__":
diff --git a/aws_doc_sdk_examples_tools/agent/update_doc_gen.py b/aws_doc_sdk_examples_tools/agent/update_doc_gen.py
index d394a56..d58ab72 100644
--- a/aws_doc_sdk_examples_tools/agent/update_doc_gen.py
+++ b/aws_doc_sdk_examples_tools/agent/update_doc_gen.py
@@ -36,17 +36,8 @@ 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, iam_updates_path: Path) -> None:
+def update(doc_gen_root: Path, iam_updates_path: Path) -> DocGen:
     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()
-        ]
-    )
+    return doc_gen