class Agent:
    def __init__(self, name: str, role: str, system_prompt: str, tools: list):
        self.name = name
        self.role = role
        self.system_prompt = system_prompt
        self.tools = tools

SWARM_ROSTER = {
    "DATA_SCIENTIST": Agent(
        name="DATA_SCIENTIST",
        role="Code Specialist",
        system_prompt="You are a Data Scientist. Write Python code and propose it via execute_python.",
        tools=["execute_python", "reflect_database_schema", "query_database"]
    ),
    "SECURITY_AUDITOR": Agent(
        name="SECURITY_AUDITOR",
        role="File & Log Analysis",
        system_prompt="You are a Security Auditor. Use read_file to analyze local documents.",
        tools=["read_file"]
    ),
    "RESEARCHER": Agent(
        name="RESEARCHER",
        role="General & Memory Intelligence",
        system_prompt="You are an Enterprise Researcher. Use ingest_knowledge_base to memorize documents, and search_knowledge_base to recall them.",
        tools=["ingest_knowledge_base", "search_knowledge_base", "search_web", "scrape_url"]
    )
}

def get_agent(agent_name: str) -> Agent:
    return SWARM_ROSTER.get(agent_name, SWARM_ROSTER["RESEARCHER"])
