Update crew.py

crew.py

@@ -6,6 +6,8 @@
 import os
 import pandas as pd
 from crewai import Agent, Crew, Process, Task
+from crewai.agents.agent_builder.base_agent import BaseAgent
+from crewai.project import CrewBase, agent, crew, task
 from crewai.tools import tool
 #from crewai_tools import StagehandTool
 from google import genai
@@ -13,6 +15,7 @@ from google.genai import types
 from openinference.instrumentation.crewai import CrewAIInstrumentor
 from phoenix.otel import register
 from tools import add, subtract, multiply, divide, modulus
+from typing import List
 from utils import read_file_json, read_docx_text, read_pptx_text, is_ext
 
 ## LLMs
@@ -46,9 +49,11 @@ tracer_provider = register(
 
 #CrewAIInstrumentor().instrument(tracer_provider=tracer_provider)
 
-def run_crew(question, file_path):
-    # Tools
- 
+@CrewBase
+class GAIACrew():
+	agents: List[BaseAgent]
+	tasks: List[Task]
+
     @tool("Web Search Tool")
     def web_search_tool(question: str) -> str:
         """Given a question only, search the web to answer the question.
@@ -311,175 +316,122 @@ def run_crew(question, file_path):
         except Exception as e:
             raise RuntimeError(f"Processing failed: {str(e)}")
                         
-    # Agents
-    
-    web_search_agent = Agent(
-        role="Web Search Agent",
-        goal="Given a question only, search the web and answer the question: {question}",
-        backstory="As an expert web search assistant, you search the web to answer the question.",
-        allow_delegation=False,
-        llm=AGENT_MODEL,
-        max_iter=2,
-        tools=[web_search_tool],
-        verbose=True
-    )
-
-    #with StagehandTool(
-    #    api_key=os.environ["BROWSERBASE_API_KEY"],
-    #    project_id=os.environ["BROWSERBASE_PROJECT_ID"],
-    #    model_api_key=os.environ["OPENAI_API_KEY"],
-    #    model_name="gpt-4.1"
-    #) as web_browser_tool:
-    #    web_browser_agent = Agent(
-    #        role="Web Browser Agent",
-    #        goal="Given a question and URL, load the URL and act, extract, navigate, observe, and answer the question: {question}",
-    #        backstory="As an expert browser assistant, you load the URL and act, extract, navigate, observe to answer the question.",
-    #        allow_delegation=False,
-    #        llm=AGENT_MODEL,
-    #        max_iter=2,
-    #        tools=[web_browser_tool],
-    #        verbose=True
-    #    )
-    
-    image_analysis_agent = Agent(
-        role="Image Analysis Agent",
-        goal="Given a question and image file, analyze the image and answer the question: {question}",
-        backstory="As an expert image analysis assistant, you analyze the image to answer the question.",
-        allow_delegation=False,
-        llm=AGENT_MODEL,
-        max_iter=2,
-        tools=[image_analysis_tool],
-        verbose=True
-    )
-
-    audio_analysis_agent = Agent(
-        role="Audio Analysis Agent",
-        goal="Given a question and audio file, analyze the audio and answer the question: {question}",
-        backstory="As an expert audio analysis assistant, you analyze the audio to answer the question.",
-        allow_delegation=False,
-        llm=AGENT_MODEL,
-        max_iter=2,
-        tools=[audio_analysis_tool],
-        verbose=True
-    )
-
-    video_analysis_agent = Agent(
-        role="Video Analysis Agent",
-        goal="Given a question and video file, analyze the video and answer the question: {question}",
-        backstory="As an expert video analysis assistant, you analyze the video file to answer the question.",
-        allow_delegation=False,
-        llm=AGENT_MODEL,
-        max_iter=2,
-        tools=[video_analysis_tool],
-        verbose=True
-    )
+	@agent
+	def web_search_agent(self) -> Agent:
+		return Agent(
+            allow_delegation=False,
+            llm=AGENT_MODEL,
+            max_iter=2,
+            tools=[web_search_tool],
+            verbose=True
+		)
     
-    youtube_analysis_agent = Agent(
-        role="YouTube Analysis Agent",
-        goal="Given a question and YouTube URL, analyze the video and answer the question: {question}",
-        backstory="As an expert YouTube analysis assistant, you analyze the video to answer the question.",
-        allow_delegation=False,
-        llm=AGENT_MODEL,
-        max_iter=2,
-        tools=[youtube_analysis_tool],
-        verbose=True
-    )
-
-    document_analysis_agent = Agent(
-        role="Document Analysis Agent",
-        goal="Given a question and document file, analyze the document and answer the question: {question}",
-        backstory="As an expert document analysis assistant, you analyze the document to answer the question.",
-        allow_delegation=False,
-        llm=AGENT_MODEL,
-        max_iter=2,
-        tools=[document_analysis_tool],
-        verbose=True
-    )
-
+	@agent
+	def image_analysis_agent(self) -> Agent:
+		return Agent(
+            allow_delegation=False,
+            llm=AGENT_MODEL,
+            max_iter=2,
+            tools=[image_analysis_tool],
+            verbose=True
+        )
+
+	@agent
+	def audio_analysis_agent(self) -> Agent:
+		return Agent(
+            allow_delegation=False,
+            llm=AGENT_MODEL,
+            max_iter=2,
+            tools=[audio_analysis_tool],
+            verbose=True
+        )
+
+	@agent
+	def video_analysis_agent(self) -> Agent:
+		return Agent(
+            allow_delegation=False,
+            llm=AGENT_MODEL,
+            max_iter=2,
+            tools=[video_analysis_tool],
+            verbose=True
+        )
+
+	@agent
+	def youtube_analysis_agent(self) -> Agent:
+		return Agent(
+            allow_delegation=False,
+            llm=AGENT_MODEL,
+            max_iter=2,
+            tools=[youtube_analysis_tool],
+            verbose=True
+        )
+
+	@agent
+	def document_analysis_agent(self) -> Agent:
+		return Agent(
+            allow_delegation=False,
+            llm=AGENT_MODEL,
+            max_iter=2,
+            tools=[document_analysis_tool],
+            verbose=True
+        )
+
+	@agent
+	def arithmetic_agent(self) -> Agent:
+		return Agent(
     arithmetic_agent = Agent(
-        role="Arithmetic Agent",
-        goal="Given a question and two numbers, perform the calculation and answer the question: {question}",
-        backstory="As an expert arithmetic assistant, you perform the calculation to answer the question.",
         allow_delegation=False,
         llm=AGENT_MODEL,
         max_iter=2,
         tools=[arithmetic_tool],
         verbose=True
     )
-        
-    code_execution_agent = Agent(
-        role="Code Execution Agent",
-        goal="Given a question and Python file, execute the file to answer the question: {question}",
-        backstory="As an expert Python code execution assistant, you execute code to answer the question.",
-        allow_delegation=False,
-        llm=AGENT_MODEL,
-        max_iter=3,
-        tools=[code_execution_tool],
-        verbose=True
-    )
-
-    code_generation_agent = Agent(
-        role="Code Generation Agent",
-        goal="Given a question and JSON data, generate and execute code to answer the question: {question}",
-        backstory="As an expert Python code generation assistant, you generate and execute code to answer the question.",
-        allow_delegation=False,
-        llm=AGENT_MODEL,
-        max_iter=3,
-        tools=[code_generation_tool],
-        verbose=True
-    )
 
+	@agent
+	def code_execution_agent(self) -> Agent:
+		return Agent(
+            allow_delegation=False,
+            llm=AGENT_MODEL,
+            max_iter=3,
+            tools=[code_execution_tool],
+            verbose=True
+        )
+
+	@agent
+	def code_generation_agent(self) -> Agent:
+		return Agent(
+            allow_delegation=False,
+            llm=AGENT_MODEL,
+            max_iter=3,
+            tools=[code_generation_tool],
+            verbose=True
+        )
+
+	@agent
+	def image_analysis_agent(self) -> Agent:
+		return Agent(
     manager_agent = Agent(
-        role="Manager Agent",
-        goal="Answer the following question. If needed, delegate to one of your coworkers. Question: {question}",
-        backstory="As an expert manager assistant, you answer the question.",
         allow_delegation=True,
         llm=MANAGER_MODEL,
         max_iter=5,
         verbose=True
     )
 
-    # Task
-
-    manager_task = Task(
-        agent=manager_agent,
-        description="Answer the following question. If needed, delegate to one of your coworkers:\n"
-                    "- Web Search Agent requires a question only.\n"
-                    "- Web Browser Agent requires a question and **URL**.\n"
-                    "- Image Analysis Agent requires a question and **.png, .jpeg, .webp, .heic, or .heif image file**.\n"
-                    "- Audio Analysis Agent requires a question and **.wav, .mp3, .aiff, .aac, .ogg, or .flac audio file**.\n"
-                    "- Video Analysis Agent requires a question and **.mp4, .mpeg, .mov, .avi, .x-flv, .mpg, .webm, .wmv, or .3gpp video file**.\n"
-                    "- YouTube Analysis Agent requires a question and **YouTube URL**.\n"
-                    "- Document Analysis Agent requires a question and **.docx, .pptx, .pdf, .txt, .html, css, .js, .md, .xml, or .rtf document file**.\n"
-                    "- Arithmetic Agent requires a question and **two numbers to add, subtract, multiply, divide, or get the modulus**. "
-                    "  In case there are more than two numbers, use the Code Generation Agent instead.\n"
-                    "- Code Execution Agent requires a question and **Python file**.\n"
-                    "- Code Generation Agent requires a question and **JSON data**.\n"
-                    "In case you cannot answer the question and there is not a good coworker, delegate to the Code Generation Agent.\n"
-                    "Question: {question}",
-        expected_output="The answer to the question."
-    )
-    
-    # Crew
-    
-    crew = Crew(
-        agents=[web_search_agent, 
-                #web_browser_agent, 
-                image_analysis_agent, 
-                audio_analysis_agent, 
-                video_analysis_agent, 
-                youtube_analysis_agent, 
-                document_analysis_agent, 
-                arithmetic_agent, 
-                code_execution_agent, 
-                code_generation_agent],
-        manager_agent=manager_agent,
-        tasks=[manager_task],
-        verbose=True
-    )
+	@task
+	def manager_task(self) -> Task:
+		return Task(
+			config=self.tasks_config['manager_task']
+		)
 
-    # Process
+	@crew
+	def crew(self) -> Crew:
+        return Crew(
+            agents=self.agents,
+            tasks=self.tasks,
+            verbose=True
+        )
 
+def run_crew(question, file_path):
     final_question = question
     
     if file_path:
@@ -489,7 +441,7 @@ def run_crew(question, file_path):
         else:
             final_question = f"{question} File path: {file_path}."
     
-    answer = crew.kickoff(inputs={"question": final_question})
+    answer = GAIACrew().crew().kickoff(inputs={"question": final_question})
     final_answer = get_final_answer(FINAL_ANSWER_MODEL, question, str(answer))
 
     print(f"=> Initial question: {question}")