131 lines
4.2 KiB
Python
131 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate Ansible playbooks from plain English descriptions using Ollama LLM.
|
|
|
|
Usage:
|
|
python3 scripts/ansible-generator.py "Install and configure nginx on all debian hosts"
|
|
python3 scripts/ansible-generator.py --name nginx-setup "Install nginx on debian hosts"
|
|
python3 scripts/ansible-generator.py --dry-run "Restart Docker on atlantis and calypso"
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
from lib.ollama import ollama_generate, ollama_available, OllamaUnavailableError
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
REPO_DIR = Path("/home/homelab/organized/repos/homelab")
|
|
PLAYBOOKS_DIR = REPO_DIR / "ansible" / "playbooks"
|
|
GENERATED_DIR = PLAYBOOKS_DIR / "generated"
|
|
INVENTORY_PATH = REPO_DIR / "ansible" / "inventory.yml"
|
|
|
|
# Reference playbooks for style
|
|
STYLE_REFS = [
|
|
PLAYBOOKS_DIR / "deploy_atlantis.yml",
|
|
PLAYBOOKS_DIR / "update_system.yml",
|
|
]
|
|
|
|
|
|
def load_style_examples() -> str:
|
|
"""Load first 50 lines of reference playbooks for style guidance."""
|
|
examples = []
|
|
for ref in STYLE_REFS:
|
|
if ref.exists():
|
|
lines = ref.read_text().splitlines()[:50]
|
|
examples.append(f"# {ref.name}\n" + "\n".join(lines))
|
|
return "\n\n".join(examples)
|
|
|
|
|
|
def load_inventory() -> str:
|
|
if INVENTORY_PATH.exists():
|
|
return INVENTORY_PATH.read_text()
|
|
return "(inventory not found)"
|
|
|
|
|
|
def slugify(text: str) -> str:
|
|
"""Convert description to a filename-safe slug."""
|
|
slug = re.sub(r"[^a-z0-9]+", "-", text.lower().strip())
|
|
slug = slug.strip("-")
|
|
return slug[:60] if slug else "playbook"
|
|
|
|
|
|
def generate_playbook(description: str) -> str:
|
|
examples = load_style_examples()
|
|
inventory = load_inventory()
|
|
|
|
prompt = (
|
|
f"Generate an Ansible playbook for: {description}\n\n"
|
|
f"Follow this style:\n{examples}\n\n"
|
|
f"Use this inventory:\n{inventory}\n\n"
|
|
"Output valid YAML only. Include proper hosts, tasks, handlers, and tags.\n"
|
|
"Do not wrap the YAML in code fences. Output only the raw YAML document starting with ---."
|
|
)
|
|
return ollama_generate(prompt, num_predict=3000, timeout=180)
|
|
|
|
|
|
def extract_yaml(raw: str) -> str:
|
|
"""Strip any markdown code fences the LLM may have added."""
|
|
# Remove ```yaml ... ``` wrapping
|
|
raw = re.sub(r"^```(?:ya?ml)?\s*\n", "", raw.strip())
|
|
raw = re.sub(r"\n```\s*$", "", raw.strip())
|
|
return raw.strip()
|
|
|
|
|
|
def validate_yaml(content: str) -> bool:
|
|
"""Check that content is valid YAML."""
|
|
try:
|
|
yaml.safe_load(content)
|
|
return True
|
|
except yaml.YAMLError as e:
|
|
log.error("Generated YAML is invalid: %s", e)
|
|
return False
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Generate Ansible playbooks from plain English via LLM")
|
|
parser.add_argument("description", help="Plain English description of the desired playbook")
|
|
parser.add_argument("--name", type=str, help="Output filename (without .yml extension)")
|
|
parser.add_argument("--dry-run", action="store_true", help="Print playbook without saving")
|
|
parser.add_argument("--verbose", "-v", action="store_true", help="Enable debug logging")
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(
|
|
level=logging.DEBUG if args.verbose else logging.INFO,
|
|
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
|
)
|
|
|
|
if not ollama_available():
|
|
log.error("Ollama is not reachable — aborting")
|
|
sys.exit(1)
|
|
|
|
log.info("Generating playbook for: %s", args.description)
|
|
raw = generate_playbook(args.description)
|
|
playbook_yaml = extract_yaml(raw)
|
|
|
|
if not validate_yaml(playbook_yaml):
|
|
log.error("LLM output failed YAML validation. Raw output:\n%s", raw)
|
|
sys.exit(1)
|
|
|
|
if args.dry_run:
|
|
print(playbook_yaml)
|
|
return
|
|
|
|
filename = args.name if args.name else slugify(args.description)
|
|
if not filename.endswith(".yml"):
|
|
filename += ".yml"
|
|
|
|
GENERATED_DIR.mkdir(parents=True, exist_ok=True)
|
|
out_path = GENERATED_DIR / filename
|
|
out_path.write_text(playbook_yaml + "\n")
|
|
log.info("Saved playbook to %s", out_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|