Update crew.py

crew.py

@@ -10,17 +10,19 @@ from google import genai
 from openai import OpenAI
 from openinference.instrumentation.crewai import CrewAIInstrumentor
 from phoenix.otel import register
-#from pytubefix import YouTube
-from util import get_final_answer #, get_img_b64, get_imgs_b64
+from util import get_final_answer
 
 ## LLMs
 
 MANAGER_MODEL      = "gpt-4.1"
 AGENT_MODEL        = "gpt-4.1"
+
 FINAL_ANSWER_MODEL = "gpt-4.5-preview"
-AUDIO_MODEL        = "gpt-4o-transcribe"
-IMAGE_MODEL        = "gpt-4.1"
-VIDEO_MODEL        = "gpt-4.1-mini"
+
+IMAGE_MODEL        = "gemini-2.0-flash"
+AUDIO_MODEL        = "gemini-2.0-flash"
+VIDEO_MODEL        = "gemini-2.0-flash"
+YOUTUBE_MODEL      = "gemini-2.0-flash"
 
 # LLM evaluation
 
@@ -37,162 +39,115 @@ tracer_provider = register(
 CrewAIInstrumentor().instrument(tracer_provider=tracer_provider)
 
 def run_crew(question, file_path):
-    # Custom tools
+    # Tools
 
-    @tool("Audio Analysis Tool")
-    def audio_analysis_tool(question: str, file_path: str) -> str:
-        """Answer a question about an audio file.
+    @tool("Image Analysis Tool")
+    def image_analysis_tool(question: str, file_path: str) -> str:
+        """Answer a question about an image file.
     
            Args:
-               question (str): Question about the audio file
-               file_path (str): Path of the audio file
+               question (str): Question about an image file
+               file_path (str): The image file path
                 
            Returns:
-               str: Answer to the question about the audio file
+               str: Answer to the question about the image file
                 
            Raises:
-               FileNotFoundError: If the audio file does not exist
                RuntimeError: If processing fails"""
-        if not os.path.exists(file_path):
-            raise FileNotFoundError(f"Audio file not found: {file_path}")
-            
         try:
-            client = OpenAI()
+            client = genai.Client(api_key="GOOGLE_API_KEY")
 
-            transcript = client.audio.transcriptions.create(
-                file=open(file_path, "rb"),
-                model=AUDIO_MODEL,
-                prompt=question
+            file = client.files.upload(file=file_path)
+
+            response = client.models.generate_content(
+                model=IMAGE_MODEL, 
+                contents=[file, question]
             )
-            
-            return transcript.text
+          
+            return response.text
         except Exception as e:
-            raise RuntimeError(f"Failed to process audio: {str(e)}")
+            raise RuntimeError(f"Processing failed: {str(e)}")
 
-    @tool("Image Analysis Tool")
-    def image_analysis_tool(question: str, file_path: str) -> str:
-        """Answer a question about an image file.
+    @tool("Audio Analysis Tool")
+    def audio_analysis_tool(question: str, file_path: str) -> str:
+        """Answer a question about an audio file.
     
            Args:
-               question (str): Question about the image file
-               file_path (str): Path of the image file
+               question (str): Question about an audio file
+               file_path (str): The audio file path
                 
            Returns:
-               str: Answer to the question about the image file
+               str: Answer to the question about the audio file
                 
            Raises:
-               FileNotFoundError: If the image file does not exist
                RuntimeError: If processing fails"""
-        if not os.path.exists(file_path):
-            raise FileNotFoundError(f"Image file not found: {file_path}")
-            
         try:
-            # Get image
-            
-            img_b64 = get_img_b64(file_path)
+            client = genai.Client(api_key="GOOGLE_API_KEY")
 
-            # OpenAI
+            file = client.files.upload(file=file_path)
             
-            client = OpenAI()
-
-            completion = client.chat.completions.create(
-                messages=[{"role": "user",
-                           "content": [{"type": "text", "text": question},
-                                       {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}}]}],
-                model=IMAGE_MODEL
+            response = client.models.generate_content(
+                model=AUDIO_MODEL, 
+                contents=[file, question]
             )
-        
-            return completion.choices[0].message.content
+          
+            return response.text
         except Exception as e:
-            raise RuntimeError(f"Failed to process image: {str(e)}")
-    
-    @tool("YouTube Audio Analysis Tool")
-    def youtube_audio_analysis_tool(question: str, url: str) -> str:
-        """Answer an audio question about a YouTube video.
+            raise RuntimeError(f"Processing failed: {str(e)}")
+
+    @tool("Video Analysis Tool")
+    def video_analysis_tool(question: str, file_path: str) -> str:
+        """Answer a question about a video file.
     
            Args:
-               question (str): Audio question about YouTube video
-               url (str): YouTube URL
+               question (str): Question about a video file
+               file_path (str): The video file path
                 
            Returns:
-               str: Answer to the audio question about YouTube video
+               str: Answer to the question about the video file
                 
            Raises:
                RuntimeError: If processing fails"""
         try:
-            # YouTube (hack to deal with access issues)
-            
-            if url.endswith("1htKBjuUWec"):
-                file_path = "data/1htKBjuUWec.mp4"
-            else:
-                raise RuntimeError()
+            client = genai.Client(api_key="GOOGLE_API_KEY")
 
-            #file_path = "audio.mp4"
-            #yt = YouTube(url, use_oauth=True, allow_oauth_cache=True)
-            #stream = yt.streams.filter(only_audio=True).first()
-            #tream.download(filename=file_path)
-            
-            # OpenAI
+            file = client.files.upload(file=file_path)
             
-            client = OpenAI()
-
-            transcription = client.audio.transcriptions.create(
-                file=open(file_path, "rb"),
-                model=AUDIO_MODEL,
-                prompt=question
+            response = client.models.generate_content(
+                model=VIDEO_MODEL, 
+                contents=[file, question]
             )
-            
-            return transcription.text
-        except Exception as e:
-            raise RuntimeError(f"Failed to process audio: {str(e)}")
 
-    @tool("YouTube Image Analysis Tool")
-    def youtube_image_analysis_tool(question: str, url: str) -> str:
-        """Answer an image question about a YouTube video.
+            return response.text
+        except Exception as e:
+            raise RuntimeError(f"Processing failed: {str(e)}")
+            
+    @tool("YouTube Analysis Tool")
+    def youtube_analysis_tool(question: str, url: str) -> str:
+        """Answer a question about a YouTube video.
     
            Args:
-               question (str): Image question about YouTube video
-               url (str): YouTube URL
+               question (str): Question about a YouTube video
+               url (str): The YouTube video URL
                 
            Returns:
-               str: Answer to the image question about YouTube video
+               str: Answer to the question about the YouTube video
                 
            Raises:
                RuntimeError: If processing fails"""
         try:
-            # YouTube (hack to deal with access issues)
-                
-            if url.endswith("L1vXCYZAYYM"):
-                file_path = "data/L1vXCYZAYYM.mp4"
-            else:
-                raise RuntimeError()
-
-            #file_path = "video.mp4"
-            #yt = YouTube(url, use_oauth=True, allow_oauth_cache=True)
-            #stream = yt.streams.get_highest_resolution()
-            #stream.download(filename=file_path)
-
-            # Get images
-
-            imgs_b64 = get_imgs_b64(file_path)
-
-            # OpenAI
-
-            client = OpenAI()
-            
-            response = client.responses.create(
-                input=[{"role": "user",
-                        "content": [{"type": "input_text", "text": (question)},
-                                    *[{"type": "input_image", "image_url": f"data:image/jpeg;base64,{img_b64}"} for img_b64 in imgs_b64]]}],
-                model=VIDEO_MODEL
+            client = genai.Client(api_key="GOOGLE_API_KEY")
+
+            return client.models.generate_content(
+                model=YOUTUBE_MODEL,
+                contents=types.Content(
+                    parts=[types.Part(file_data=types.FileData(file_uri=url)),
+                           types.Part(text=question)]
+                )
             )
-           
-            return response.output_text
         except Exception as e:
-            raise RuntimeError(f"Failed to process video: {str(e)}")
+            raise RuntimeError(f"Processing failed: {str(e)}")
     
-    # Built-in tools
-
     web_search_tool = SerperDevTool()
     web_rag_tool = WebsiteSearchTool()
     python_coding_tool = CodeInterpreterTool()
@@ -243,10 +198,10 @@ def run_crew(question, file_path):
         verbose=False
     )
 
-    youtube_image_analysis_agent = Agent(
-        role="YouTube Image Analysis Agent",
-        goal="Analyze YouTube video to help answer image question \"{question}\"",
-        backstory="As an expert YouTube image analysis assistant, you analyze the video to help answer the question.",
+    youtube_analysis_agent = Agent(
+        role="YouTube Analysis Agent",
+        goal="Analyze YouTube video to help answer question \"{question}\"",
+        backstory="As an expert YouTube analysis assistant, you analyze the video to help answer the question.",
         allow_delegation=False,
         llm=AGENT_MODEL,
         max_iter=2,
@@ -289,10 +244,10 @@ def run_crew(question, file_path):
     
     crew = Crew(
         agents=[web_search_agent, 
-                audio_analysis_agent, 
                 image_analysis_agent, 
-                youtube_audio_analysis_agent, 
-                youtube_image_analysis_agent, 
+                audio_analysis_agent, 
+                video_analysis_agent, 
+                youtube_analysis_agent, 
                 python_coding_agent],
         manager_agent=manager_agent,
         tasks=[manager_task],