Teams in Legion

from legion.agents import agent
from legion.groups.decorators import team
from legion.interface.decorators import tool
from typing import Annotated, List
from pydantic import Field

# Define specialized agents
@agent(
    model="openai:gpt-4o-mini",
    temperature=0.3
)
class Researcher:
    """You are a research specialist who excels at finding and organizing information.
    
    Your role is to gather relevant facts and data on a given topic.
    You should be thorough and objective in your research.
    """
    
    @tool
    def search_information(
        self,
        query: Annotated[str, Field(description="Search query")]
    ) -> List[str]:
        """Search for information on a given query."""
        # In a real implementation, this would use a search API
        if "climate change" in query.lower():
            return [
                "Global temperatures have risen by 1.1°C since pre-industrial times.",
                "Sea levels are rising at a rate of 3.7mm per year.",
                "Arctic sea ice is declining at a rate of 13.1% per decade.",
                "CO2 levels in the atmosphere are at their highest in 800,000 years."
            ]
        else:
            return ["No specific information found for this query."]

@agent(
    model="openai:gpt-4o-mini",
    temperature=0.7
)
class Writer:
    """You are a skilled writer who excels at creating engaging and informative content.
    
    Your role is to transform raw information into well-structured, compelling narratives.
    You should adapt your writing style to the target audience and purpose.
    """
    
    @tool
    def draft_content(
        self,
        topic: Annotated[str, Field(description="Content topic")],
        facts: Annotated[List[str], Field(description="Facts to include")],
        style: Annotated[str, Field(description="Writing style (formal, casual, technical)")] = "formal"
    ) -> str:
        """Draft content based on provided facts and style."""
        # In a real implementation, this would call an LLM
        intro = f"Here's a draft on {topic} in a {style} style:\n\n"
        
        if style == "formal":
            content = f"# {topic.title()}\n\n"
            content += "## Introduction\n\n"
            content += f"This document provides an overview of key facts regarding {topic}.\n\n"
            content += "## Key Findings\n\n"
            for fact in facts:
                content += f"- {fact}\n"
        else:
            content = f"# Let's Talk About {topic.title()}!\n\n"
            content += f"Hey there! Today we're diving into {topic} - here's what you need to know:\n\n"
            for fact in facts:
                content += f"* {fact}\n"
                
        return intro + content

# Create a team of agents
@team
class ContentCreationTeam:
    """A team that collaborates to research and write content on various topics."""
    
    # Define team members
    members = {
        "researcher": Researcher(),
        "writer": Writer()
    }
    
    # Define team workflow
    async def process_topic(self, topic: str, style: str = "formal"):
        # Step 1: Research the topic
        research_results = await self.members["researcher"].aprocess(
            f"Find information about {topic}"
        )
        facts = research_results.tool_results["search_information"]
        
        # Step 2: Write content based on research
        writing_results = await self.members["writer"].aprocess(
            f"Write content about {topic} using these facts: {facts}"
        )
        content = writing_results.tool_results["draft_content"]
        
        return content

# Use the team
content_team = ContentCreationTeam()
article = await content_team.process_topic("climate change", style="casual")
print(article)

@agent(
    model="openai:gpt-4o",
    temperature=0.5
)
class ProjectManager:
    """You are a project manager who coordinates the work of a team.
    
    Your role is to break down complex tasks, assign them to appropriate team members,
    and ensure the overall quality and coherence of the final product.
    """
    
    @tool
    def create_project_plan(
        self,
        project: Annotated[str, Field(description="Project description")],
        team_members: Annotated[List[str], Field(description="Available team members")]
    ) -> dict:
        """Create a project plan with tasks assigned to team members."""
        # Implementation details...
        return {
            "tasks": [
                {"id": 1, "description": "Research phase", "assigned_to": "researcher"},
                {"id": 2, "description": "Content creation", "assigned_to": "writer"},
                {"id": 3, "description": "Quality review", "assigned_to": "editor"}
            ]
        }

