gns3-server/gns3server/agent/gns3_copilot/skills/loader.py

# SPDX-License-Identifier: GPL-3.0-or-later
#
# GNS3-Copilot - AI-powered Network Lab Assistant for GNS3
#
# This file is part of GNS3-Copilot project.
#
# GNS3-Copilot is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# GNS3-Copilot is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License
# along with GNS3-Copilot. If not, see <https://www.gnu.org/licenses/>.
#
# Copyright (C) 2025 Yue Guobin
# Author: Yue Guobin
#
# Project Home: https://github.com/yueguobin/gns3-copilot
#

"""
Skills Loader

This module provides functionality to load skills from YAML files.
Supports loading injection skills and device/feature skills.
"""

import logging
import os
from pathlib import Path
from typing import Any, Dict

try:
    import yaml
except ImportError:
    yaml = None

logger = logging.getLogger(__name__)


class SkillsLoader:
    """
    Load skills from YAML files in the skills directory.

    This loader reads YAML files and converts them to the skill dictionary
    format expected by the skills registry.
    """

    def __init__(self, skills_dir: str):
        """
        Initialize the skills loader.

        Args:
            skills_dir: Path to the skills directory containing YAML files
        """
        self.skills_dir = Path(skills_dir)
        if not self.skills_dir.exists():
            logger.warning(f"Skills directory does not exist: {self.skills_dir}")

    def load_injection_skills(self) -> Dict[str, Dict[str, Any]]:
        """
        Load all injection skills from YAML files.

        Returns:
            Dictionary mapping skill keys to skill definitions
        """
        if yaml is None:
            logger.error("PyYAML is not installed. Cannot load skills from YAML.")
            return {}

        skills = {}
        injection_dir = self.skills_dir / "injection"

        if not injection_dir.exists():
            logger.warning(f"Injection skills directory not found: {injection_dir}")
            return {}

        for yaml_file in injection_dir.glob("*.yaml"):
            try:
                skill_data = self._load_yaml(yaml_file)
                if not skill_data:
                    logger.warning(f"Skipping empty YAML file: {yaml_file}")
                    continue
                # Generate key from filename (e.g., "ospf_issues.yaml" -> "injection_ospf")
                skill_key = f"injection_{yaml_file.stem}"
                skills[skill_key] = skill_data
                logger.debug(f"Loaded injection skill: {skill_key} from {yaml_file}")
            except Exception as e:
                logger.error(f"Failed to load skill from {yaml_file}: {e}")

        logger.debug(f"Loaded {len(skills)} injection skills from {injection_dir}")
        return skills

    def load_device_skills(self) -> Dict[str, Dict[str, Any]]:
        """
        Load all device skills from YAML files.

        Supports two layouts:
        - Single file: device/<device>.yaml
        - Split directory: device/<device>/_base.yaml + device/<device>/<topic>.yaml
          (topic files are merged into the base skill under "topics",
          keyed by their "topic" field)

        Returns:
            Dictionary mapping skill keys to skill definitions
        """
        if yaml is None:
            logger.error("PyYAML is not installed. Cannot load skills from YAML.")
            return {}

        skills: Dict[str, Dict[str, Any]] = {}
        device_dir = self.skills_dir / "device"

        if not device_dir.exists():
            logger.warning(f"Device skills directory not found: {device_dir}")
            return {}

        for entry in sorted(device_dir.iterdir()):
            if entry.is_file() and entry.suffix == ".yaml":
                self._load_single_device_skill(skills, entry)
            elif entry.is_dir():
                self._load_split_device_skill(skills, entry)

        logger.debug(f"Loaded {len(skills)} device skills from device directory")
        return skills

    def _load_single_device_skill(self, skills: Dict[str, Dict[str, Any]], yaml_file: Path) -> None:
        """
        Load a single-file device skill into the skills dictionary.
        """
        try:
            skill_data = self._load_yaml(yaml_file)
            if not skill_data:
                logger.warning(f"Skipping empty YAML file: {yaml_file}")
                return
            # Use device_type from YAML content as the key
            # Fallback to filename stem if device_type not present
            skill_key = skill_data.get("device_type")
            if not skill_key:
                skill_key = yaml_file.stem
                logger.warning(f"No device_type in {yaml_file}, using filename '{skill_key}' as key")
            skills[skill_key] = skill_data
            logger.debug(f"Loaded device skill: {skill_key} from {yaml_file}")
        except Exception as e:
            logger.error(f"Failed to load skill from {yaml_file}: {e}")

    def _load_split_device_skill(self, skills: Dict[str, Dict[str, Any]], device_path: Path) -> None:
        """
        Load a split device skill (directory with _base.yaml + topic files).

        The base file provides the device-level skill; every other YAML file
        in the directory is a protocol topic merged under "topics".
        """
        base_file = device_path / "_base.yaml"
        if not base_file.exists():
            logger.error(f"No _base.yaml in device directory: {device_path}, skipping")
            return

        try:
            base_data = self._load_yaml(base_file)
        except Exception as e:
            logger.error(f"Failed to load skill from {base_file}: {e}")
            return
        if not base_data:
            logger.warning(f"Skipping empty YAML file: {base_file}")
            return

        skill_key = base_data.get("device_type")
        if not skill_key:
            skill_key = device_path.name
            logger.warning(f"No device_type in {base_file}, using directory name '{skill_key}' as key")

        # Seed topics from the base file (if any), then merge topic files
        base_topics = base_data.get("topics")
        topics: Dict[str, Any] = dict(base_topics) if isinstance(base_topics, dict) else {}

        for yaml_file in sorted(device_path.glob("*.yaml")):
            if yaml_file.name == "_base.yaml":
                continue
            try:
                topic_data = self._load_yaml(yaml_file)
                if not topic_data:
                    logger.warning(f"Skipping empty YAML file: {yaml_file}")
                    continue

                topic_device_type = topic_data.pop("device_type", None)
                if topic_device_type is not None and topic_device_type != skill_key:
                    logger.error(
                        f"device_type mismatch in {yaml_file}: '{topic_device_type}' "
                        f"!= base '{skill_key}', skipping topic"
                    )
                    continue

                topic_key = topic_data.pop("topic", None)
                if not topic_key:
                    topic_key = yaml_file.stem
                    logger.warning(f"No 'topic' field in {yaml_file}, using filename '{topic_key}'")
                if topic_key in topics:
                    logger.warning(f"Duplicate topic '{topic_key}' in {device_path.name} (from {yaml_file}), overwriting")

                # category/topics belong to the base skill only
                topic_data.pop("category", None)
                topic_data.pop("topics", None)

                topics[topic_key] = topic_data
                logger.debug(f"Loaded device topic: {skill_key}/{topic_key} from {yaml_file}")
            except Exception as e:
                logger.error(f"Failed to load topic from {yaml_file}: {e}")

        if topics:
            base_data["topics"] = topics
        skills[skill_key] = base_data
        logger.debug(f"Loaded device skill: {skill_key} from {device_path} ({len(topics)} topics)")

    def load_feature_skills(self) -> Dict[str, Dict[str, Any]]:
        """
        Load all feature skills from YAML files.

        Feature skills are network planning and design functionalities
        (e.g., topology planner), not device-specific features.

        Returns:
            Dictionary mapping skill keys to skill definitions
        """
        if yaml is None:
            logger.error("PyYAML is not installed. Cannot load skills from YAML.")
            return {}

        skills = {}
        feature_dir = self.skills_dir / "feature"

        if not feature_dir.exists():
            logger.warning(f"Feature skills directory not found: {feature_dir}")
            return {}

        for yaml_file in feature_dir.glob("*.yaml"):
            try:
                skill_data = self._load_yaml(yaml_file)
                if not skill_data:
                    logger.warning(f"Skipping empty YAML file: {yaml_file}")
                    continue
                # Use device_type from YAML content as the key
                # Fallback to filename stem if device_type not present
                skill_key = skill_data.get("device_type") if isinstance(skill_data, dict) else None
                if not skill_key:
                    skill_key = yaml_file.stem
                    logger.warning(f"No device_type in {yaml_file}, using filename '{skill_key}' as key")
                skills[skill_key] = skill_data
                logger.debug(f"Loaded feature skill: {skill_key} from {yaml_file}")
            except Exception as e:
                logger.error(f"Failed to load feature skill from {yaml_file}: {e}")

        logger.debug(f"Loaded {len(skills)} feature skills from feature directory")
        return skills

    def load_prompt(self, prompt_name: str) -> str:
        """
        Load a prompt from a markdown file.

        Args:
            prompt_name: Name of the prompt file (without .md extension)

        Returns:
            Prompt content as string, or empty string if not found
        """
        prompts_dir = self.skills_dir / "prompts"
        prompt_file = prompts_dir / f"{prompt_name}.md"

        if not prompt_file.exists():
            logger.warning(f"Prompt file not found: {prompt_file}")
            return ""

        try:
            with open(prompt_file, "r", encoding="utf-8") as f:
                content = f.read()
            logger.debug(f"Loaded prompt: {prompt_name} from {prompt_file}")
            return content
        except Exception as e:
            logger.error(f"Failed to load prompt from {prompt_file}: {e}")
            return ""

    def load_forbidden_commands(self) -> list:
        """
        Load forbidden commands from the config directory.

        Returns:
            List of forbidden command patterns, or empty list if not found
        """
        config_dir = self.skills_dir / "config"
        config_file = config_dir / "forbidden_commands.txt"

        if not config_file.exists():
            logger.warning(f"Forbidden commands file not found: {config_file}")
            return []

        try:
            commands = []
            with open(config_file, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith("#"):
                        continue
                    commands.append(line.lower())
            logger.debug(f"Loaded {len(commands)} forbidden commands from {config_file}")
            return commands
        except Exception as e:
            logger.error(f"Failed to load forbidden commands from {config_file}: {e}")
            return []

    def load_packet_analysis_protocols(self) -> Dict[str, Dict[str, Any]]:
        """
        Load packet analysis protocol definitions from YAML files.

        Returns:
            Dictionary mapping protocol keys to protocol definitions.
            Each protocol contains tshark_field, display_filter, and filter_examples.
        """
        if yaml is None:
            logger.error("PyYAML is not installed. Cannot load packet analysis protocols.")
            return {}

        protocols = {}
        packet_analysis_dir = self.skills_dir / "packet_analysis"

        if not packet_analysis_dir.exists():
            logger.debug(f"Packet analysis directory not found: {packet_analysis_dir}")
            return {}

        for yaml_file in packet_analysis_dir.glob("*.yaml"):
            try:
                protocol_data = self._load_yaml(yaml_file)
                if not protocol_data:
                    logger.warning(f"Skipping empty YAML file: {yaml_file}")
                    continue

                # Use protocol_key field from YAML as the key
                protocol_key = protocol_data.get("protocol_key")
                if not protocol_key:
                    logger.warning(f"No 'protocol_key' field in {yaml_file}, skipping")
                    continue

                protocols[protocol_key] = protocol_data
                logger.debug(f"Loaded packet analysis protocol: {protocol_key} from {yaml_file}")
            except Exception as e:
                logger.error(f"Failed to load protocol from {yaml_file}: {e}")

        logger.debug(f"Loaded {len(protocols)} packet analysis protocols from {packet_analysis_dir}")
        return protocols

    def _load_yaml(self, file_path: Path) -> Dict[str, Any]:
        """
        Load a YAML file and return its content.

        Args:
            file_path: Path to the YAML file

        Returns:
            Parsed YAML content as dictionary, or empty dict if file is empty
        """
        with open(file_path, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f)
        if not isinstance(data, dict):
            logger.warning(f"YAML file {file_path} is empty or has invalid format")
            return {}
        return data

    def validate_skill_format(self, skill_data: Dict[str, Any]) -> bool:
        """
        Validate that a skill dictionary has the required fields.

        Args:
            skill_data: Skill dictionary to validate

        Returns:
            True if valid, False otherwise
        """
        required_fields = ["name", "description", "issues"]
        for field in required_fields:
            if field not in skill_data:
                logger.error(f"Skill missing required field: {field}")
                return False

        if not isinstance(skill_data["issues"], dict):
            logger.error("Skill 'issues' field must be a dictionary")
            return False

        for issue_key, issue_data in skill_data["issues"].items():
            if not isinstance(issue_data, dict):
                logger.error(f"Issue '{issue_key}' must be a dictionary")
                return False

            issue_required = ["name", "description"]
            for field in issue_required:
                if field not in issue_data:
                    logger.error(f"Issue '{issue_key}' missing required field: {field}")
                    return False

        return True