!pip install crewai crewai-tools langchain-google-genai python-dotenv

import os
from typing import List, Dict, Any
from dataclasses import dataclass
from enum import Enum


from crewai import Agent, Task, Crew, Process
from crewai_tools import SerperDevTool, WebsiteSearchTool
from langchain_google_genai import ChatGoogleGenerativeAI
from dotenv import load_dotenv


class TaskPriority(Enum):
   LOW = 1
   MEDIUM = 2
   HIGH = 3
   CRITICAL = 4

@dataclass
class TaskConfig:
   description: str
   expected_output: str
   priority: TaskPriority
   max_execution_time: int = 300 
   requires_human_input: bool = False


class SupervisorFramework:
   """
   Advanced Supervisor Agent Framework using CrewAI
   Manages multiple specialized agents with hierarchical coordination
   """
  
   def __init__(self, gemini_api_key: str, serper_api_key: str = None):
       """
       Initialize the supervisor framework
      
       Args:
           gemini_api_key: Google Gemini API key (free tier)
           serper_api_key: Serper API key for web search (optional, has free tier)
       """
       os.environ["GOOGLE_API_KEY"] = gemini_api_key
       if serper_api_key:
           os.environ["SERPER_API_KEY"] = serper_api_key
      
       self.llm = ChatGoogleGenerativeAI(
           model="gemini-1.5-flash",
           temperature=0.7,
           max_tokens=2048
       )
      
       self.tools = []
       if serper_api_key:
           self.tools.append(SerperDevTool())
      
       self.agents = {}
       self.supervisor = None
       self.crew = None
      
       print(" SupervisorFramework initialized successfully!")
       print(f" LLM Model: {self.llm.model}")
       print(f"  Available tools: {len(self.tools)}")


   def create_research_agent(self) -> Agent:
       """Create a specialized research agent"""
       return Agent(
           role="Senior Research Analyst",
           goal="Conduct comprehensive research and gather accurate information on any given topic",
           backstory="""You are an expert research analyst with years of experience in
           information gathering, fact-checking, and synthesizing complex data from multiple sources.
           You excel at finding reliable sources and presenting well-structured research findings.""",
           verbose=True,
           allow_delegation=False,
           llm=self.llm,
           tools=self.tools,
           max_iter=3,
           memory=True
       )


   def create_analyst_agent(self) -> Agent:
       """Create a specialized data analyst agent"""
       return Agent(
           role="Strategic Data Analyst",
           goal="Analyze data, identify patterns, and provide actionable insights",
           backstory="""You are a strategic data analyst with expertise in statistical analysis,
           pattern recognition, and business intelligence. You transform raw data and research
           into meaningful insights that drive decision-making.""",
           verbose=True,
           allow_delegation=False,
           llm=self.llm,
           max_iter=3,
           memory=True
       )


   def create_writer_agent(self) -> Agent:
       """Create a specialized content writer agent"""
       return Agent(
           role="Expert Technical Writer",
           goal="Create clear, engaging, and well-structured written content",
           backstory="""You are an expert technical writer with a talent for making complex
           information accessible and engaging. You specialize in creating documentation,
           reports, and content that effectively communicates insights to diverse audiences.""",
           verbose=True,
           allow_delegation=False,
           llm=self.llm,
           max_iter=3,
           memory=True
       )


   def create_reviewer_agent(self) -> Agent:
       """Create a quality assurance reviewer agent"""
       return Agent(
           role="Quality Assurance Reviewer",
           goal="Review, validate, and improve the quality of all deliverables",
           backstory="""You are a meticulous quality assurance expert with an eye for detail
           and a commitment to excellence. You ensure all work meets high standards of accuracy,
           completeness, and clarity before final delivery.""",
           verbose=True,
           allow_delegation=False,
           llm=self.llm,
           max_iter=2,
           memory=True
       )


   def create_supervisor_agent(self) -> Agent:
       """Create the main supervisor agent"""
       return Agent(
           role="Project Supervisor & Coordinator",
           goal="Coordinate team efforts, manage workflows, and ensure project success",
           backstory="""You are an experienced project supervisor with expertise in team
           coordination, workflow optimization, and quality management. You ensure that all
           team members work efficiently towards common goals and maintain high standards
           throughout the project lifecycle.""",
           verbose=True,
           allow_delegation=True,
           llm=self.llm,
           max_iter=2,
           memory=True
       )


   def setup_agents(self):
       """Initialize all agents in the framework"""
       print(" Setting up specialized agents...")
      
       self.agents = {
           'researcher': self.create_research_agent(),
           'analyst': self.create_analyst_agent(),
           'writer': self.create_writer_agent(),
           'reviewer': self.create_reviewer_agent()
       }
      
       self.supervisor = self.create_supervisor_agent()
      
       print(f" Created {len(self.agents)} specialized agents + 1 supervisor")
      
       for role, agent in self.agents.items():
           print(f"   └── {role.title()}: {agent.role}")


   def create_task_workflow(self, topic: str, task_configs: Dict[str, TaskConfig]) -> List[Task]:
       """
       Create a comprehensive task workflow
      
       Args:
           topic: Main topic/project focus
           task_configs: Dictionary of task configurations
          
       Returns:
           List of CrewAI Task objects
       """
       tasks = []
      
       if 'research' in task_configs:
           config = task_configs['research']
           research_task = Task(
               description=f"{config.description} Focus on: {topic}",
               expected_output=config.expected_output,
               agent=self.agents['researcher']
           )
           tasks.append(research_task)
      
       if 'analysis' in task_configs:
           config = task_configs['analysis']
           analysis_task = Task(
               description=f"{config.description} Analyze the research findings about: {topic}",
               expected_output=config.expected_output,
               agent=self.agents['analyst'],
               context=tasks 
           )
           tasks.append(analysis_task)
      
       if 'writing' in task_configs:
           config = task_configs['writing']
           writing_task = Task(
               description=f"{config.description} Create content about: {topic}",
               expected_output=config.expected_output,
               agent=self.agents['writer'],
               context=tasks
           )
           tasks.append(writing_task)
      
       if 'review' in task_configs:
           config = task_configs['review']
           review_task = Task(
               description=f"{config.description} Review all work related to: {topic}",
               expected_output=config.expected_output,
               agent=self.agents['reviewer'],
               context=tasks
           )
           tasks.append(review_task)
      
       supervisor_task = Task(
           description=f"""As the project supervisor, coordinate the entire workflow for: {topic}.
           Monitor progress, ensure quality standards, resolve any conflicts between agents,
           and provide final project oversight. Ensure all deliverables meet requirements.""",
           expected_output="""A comprehensive project summary including:
           - Executive summary of all completed work
           - Quality assessment of deliverables
           - Recommendations for improvements or next steps
           - Final project status report""",
           agent=self.supervisor,
           context=tasks
       )
       tasks.append(supervisor_task)
      
       return tasks


   def execute_project(self,
                      topic: str,
                      task_configs: Dict[str, TaskConfig],
                      process_type: Process = Process.hierarchical) -> Dict[str, Any]:
       """
       Execute a complete project using the supervisor framework
      
       Args:
           topic: Main project topic
           task_configs: Task configurations
           process_type: CrewAI process type (hierarchical recommended for supervisor)
          
       Returns:
           Dictionary containing execution results
       """
       print(f" Starting project execution: {topic}")
       print(f" Process type: {process_type.value}")
      
       if not self.agents or not self.supervisor:
           self.setup_agents()
      
       tasks = self.create_task_workflow(topic, task_configs)
       print(f" Created {len(tasks)} tasks in workflow")
      
       crew_agents = list(self.agents.values()) + [self.supervisor]
      
       self.crew = Crew(
           agents=crew_agents,
           tasks=tasks,
           process=process_type,
           manager_llm=self.llm, 
           verbose=True,
           memory=True
       )
      
       print(" Executing project...")
       try:
           result = self.crew.kickoff()
          
           return {
               'status': 'success',
               'result': result,
               'topic': topic,
               'tasks_completed': len(tasks),
               'agents_involved': len(crew_agents)
           }
          
       except Exception as e:
           print(f" Error during execution: {str(e)}")
           return {
               'status': 'error',
               'error': str(e),
               'topic': topic
           }


   def get_crew_usage_metrics(self) -> Dict[str, Any]:
       """Get usage metrics from the crew"""
       if not self.crew:
           return {'error': 'No crew execution found'}
      
       try:
           return {
               'total_tokens_used': getattr(self.crew, 'total_tokens_used', 'Not available'),
               'total_cost': getattr(self.crew, 'total_cost', 'Not available'),
               'execution_time': getattr(self.crew, 'execution_time', 'Not available')
           }
       except:
           return {'note': 'Metrics not available for this execution'}