@team
class HierarchicalTeam:
    """A team with a hierarchical structure where a project manager coordinates other agents."""
    
    # Define team members
    members = {
        "manager": ProjectManager(),
        "researcher": Researcher(),
        "writer": Writer(),
        "editor": Editor()
    }
    
    async def execute_project(self, project_description: str):
        # Step 1: Manager creates a project plan
        plan_results = await self.members["manager"].aprocess(
            f"Create a plan for this project: {project_description}"
        )
        plan = plan_results.tool_results["create_project_plan"]
        
        # Step 2: Execute tasks according to the plan
        results = {}
        for task in plan["tasks"]:
            agent_name = task["assigned_to"]
            task_description = task["description"]
            
            # Assign task to the appropriate agent
            agent_result = await self.members[agent_name].aprocess(
                f"Complete this task: {task_description}. Project context: {project_description}"
            )
            
            results[task["id"]] = agent_result
            
        # Step 3: Manager reviews and finalizes the project
        final_result = await self.members["manager"].aprocess(
            f"Review and finalize the project results: {results}"
        )
        
        return final_result

@team
class CollaborativeSolvingTeam:
    """A team that collaboratively solves complex problems through iterative refinement."""
    
    # Define team members with different expertise
    members = {
        "creative_thinker": CreativeThinker(),
        "critical_analyzer": CriticalAnalyzer(),
        "domain_expert": DomainExpert()
    }
    
    async def solve_problem(self, problem_statement: str, max_iterations: int = 3):
        # Initial solution proposal
        solution = await self.members["creative_thinker"].aprocess(
            f"Propose an initial solution to this problem: {problem_statement}"
        )
        current_solution = solution.content
        
        # Iterative refinement
        for i in range(max_iterations):
            # Critical analysis of current solution
            critique = await self.members["critical_analyzer"].aprocess(
                f"Critically analyze this solution to the problem '{problem_statement}': {current_solution}"
            )
            
            # Domain expert review
            expert_review = await self.members["domain_expert"].aprocess(
                f"Review this solution from a domain perspective. Problem: {problem_statement}, Solution: {current_solution}, Critique: {critique.content}"
            )
            
            # Refinement based on feedback
            refined_solution = await self.members["creative_thinker"].aprocess(
                f"Refine your solution based on this feedback. Original problem: {problem_statement}, Current solution: {current_solution}, Critique: {critique.content}, Expert review: {expert_review.content}"
            )
            
            current_solution = refined_solution.content
            
        return current_solution

@agent(
    model="openai:gpt-4o",
    temperature=0.7
)
class ProponentAgent:
    """You argue in favor of a position, highlighting its strengths and benefits."""
    
    @tool
    def argue_for(self, topic: str) -> str:
        """Present arguments in favor of a position."""
        # Implementation details...
        return "Arguments in favor of the position..."

@agent(
    model="openai:gpt-4o",
    temperature=0.7
)
class OpponentAgent:
    """You argue against a position, highlighting its weaknesses and drawbacks."""
    
    @tool
    def argue_against(self, topic: str) -> str:
        """Present arguments against a position."""
        # Implementation details...
        return "Arguments against the position..."

@agent(
    model="openai:gpt-4o",
    temperature=0.3
)
class MediatorAgent:
    """You synthesize different perspectives to reach a balanced conclusion."""
    
    @tool
    def synthesize(self, topic: str, arguments_for: str, arguments_against: str) -> str:
        """Synthesize different perspectives into a balanced conclusion."""
        # Implementation details...
        return "Balanced synthesis of the arguments..."

@team
class DebateTeam:
    """A team that debates topics from different perspectives to reach balanced conclusions."""
    
    # Define team members
    members = {
        "proponent": ProponentAgent(),
        "opponent": OpponentAgent(),
        "mediator": MediatorAgent()
    }
    
    async def debate_topic(self, topic: str, rounds: int = 2):
        arguments_for = ""
        arguments_against = ""
        
        # Multiple rounds of debate
        for i in range(rounds):
            # Proponent presents arguments
            pro_result = await self.members["proponent"].aprocess(
                f"Present arguments in favor of {topic}. Previous arguments against: {arguments_against}"
            )
            arguments_for = pro_result.tool_results["argue_for"]
            
            # Opponent presents counter-arguments
            con_result = await self.members["opponent"].aprocess(
                f"Present arguments against {topic}. Previous arguments for: {arguments_for}"
            )
            arguments_against = con_result.tool_results["argue_against"]
        
        # Mediator synthesizes the debate
        synthesis_result = await self.members["mediator"].aprocess(
            f"Synthesize these perspectives on {topic}. Arguments for: {arguments_for}. Arguments against: {arguments_against}."
        )
        
        return synthesis_result.tool_results["synthesize"]