agent/parsing.py

import json
from typing import List, Dict


def to_pretty_json(data) -> str:
    return json.dumps(data, indent=2, sort_keys=True)


# Known presets from the Track 1 MCP server:
PRESET_KEYWORDS = [
    {
        "change_type": "vlan",
        "preset_id": "leaf_tor_vlan_stage",
        "keywords": ["stage", "staging", "test"],
        "default": False,
    },
    {
        "change_type": "vlan",
        "preset_id": "leaf_tor_vlan_commit",
        "keywords": ["commit", "production", "prod", "rollout"],
        "default": False,
    },
    {
        "change_type": "interface",
        "preset_id": "tor_uplink_enable",
        "keywords": ["enable uplink", "bring up uplink", "turn on uplink"],
        "default": False,
    },
    {
        "change_type": "interface",
        "preset_id": "tor_uplink_shutdown",
        "keywords": ["shutdown uplink", "shut down uplink", "disable uplink"],
        "default": False,
    },
    {
        "change_type": "bgp_neighbor",
        "preset_id": "leaf_bgp_fabric_neighbor_add",
        "keywords": ["add bgp", "new neighbor", "add neighbor"],
        "default": False,
    },
    {
        "change_type": "bgp_neighbor",
        "preset_id": "leaf_bgp_fabric_neighbor_remove",
        "keywords": ["remove bgp", "delete neighbor", "remove neighbor"],
        "default": False,
    },
]


def _guess_preset(change_text: str) -> Dict[str, str]:
    text = change_text.lower()

    # First pass: explicit keywords.
    for preset in PRESET_KEYWORDS:
        for kw in preset["keywords"]:
            if kw in text:
                return {
                    "change_type": preset["change_type"],
                    "preset_id": preset["preset_id"],
                    "mode": "lightning",
                }

    # Second pass: vague cues (VLAN vs interface vs bgp)
    if "vlan" in text:
        return {
            "change_type": "vlan",
            "preset_id": "leaf_tor_vlan_stage",
            "mode": "lightning",
        }
    if "uplink" in text or "interface" in text:
        return {
            "change_type": "interface",
            "preset_id": "tor_uplink_enable",
            "mode": "lightning",
        }
    if "bgp" in text or "neighbor" in text:
        return {
            "change_type": "bgp_neighbor",
            "preset_id": "leaf_bgp_fabric_neighbor_add",
            "mode": "lightning",
        }

    # Fallback: safest-ish default.
    return {
        "change_type": "vlan",
        "preset_id": "leaf_tor_vlan_stage",
        "mode": "lightning",
    }


def _guess_device(change_text: str) -> str:
    """
    Ultra-naive device extraction: look for leaf-XX / tor-XX / core-XX tokens,
    otherwise default to leaf-01.
    """
    text = change_text.lower()
    for prefix in ["leaf-", "tor-", "core-"]:
        idx = text.find(prefix)
        if idx != -1:
            # grab token like leaf-01
            token = text[idx:].split()[0].strip(",.;:")
            return token
    return "leaf-01"


def parse_change_request(change_text: str) -> List[Dict]:
    """
    For now, treat the entire request as one step.

    Later you can split into multiple steps by sentence/semicolon/etc.
    """
    base = _guess_preset(change_text)
    device = _guess_device(change_text)

    step_json = {
        "device": device,
        "action": f"{base['change_type']}_change",
        "change_type": base["change_type"],
        "preset_id": base["preset_id"],
        "mode": base["mode"],
        "raw_text": change_text.strip(),
    }

    step = {
        "description": change_text.strip(),
        "json": step_json,
    }

    return [step]


def build_diff_summaries(steps: List[Dict]) -> List[Dict]:
    """
    Produces simple 'diff-ish' summaries per step.
    """
    summaries = []
    for step in steps:
        j = step["json"]
        device = j.get("device", "device")
        change_type = j.get("change_type", "change")
        preset_id = j.get("preset_id", "preset")

        title = f"{change_type} change on {device} ({preset_id})"
        body = (
            f"# BEFORE (conceptual)\n"
            f"# - configuration in stable state for {change_type} preset '{preset_id}'\n\n"
            f"# AFTER (conceptual)\n"
            f"# - configuration reflecting applied preset '{preset_id}' on {device}\n"
        )

        summaries.append({"title": title, "body": body})
    return summaries


def build_rollback_plan(steps: List[Dict]) -> List[Dict]:
    """
    Invert each step in reverse order.
    """
    rollback_steps = []
    for step in reversed(steps):
        j = step["json"]
        device = j.get("device", "device")
        change_type = j.get("change_type", "change")
        preset_id = j.get("preset_id", "preset")

        description = f"Rollback {change_type} change on {device} (preset {preset_id})"

        commands = [
            f"# rollback placeholder for {change_type} change on {device}",
            f"# original preset: {preset_id}",
        ]

        rollback_steps.append(
            {
                "description": description,
                "commands": commands,
            }
        )

    return rollback_steps