supervisor_framework_crewai_gemini_marktechpost.py

# -*- coding: utf-8 -*-
"""supervisor_framework_crewai_gemini_Marktechpost.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1VzZ9EwQ1BZWKsLGtUIIWW6wx5tu_VgpC
"""

!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
    )

    topic = "The Future of Artificial Intelligence in Healthcare"
    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()