kaizen-agentic/src/kaizen_agentic/registry.py

"""Agent registry and management functionality."""

import re
import yaml
from pathlib import Path
from typing import Dict, List, Optional, Set
from dataclasses import dataclass
from enum import Enum


class AgentCategory(Enum):
    """Categories of agents for organization."""

    PROJECT_MANAGEMENT = "project-management"
    DEVELOPMENT_PROCESS = "development-process"
    CODE_QUALITY = "code-quality"
    INFRASTRUCTURE = "infrastructure"
    TESTING = "testing"
    DOCUMENTATION = "documentation"
    META = "meta"


@dataclass
class AgentDefinition:
    """Represents an agent definition with metadata."""

    name: str
    description: str
    file_path: Path
    category: AgentCategory
    dependencies: Set[str]
    model: Optional[str] = None
    memory: Optional[str] = None  # "enabled" (default) | "disabled"

    @staticmethod
    def _read_frontmatter(file_path: Path) -> dict:
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
        frontmatter_match = re.match(r"^---\n(.*?)\n---\n", content, re.DOTALL)
        if not frontmatter_match:
            raise ValueError(f"No YAML frontmatter found in {file_path}")
        frontmatter = yaml.safe_load(frontmatter_match.group(1))
        if not isinstance(frontmatter, dict) or "name" not in frontmatter:
            raise ValueError(f"Invalid frontmatter in {file_path}")
        return frontmatter

    @classmethod
    def from_file(cls, file_path: Path) -> "AgentDefinition":
        """Create AgentDefinition from a markdown file."""
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()

        # Extract YAML frontmatter
        frontmatter_match = re.match(r"^---\n(.*?)\n---\n", content, re.DOTALL)
        if not frontmatter_match:
            raise ValueError(f"No YAML frontmatter found in {file_path}")

        frontmatter = yaml.safe_load(frontmatter_match.group(1))

        # Extract dependencies from frontmatter and content
        dependencies = cls._extract_dependencies(content, frontmatter)

        # Determine category from name or content
        category = cls._determine_category(frontmatter["name"], content)

        return cls(
            name=frontmatter["name"],
            description=frontmatter["description"],
            file_path=file_path,
            category=category,
            dependencies=dependencies,
            model=frontmatter.get("model"),
            memory=frontmatter.get("memory"),
        )

    @staticmethod
    def _extract_dependencies(content: str, frontmatter: dict) -> Set[str]:
        """Extract agent dependencies from frontmatter and content."""
        dependencies = set()

        # Check frontmatter for explicit dependencies
        for key in ["dependencies", "depends_on", "requires"]:
            if key in frontmatter:
                deps = frontmatter[key]
                if isinstance(deps, list):
                    dependencies.update(deps)
                elif isinstance(deps, str):
                    # Handle comma-separated string
                    dependencies.update([d.strip() for d in deps.split(",")])

        # Look for explicit dependencies in content
        dep_patterns = [
            r"depends_on:\s*\[(.*?)\]",
            r"requires:\s*\[(.*?)\]",
            r"dependencies:\s*\[(.*?)\]",
        ]

        for pattern in dep_patterns:
            matches = re.findall(pattern, content, re.IGNORECASE)
            for match in matches:
                if isinstance(match, str):
                    deps = [d.strip().strip("\"'") for d in match.split(",")]
                    dependencies.update(deps)

        # Look for specific agent references in content (more precise)
        # Only look for full agent names like "todo-keeper agent" or "uses changelog-keeper"
        agent_patterns = [
            r"uses?\s+(\w+(?:-\w+)*-(?:keeper|agent|workflow|helper|manager))",
            r"depends?\s+on\s+(\w+(?:-\w+)*-(?:keeper|agent|workflow|helper|manager))",
            r"requires?\s+(\w+(?:-\w+)*-(?:keeper|agent|workflow|helper|manager))",
        ]

        for pattern in agent_patterns:
            matches = re.findall(pattern, content.lower())
            for match in matches:
                if match not in ["optimization", "agentic", "driven", "assisted"]:
                    dependencies.add(match.replace("-", "_"))

        return dependencies

    @staticmethod
    def _determine_category(name: str, content: str) -> AgentCategory:
        """Determine agent category based on name and content."""
        name_lower = name.lower()

        # Project management agents
        project_keywords = ["todo", "changelog", "contributing", "project"]
        if any(keyword in name_lower for keyword in project_keywords):
            return AgentCategory.PROJECT_MANAGEMENT

        # Testing agents
        if any(keyword in name_lower for keyword in ["test", "tdd"]):
            return AgentCategory.TESTING

        # Code quality agents
        if any(
            keyword in name_lower for keyword in ["refactor", "optimization", "code"]
        ):
            return AgentCategory.CODE_QUALITY

        # Documentation agents
        if any(keyword in name_lower for keyword in ["documentation", "claude"]):
            return AgentCategory.DOCUMENTATION

        # Meta agents (coaching, cross-agent orchestration)
        if any(keyword in name_lower for keyword in ["coach", "meta"]):
            return AgentCategory.META

        # Infrastructure agents
        if any(
            keyword in name_lower
            for keyword in ["setup", "repository", "tooling", "sys-medic", "medic"]
        ):
            return AgentCategory.INFRASTRUCTURE

        # Development process agents
        if any(
            keyword in name_lower
            for keyword in ["workflow", "requirements", "maintenance"]
        ):
            return AgentCategory.DEVELOPMENT_PROCESS

        # Default fallback
        return AgentCategory.DEVELOPMENT_PROCESS


