from dataclasses import dataclass
from typing import List, Dict, Any, 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'
Reference in New Issue View Git Blame Copy Permalink