training/prepare_data.py

"""
Prepare training data from AWS DevOps Agent documentation markdown files.

Reads markdown files from training/data/raw/ and generates instruction-response
pairs in JSONL format for fine-tuning.

Generates multiple Q&A pair types:
1. Full-page summary pairs
2. Section-level Q&A pairs
3. Targeted question variants per section
"""

import json
import re
from pathlib import Path

RAW_DIR = Path("training/data/raw")
OUTPUT_FILE = Path("training/data/train.jsonl")

SYSTEM_MSG = """/no_think
You are Tech Advisor, an expert on AWS cloud services with deep knowledge of AWS DevOps Agent.

You have comprehensive knowledge of AWS DevOps Agent including:
- What it is and how it works (Agent Spaces, topology, dual-console architecture)
- Key features: autonomous incident response, proactive prevention, on-demand SRE tasks
- Integrations: CloudWatch, Datadog, Dynatrace, New Relic, Splunk, Grafana, PagerDuty, GitHub, GitLab, Azure DevOps, ServiceNow, Slack
- GA features: Azure/on-prem support, Triage Agent, Learned/Custom Skills, Code Indexing, Private Connections
- Pricing: $0.0083 per agent-second, free trial details, AWS Support credits
- Getting started: Agent Spaces, connecting tools, running investigations
- Security: encryption, customer managed keys, IdP integration, CloudTrail auditing
- Available regions: US East, US West, Frankfurt, Ireland, Sydney, Tokyo

Be concise and structured. Use bullet points where appropriate. Provide accurate, detailed answers."""


def extract_sections(markdown: str) -> list[dict]:
    """Split a markdown document into sections by headers."""
    sections = []
    current_title = "Overview"
    current_content = []
    current_level = 0

    for line in markdown.split("\n"):
        header_match = re.match(r'^(#{1,4})\s+(.+)', line)
        if header_match:
            if current_content:
                content_text = "\n".join(current_content).strip()
                if content_text:
                    sections.append({
                        "title": current_title,
                        "content": content_text,
                        "level": current_level,
                    })
            current_level = len(header_match.group(1))
            current_title = header_match.group(2).strip()
            current_content = []
        else:
            current_content.append(line)

    if current_content:
        content_text = "\n".join(current_content).strip()
        if content_text:
            sections.append({
                "title": current_title,
                "content": content_text,
                "level": current_level,
            })

    return [s for s in sections if len(s["content"]) > 30]


def clean_topic_name(filename: str) -> str:
    """Convert filename to a readable topic name."""
    name = filename.replace(".md", "")
    name = re.sub(r'^(about-aws-devops-agent-|aws-devops-agent-|getting-started-with-aws-devops-agent-|working-with-devops-agent-|configuring-capabilities-for-aws-devops-agent-|connecting-telemetry-sources-|connecting-to-ticketing-and-chat-|connecting-to-cicd-pipelines-|connecting-azure-|custom-agents-|interfacing-with-the-devops-agent-|integrating-devops-agent-into-event-driven-applications-using-amazon-eventbridge-)', '', name)
    name = name.replace("-", " ").replace("_", " ")
    return name


def generate_question_variants(title: str, topic: str) -> list[str]:
    """Generate natural question variants for a section."""
    title_lower = title.lower()
    questions = []

    if any(w in title_lower for w in ["what is", "about", "overview"]):
        questions.extend([
            f"What is {topic} in AWS DevOps Agent?",
            f"Explain {topic} in AWS DevOps Agent.",
            f"Tell me about {topic}.",
        ])
    elif any(w in title_lower for w in ["getting started", "creating", "setup", "setting up"]):
        questions.extend([
            f"How do I set up {topic} in AWS DevOps Agent?",
            f"Walk me through {title.lower()} for AWS DevOps Agent.",
            f"What are the steps to {title.lower()}?",
        ])
    elif any(w in title_lower for w in ["connecting", "integrat"]):
        questions.extend([
            f"How do I connect {topic} to AWS DevOps Agent?",
            f"What's the process for integrating {topic} with AWS DevOps Agent?",
            f"How does AWS DevOps Agent work with {topic}?",
        ])
    elif any(w in title_lower for w in ["security", "encryption", "iam", "authentication"]):
        questions.extend([
            f"How does {topic} work in AWS DevOps Agent?",
            f"What security features does AWS DevOps Agent provide for {topic}?",
            f"Tell me about {topic} for AWS DevOps Agent.",
        ])
    elif any(w in title_lower for w in ["pricing", "cost", "quota"]):
        questions.extend([
            f"What are the {topic} for AWS DevOps Agent?",
            f"How much does AWS DevOps Agent cost?",
            f"What are the limits and {topic} for AWS DevOps Agent?",
        ])
    else:
        questions.extend([
            f"What is {title} in AWS DevOps Agent?",
            f"Tell me about {title} in AWS DevOps Agent.",
            f"How does {title} work in AWS DevOps Agent?",
        ])

    return questions


def generate_pairs_from_doc(filepath: Path) -> list[dict]:
    """Generate training pairs from a single documentation file."""
    content = filepath.read_text().strip()
    if not content or len(content) < 50:
        return []

    topic = clean_topic_name(filepath.name)
    pairs = []

    # Pair 1: Full document as a comprehensive answer
    full_question = f"Give me a comprehensive overview of {topic} in AWS DevOps Agent."
    if len(content) > 200:
        pairs.append({
            "messages": [
                {"role": "system", "content": SYSTEM_MSG},
                {"role": "user", "content": full_question},
                {"role": "assistant", "content": content},
            ]
        })

    # Pair 2: Direct "what is" question with full content
    pairs.append({
        "messages": [
            {"role": "system", "content": SYSTEM_MSG},
            {"role": "user", "content": f"What is {topic}?"},
            {"role": "assistant", "content": content},
        ]
    })

    # Section-level pairs with question variants
    sections = extract_sections(content)
    for section in sections:
        if len(section["content"]) < 50:
            continue

        questions = generate_question_variants(section["title"], topic)
        for q in questions[:2]:
            pairs.append({
                "messages": [
                    {"role": "system", "content": SYSTEM_MSG},
                    {"role": "user", "content": q},
                    {"role": "assistant", "content": section["content"]},
                ]
            })

    return pairs


def main():
    if not RAW_DIR.exists():
        print(f"Create {RAW_DIR} and add AWS documentation markdown files.")
        RAW_DIR.mkdir(parents=True, exist_ok=True)
        return

    md_files = list(RAW_DIR.glob("*.md"))
    if not md_files:
        print(f"No markdown files found in {RAW_DIR}")
        return

    all_pairs = []
    for filepath in sorted(md_files):
        pairs = generate_pairs_from_doc(filepath)
        all_pairs.extend(pairs)
        print(f"  {filepath.name}: {len(pairs)} training pairs")

    OUTPUT_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(OUTPUT_FILE, "w") as f:
        for pair in all_pairs:
            f.write(json.dumps(pair) + "\n")

    print(f"\nTotal: {len(all_pairs)} training pairs written to {OUTPUT_FILE}")


if __name__ == "__main__":
    main()