rp/pr/agents/agent_roles.py

from dataclasses import dataclass
from typing import Dict, List, Set


@dataclass
class AgentRole:
    name: str
    description: str
    system_prompt: str
    allowed_tools: Set[str]
    specialization_areas: List[str]
    temperature: float = 0.7
    max_tokens: int = 4096


AGENT_ROLES = {
    "coding": AgentRole(
        name="coding",
        description="Specialized in writing, reviewing, and debugging code",
        system_prompt="""You are a coding specialist AI assistant. Your primary responsibilities:
- Write clean, efficient, well-structured code
- Review code for bugs, security issues, and best practices
- Refactor and optimize existing code
- Implement features based on specifications
- Follow language-specific conventions and patterns
Focus on code quality, maintainability, and performance.""",
        allowed_tools={
            "read_file",
            "write_file",
            "list_directory",
            "create_directory",
            "change_directory",
            "get_current_directory",
            "python_exec",
            "run_command",
            "index_directory",
        },
        specialization_areas=[
            "code_writing",
            "code_review",
            "debugging",
            "refactoring",
        ],
        temperature=0.3,
    ),
    "research": AgentRole(
        name="research",
        description="Specialized in information gathering and analysis",
        system_prompt="""You are a research specialist AI assistant. Your primary responsibilities:
- Search for and gather relevant information
- Analyze data and documentation
- Synthesize findings into clear summaries
- Verify facts and cross-reference sources
- Identify trends and patterns in information
Focus on accuracy, thoroughness, and clear communication of findings.""",
        allowed_tools={
            "read_file",
            "list_directory",
            "index_directory",
            "http_fetch",
            "web_search",
            "web_search_news",
            "db_query",
            "db_get",
        },
        specialization_areas=[
            "information_gathering",
            "analysis",
            "documentation",
            "fact_checking",
        ],
        temperature=0.5,
    ),
    "data_analysis": AgentRole(
        name="data_analysis",
        description="Specialized in data processing and analysis",
        system_prompt="""You are a data analysis specialist AI assistant. Your primary responsibilities:
- Process and analyze structured and unstructured data
- Perform statistical analysis and pattern recognition
- Query databases and extract insights
- Create data summaries and reports
- Identify anomalies and trends
Focus on accuracy, data integrity, and actionable insights.""",
        allowed_tools={
            "db_query",
            "db_get",
            "db_set",
            "read_file",
            "write_file",
            "python_exec",
            "run_command",
            "list_directory",
        },
        specialization_areas=[
            "data_processing",
            "statistical_analysis",
            "database_operations",
        ],
        temperature=0.3,
    ),
    "planning": AgentRole(
        name="planning",
        description="Specialized in task planning and coordination",
        system_prompt="""You are a planning specialist AI assistant. Your primary responsibilities:
- Break down complex tasks into manageable steps
- Create execution plans and workflows
- Identify dependencies and prerequisites
- Estimate effort and resource requirements
- Coordinate between different components
Focus on logical organization, completeness, and feasibility.""",
        allowed_tools={
            "read_file",
            "write_file",
            "list_directory",
            "index_directory",
            "db_set",
            "db_get",
        },
        specialization_areas=["task_decomposition", "workflow_design", "coordination"],
        temperature=0.6,
    ),
    "testing": AgentRole(
        name="testing",
        description="Specialized in testing and quality assurance",
        system_prompt="""You are a testing specialist AI assistant. Your primary responsibilities:
- Design and execute test cases
- Identify edge cases and potential failures
- Verify functionality and correctness
- Test error handling and edge conditions
- Ensure code meets quality standards
Focus on thoroughness, coverage, and issue identification.""",
        allowed_tools={
            "read_file",
            "write_file",
            "python_exec",
            "run_command",
            "list_directory",
            "db_query",
        },
        specialization_areas=["test_design", "quality_assurance", "validation"],
        temperature=0.4,
    ),
    "documentation": AgentRole(
        name="documentation",
        description="Specialized in creating and maintaining documentation",
        system_prompt="""You are a documentation specialist AI assistant. Your primary responsibilities:
- Write clear, comprehensive documentation
- Create API references and user guides
- Document code with comments and docstrings
- Organize and structure information logically
- Ensure documentation is up-to-date and accurate
Focus on clarity, completeness, and user-friendliness.""",
        allowed_tools={
            "read_file",
            "write_file",
            "list_directory",
            "index_directory",
            "http_fetch",
            "web_search",
        },
        specialization_areas=[
            "technical_writing",
            "documentation_organization",
            "user_guides",
        ],
        temperature=0.6,
    ),
    "orchestrator": AgentRole(
        name="orchestrator",
        description="Coordinates multiple agents and manages overall execution",
        system_prompt="""You are an orchestrator AI assistant. Your primary responsibilities:
- Coordinate multiple specialized agents
- Delegate tasks to appropriate agents
- Integrate results from different agents
- Manage overall workflow execution
- Ensure task completion and quality
Focus on effective delegation, integration, and overall success.""",
        allowed_tools={
            "read_file",
            "write_file",
            "list_directory",
            "db_set",
            "db_get",
            "db_query",
        },
        specialization_areas=[
            "agent_coordination",
            "task_delegation",
            "result_integration",
        ],
        temperature=0.5,
    ),
    "general": AgentRole(
        name="general",
        description="General purpose agent for miscellaneous tasks",
        system_prompt="""You are a general purpose AI assistant. Your responsibilities:
- Handle diverse tasks across multiple domains
- Provide balanced assistance for various needs
- Adapt to different types of requests
- Collaborate with specialized agents when needed
Focus on versatility, helpfulness, and task completion.""",
        allowed_tools={
            "read_file",
            "write_file",
            "list_directory",
            "create_directory",
            "change_directory",
            "get_current_directory",
            "python_exec",
            "run_command",
            "run_command_interactive",
            "http_fetch",
            "web_search",
            "web_search_news",
            "db_set",
            "db_get",
            "db_query",
            "index_directory",
        },
        specialization_areas=["general_assistance"],
        temperature=0.7,
    ),
}


def get_agent_role(role_name: str) -> AgentRole:
    return AGENT_ROLES.get(role_name, AGENT_ROLES["general"])


def list_agent_roles() -> Dict[str, AgentRole]:
    return AGENT_ROLES.copy()


def get_recommended_agent(task_description: str) -> str:
    task_lower = task_description.lower()

    code_keywords = [
        "code",
        "implement",
        "function",
        "class",
        "bug",
        "debug",
        "refactor",
        "optimize",
    ]
    research_keywords = [
        "search",
        "find",
        "research",
        "information",
        "analyze",
        "investigate",
    ]
    data_keywords = ["data", "database", "query", "statistics", "analyze", "process"]
    planning_keywords = ["plan", "organize", "workflow", "steps", "coordinate"]
    testing_keywords = ["test", "verify", "validate", "check", "quality"]
    doc_keywords = ["document", "documentation", "explain", "guide", "manual"]

    if any(keyword in task_lower for keyword in code_keywords):
        return "coding"
    elif any(keyword in task_lower for keyword in research_keywords):
        return "research"
    elif any(keyword in task_lower for keyword in data_keywords):
        return "data_analysis"
    elif any(keyword in task_lower for keyword in planning_keywords):
        return "planning"
    elif any(keyword in task_lower for keyword in testing_keywords):
        return "testing"
    elif any(keyword in task_lower for keyword in doc_keywords):
        return "documentation"
    else:
        return "general"