class AgentRegistry:
    """Registry for managing and discovering agents."""

    def __init__(self, agents_dir: Path):
        self.agents_dir = Path(agents_dir)
        self._agents: Dict[str, AgentDefinition] = {}
        self._file_index: Dict[str, Path] = {}
        self._index_agent_files()

    def _index_agent_files(self) -> None:
        """Index agent files by frontmatter name without full parse."""
        if not self.agents_dir.exists():
            return

        for agent_file in self.agents_dir.glob("agent-*.md"):
            try:
                frontmatter = AgentDefinition._read_frontmatter(agent_file)
                self._file_index[frontmatter["name"]] = agent_file
            except Exception as e:
                print(f"Warning: Failed to index agent {agent_file}: {e}")

    def get_agent_path(self, name: str) -> Optional[Path]:
        """Return the source file path for an agent (no full parse)."""
        return self._file_index.get(name)

    def get_agent(self, name: str) -> Optional[AgentDefinition]:
        """Get agent definition by name (lazy-loaded)."""
        if name in self._agents:
            return self._agents[name]
        file_path = self._file_index.get(name)
        if file_path is None:
            return None
        try:
            agent_def = AgentDefinition.from_file(file_path)
        except Exception as e:
            print(f"Warning: Failed to load agent {name}: {e}")
            return None
        self._agents[name] = agent_def
        return agent_def

    def agent_names(self) -> List[str]:
        """List indexed agent names without loading full definitions."""
        return sorted(self._file_index.keys())

    def list_agents(
        self, category: Optional[AgentCategory] = None
    ) -> List[AgentDefinition]:
        """List all agents, optionally filtered by category."""
        agents = [self.get_agent(name) for name in self.agent_names()]
        agents = [agent for agent in agents if agent is not None]
        if category:
            agents = [a for a in agents if a.category == category]
        return sorted(agents, key=lambda a: a.name)

    def get_categories(self) -> Dict[AgentCategory, List[AgentDefinition]]:
        """Get agents organized by category."""
        categories = {}
        for agent in self.list_agents():
            if agent.category not in categories:
                categories[agent.category] = []
            categories[agent.category].append(agent)

        # Sort agents within each category
        for category in categories:
            categories[category].sort(key=lambda a: a.name)

        return categories

    def resolve_dependencies(self, agent_names: List[str]) -> List[str]:
        """Resolve agent dependencies and return ordered list."""
        resolved = []
        visited = set()

        def visit(name: str):
            if name in visited:
                return
            visited.add(name)

            agent = self.get_agent(name)
            if not agent:
                print(f"Warning: Agent '{name}' not found")
                return

            # Visit dependencies first
            for dep in agent.dependencies:
                visit(dep)

            if name not in resolved:
                resolved.append(name)

        for name in agent_names:
            visit(name)

        return resolved

    def validate_agents(self) -> Dict[str, List[str]]:
        """Validate all agents and return validation errors."""
        errors = {}

        for name in self.agent_names():
            agent = self.get_agent(name)
            if agent is None:
                errors[name] = ["Failed to load agent definition"]
                continue
            agent_errors = []

            # Check for missing dependencies
            for dep in agent.dependencies:
                if dep not in self._file_index:
                    agent_errors.append(f"Missing dependency: {dep}")

            # Check file exists
            if not agent.file_path.exists():
                agent_errors.append(f"Agent file not found: {agent.file_path}")

            # Check for circular dependencies
            if self._has_circular_dependency(name):
                agent_errors.append("Circular dependency detected")

            if agent_errors:
                errors[name] = agent_errors

        return errors

    def _has_circular_dependency(
        self, agent_name: str, visited: Optional[Set[str]] = None
    ) -> bool:
        """Check if an agent has circular dependencies."""
        if visited is None:
            visited = set()

        if agent_name in visited:
            return True

        agent = self.get_agent(agent_name)
        if not agent:
            return False

        visited.add(agent_name)

        for dep in agent.dependencies:
            if self._has_circular_dependency(dep, visited.copy()):
                return True

        return False

    def get_agent_templates(self) -> Dict[str, List[str]]:
        """Get predefined agent templates for different project types."""
        return {
            "python-basic": ["setupRepository", "keepaTodofile", "keepaChangelog"],
            "python-web": [
                "setupRepository",
                "tdd-workflow",
                "code-refactoring",
                "keepaTodofile",
                "keepaChangelog",
                "keepaContributingfile",
            ],
            "python-cli": [
                "setupRepository",
                "tdd-workflow",
                "testing-efficiency",
                "claude-documentation",
                "keepaTodofile",
                "keepaChangelog",
            ],
            "python-data": [
                "setupRepository",
                "datamodel-optimization",
                "testing-efficiency",
                "requirements-engineering",
                "keepaTodofile",
                "keepaChangelog",
            ],
            "comprehensive": [agent.name for agent in self.list_agents()],
        }