mcp-agent-selection/src/mcp_agent_selection/simple_server.py

"""
MCP Agent Selection Service - Simple server for agent recommendations with roots support
"""

import asyncio
import json
import os
import yaml
from pathlib import Path
from typing import List, Dict, Optional, Set
from dataclasses import dataclass
from fastmcp import FastMCP

@dataclass
class AgentInfo:
    name: str
    emoji: str
    description: str
    tools: List[str]
    content: str
    file_path: str

@dataclass
class ProjectRoots:
    directories: List[str]
    base_path: str
    description: str = ""

class SimpleAgentLibrary:
    def __init__(self, templates_path: Path):
        self.templates_path = templates_path
        self.agents: Dict[str, AgentInfo] = {}
        self.roots: Optional[ProjectRoots] = None
    
    async def load_agents(self):
        """Load all agent templates from the directory"""
        if not self.templates_path.exists():
            print(f"⚠️  Templates path not found: {self.templates_path}")
            return
        
        for file_path in self.templates_path.glob("*.md"):
            try:
                content = file_path.read_text()
                
                # Extract YAML frontmatter
                if content.startswith("---\n"):
                    parts = content.split("---\n", 2)
                    if len(parts) >= 3:
                        frontmatter = yaml.safe_load(parts[1])
                        body = parts[2]
                        
                        agent = AgentInfo(
                            name=frontmatter.get("name", file_path.stem),
                            emoji=frontmatter.get("emoji", "🤖"),
                            description=frontmatter.get("description", ""),
                            tools=frontmatter.get("tools", []),
                            content=body,
                            file_path=str(file_path)
                        )
                        
                        self.agents[agent.name] = agent
                        
            except Exception as e:
                print(f"❌ Error loading {file_path}: {e}")
        
        print(f"✅ Loaded {len(self.agents)} agents")
    
    def set_roots(self, directories: List[str], base_path: str, description: str = ""):
        """Set project roots for focused analysis"""
        self.roots = ProjectRoots(directories, base_path, description)
    
    def get_roots(self) -> Optional[ProjectRoots]:
        """Get current project roots"""
        return self.roots
    
    def clear_roots(self):
        """Clear project roots"""
        self.roots = None
    
    def recommend_agents(self, task: str, project_context: str = "") -> List[Dict]:
        """Simple recommendation logic"""
        recommendations = []
        task_lower = task.lower()
        
        # Keywords to agent mapping
        keywords = {
            "python": ["🔮-python-mcp-expert", "🧪-testing-integration-expert"],
            "fastapi": ["🚄-fastapi-expert"],
            "docker": ["🐳-docker-infrastructure-expert"],
            "security": ["🔒-security-audit-expert"],
            "documentation": ["📖-readme-expert"],
            "subagent": ["🎭-subagent-expert"],
            "test": ["🧪-testing-integration-expert"],
            "mcp": ["🔮-python-mcp-expert"]
        }
        
        # Find matching agents
        for keyword, agent_names in keywords.items():
            if keyword in task_lower:
                for agent_name in agent_names:
                    if agent_name in self.agents:
                        agent = self.agents[agent_name]
                        recommendations.append({
                            "name": agent.name,
                            "emoji": agent.emoji,
                            "description": agent.description,
                            "confidence": 0.8 if keyword in task_lower else 0.5,
                            "reason": f"Matches keyword '{keyword}' in task description"
                        })
        
        # If no specific matches, suggest subagent expert
        if not recommendations and "🎭-subagent-expert" in self.agents:
            agent = self.agents["🎭-subagent-expert"]
            recommendations.append({
                "name": agent.name,
                "emoji": agent.emoji,
                "description": agent.description,
                "confidence": 0.6,
                "reason": "General purpose recommendation for task planning"
            })
        
        return recommendations[:5]  # Limit to top 5

# Initialize
templates_path = Path(os.getenv("AGENT_TEMPLATES_PATH", "/home/rpm/claude/claude-config/agent_templates"))
agent_library = SimpleAgentLibrary(templates_path)

# Create FastMCP server
mcp = FastMCP("MCP Agent Selection Service")

@mcp.tool()
async def set_project_roots(directories: List[str], base_path: str, description: str = "") -> str:
    """Set project roots for focused analysis"""
    agent_library.set_roots(directories, base_path, description)
    return f"✅ Set project roots: {directories} in {base_path}"

@mcp.tool()
async def get_current_roots() -> Dict:
    """Get current project roots configuration"""
    roots = agent_library.get_roots()
    if roots:
        return {
            "directories": roots.directories,
            "base_path": roots.base_path,
            "description": roots.description
        }
    return {"message": "No roots configured"}

@mcp.tool()
async def clear_project_roots() -> str:
    """Clear project roots configuration"""
    agent_library.clear_roots()
    return "✅ Cleared project roots"

@mcp.tool()
async def recommend_agents(task: str, project_context: str = "") -> List[Dict]:
    """Get agent recommendations for a task"""
    recommendations = agent_library.recommend_agents(task, project_context)
    return recommendations

@mcp.tool()
async def get_agent_content(agent_name: str) -> str:
    """Get full content of a specific agent"""
    if agent_name in agent_library.agents:
        return agent_library.agents[agent_name].content
    return f"❌ Agent '{agent_name}' not found"

@mcp.tool()
async def list_agents() -> List[Dict]:
    """List all available agents"""
    return [
        {
            "name": agent.name,
            "emoji": agent.emoji,
            "description": agent.description,
            "tools": agent.tools
        }
        for agent in agent_library.agents.values()
    ]

@mcp.tool()
async def server_stats() -> Dict:
    """Get server statistics"""
    return {
        "total_agents": len(agent_library.agents),
        "roots_configured": agent_library.roots is not None,
        "templates_path": str(agent_library.templates_path)
    }

async def main():
    """Main entry point"""
    print("🚀 Starting Claude Agent MCP Server...")
    await agent_library.load_agents()
    
    print("🔧 Available MCP tools:")
    print("  - set_project_roots")
    print("  - get_current_roots") 
    print("  - clear_project_roots")
    print("  - recommend_agents")
    print("  - get_agent_content")
    print("  - list_agents")
    print("  - server_stats")
    
    await mcp.run(transport="stdio")

def main_sync():
    """Synchronous entry point for script execution"""
    asyncio.run(main())

if __name__ == "__main__":
    main_sync()