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"


@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

    @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')
        )

    @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

        # Infrastructure agents
        if any(keyword in name_lower for keyword in ['setup', 'repository', 'tooling']):
            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._load_agents()

    def _load_agents(self):
        """Load all agents from the agents directory."""
        if not self.agents_dir.exists():
            return

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

    def get_agent(self, name: str) -> Optional[AgentDefinition]:
        """Get agent definition by name."""
        return self._agents.get(name)

    def list_agents(self, category: Optional[AgentCategory] = None) -> List[AgentDefinition]:
        """List all agents, optionally filtered by category."""
        agents = list(self._agents.values())
        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._agents.values():
            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, agent in self._agents.items():
            agent_errors = []

            # Check for missing dependencies
            for dep in agent.dependencies:
                if dep not in self._agents:
                    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()
            ]
        }