def create_sample_task_configs() -> Dict[str, TaskConfig]:
   """Create sample task configurations for demonstration"""
   return {
       'research': TaskConfig(
           description="Conduct comprehensive research on the given topic. Gather information from reliable sources and compile key findings.",
           expected_output="A detailed research report with key findings, statistics, trends, and source references (minimum 500 words).",
           priority=TaskPriority.HIGH
       ),
       'analysis': TaskConfig(
           description="Analyze the research findings to identify patterns, insights, and implications.",
           expected_output="An analytical report highlighting key insights, trends, opportunities, and potential challenges (minimum 400 words).",
           priority=TaskPriority.HIGH
       ),
       'writing': TaskConfig(
           description="Create a comprehensive, well-structured document based on the research and analysis.",
           expected_output="A professional document with clear structure, engaging content, and actionable recommendations (minimum 800 words).",
           priority=TaskPriority.MEDIUM
       ),
       'review': TaskConfig(
           description="Review all deliverables for quality, accuracy, completeness, and coherence.",
           expected_output="A quality assessment report with recommendations for improvements and final approval status.",
           priority=TaskPriority.CRITICAL
       )
   }

def demo_supervisor_framework():
   """
   Demo function to showcase the supervisor framework
   Replace 'your_gemini_api_key' with your actual API key
   """
   print(" CrewAI Supervisor Framework Demo")
   print("=" * 50)
  
   framework = SupervisorFramework(
       gemini_api_key="Use Your API Key Here", 
       serper_api_key=None 
   )
  
   task_configs = create_sample_task_configs()
  
   print(f" Demo Topic: {topic}")
   print(f" Task Configurations: {list(task_configs.keys())}")
  
   results = framework.execute_project(topic, task_configs)
  
   print("\n" + "=" * 50)
   print(" EXECUTION RESULTS")
   print("=" * 50)
  
   if results['status'] == 'success':
       print(f" Status: {results['status'].upper()}")
       print(f" Tasks Completed: {results['tasks_completed']}")
       print(f" Agents Involved: {results['agents_involved']}")
       print(f" Final Result Preview: {str(results['result'])[:200]}...")
   else:
       print(f" Status: {results['status'].upper()}")
       print(f" Error: {results['error']}")
  
   metrics = framework.get_crew_usage_metrics()
   print(f"\n Usage Metrics: {metrics}")


if __name__ == "__main__":
   demo_supervisor_framework()