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"
Reference in New Issue View Git Blame Copy Permalink