Update crew.py

crew.py

@@ -47,276 +47,276 @@ tracer_provider = register(
 
 #CrewAIInstrumentor().instrument(tracer_provider=tracer_provider)
 
-@CrewBase
-class GAIACrew():
-	agents: List[BaseAgent]
-	tasks: List[Task]
-    
-    # Tools
- 
-    @tool("Web Search Tool")
-    def web_search_tool(question: str) -> str:
-        """Given a question only, search the web to answer the question.
-    
-           Args:
-               question (str): Question to answer
-                
-           Returns:
-               str: Answer to the question
-                
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+# Tools
+
+@tool("Web Search Tool")
+def web_search_tool(question: str) -> str:
+    """Given a question only, search the web to answer the question.
+
+       Args:
+           question (str): Question to answer
+            
+       Returns:
+           str: Answer to the question
             
-            response = client.models.generate_content(
-                model=WEB_SEARCH_MODEL,
-                contents=question,
-                config=types.GenerateContentConfig(
-                    tools=[types.Tool(google_search=types.GoogleSearchRetrieval())]
-                )
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+        
+        response = client.models.generate_content(
+            model=WEB_SEARCH_MODEL,
+            contents=question,
+            config=types.GenerateContentConfig(
+                tools=[types.Tool(google_search=types.GoogleSearchRetrieval())]
             )
+        )
 
-            return response.text
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
-    
-    @tool("Image Analysis Tool")
-    def image_analysis_tool(question: str, file_path: str) -> str:
-        """Given a question and image file, analyze the image to answer the question.
-    
-           Args:
-               question (str): Question about an image file
-               file_path (str): The image file path
-                
-           Returns:
-               str: Answer to the question about the image file
-                
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-            
-            file = client.files.upload(file=file_path)
+        return response.text
+    except Exception as e:
+        raise RuntimeError(f"Processing failed: {str(e)}")
 
-            response = client.models.generate_content(
-                model=IMAGE_ANALYSIS_MODEL,
-                contents=[file, question]
-            )
-          
-            return response.text
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
-
-    @tool("Audio Analysis Tool")
-    def audio_analysis_tool(question: str, file_path: str) -> str:
-        """Given a question and audio file, analyze the audio to answer the question.
-    
-           Args:
-               question (str): Question about an audio file
-               file_path (str): The audio file path
-                
-           Returns:
-               str: Answer to the question about the audio file
-                
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+@tool("Image Analysis Tool")
+def image_analysis_tool(question: str, file_path: str) -> str:
+    """Given a question and image file, analyze the image to answer the question.
 
-            file = client.files.upload(file=file_path)
+       Args:
+           question (str): Question about an image file
+           file_path (str): The image file path
             
-            response = client.models.generate_content(
-                model=AUDIO_ANALYSIS_MODEL, 
-                contents=[file, question]
-            )
-          
-            return response.text
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
-
-    @tool("Video Analysis Tool")
-    def video_analysis_tool(question: str, file_path: str) -> str:
-        """Given a question and video file, analyze the video to answer the question.
-    
-           Args:
-               question (str): Question about a video file
-               file_path (str): The video file path
-                
-           Returns:
-               str: Answer to the question about the video file
-                
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+       Returns:
+           str: Answer to the question about the image file
+            
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+        
+        file = client.files.upload(file=file_path)
 
-            file = client.files.upload(file=file_path)
+        response = client.models.generate_content(
+            model=IMAGE_ANALYSIS_MODEL,
+            contents=[file, question]
+        )
+      
+        return response.text
+    except Exception as e:
+        raise RuntimeError(f"Processing failed: {str(e)}")
+
+@tool("Audio Analysis Tool")
+def audio_analysis_tool(question: str, file_path: str) -> str:
+    """Given a question and audio file, analyze the audio to answer the question.
+
+       Args:
+           question (str): Question about an audio file
+           file_path (str): The audio file path
             
-            response = client.models.generate_content(
-                model=VIDEO_ANALYSIS_MODEL, 
-                contents=[file, question]
-            )
+       Returns:
+           str: Answer to the question about the audio file
+            
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 
-            return response.text
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
+        file = client.files.upload(file=file_path)
+        
+        response = client.models.generate_content(
+            model=AUDIO_ANALYSIS_MODEL, 
+            contents=[file, question]
+        )
+      
+        return response.text
+    except Exception as e:
+        raise RuntimeError(f"Processing failed: {str(e)}")
+
+@tool("Video Analysis Tool")
+def video_analysis_tool(question: str, file_path: str) -> str:
+    """Given a question and video file, analyze the video to answer the question.
+
+       Args:
+           question (str): Question about a video file
+           file_path (str): The video file path
             
-    @tool("YouTube Analysis Tool")
-    def youtube_analysis_tool(question: str, url: str) -> str:
-        """Given a question and YouTube URL, analyze the video to answer the question.
-    
-           Args:
-               question (str): Question about a YouTube video
-               url (str): The YouTube URL
-                
-           Returns:
-               str: Answer to the question about the YouTube video
-                
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-
-            return client.models.generate_content(
-                model=YOUTUBE_ANALYSIS_MODEL,
-                contents=types.Content(
-                    parts=[types.Part(file_data=types.FileData(file_uri=url)),
-                           types.Part(text=question)]
-                )
-            )
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
+       Returns:
+           str: Answer to the question about the video file
+            
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 
-    @tool("Document Analysis Tool")
-    def document_analysis_tool(question: str, file_path: str) -> str:
-        """Given a question and document file, analyze the document to answer the question.
-    
-           Args:
-               question (str): Question about a document file
-               file_path (str): The document file path
-                
-           Returns:
-               str: Answer to the question about the document file
-                
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+        file = client.files.upload(file=file_path)
+        
+        response = client.models.generate_content(
+            model=VIDEO_ANALYSIS_MODEL, 
+            contents=[file, question]
+        )
 
-            contents = []
+        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:
+    """Given a question and YouTube URL, analyze the video to answer the question.
+
+       Args:
+           question (str): Question about a YouTube video
+           url (str): The YouTube URL
             
-            if is_ext(file_path, ".docx"):
-                text_data = read_docx_text(file_path)
-                contents = [f"{question}\n{text_data}"]
-                print(f"=> Text data:\n{text_data}")
-            elif is_ext(file_path, ".pptx"):
-                text_data = read_pptx_text(file_path)
-                contents = [f"{question}\n{text_data}"]
-                print(f"=> Text data:\n{text_data}")
-            else:
-                file = client.files.upload(file=file_path)
-                contents = [file, question]
+       Returns:
+           str: Answer to the question about the YouTube video
             
-            response = client.models.generate_content(
-                model=DOCUMENT_ANALYSIS_MODEL,
-                contents=contents
-            )
-          
-            return response.text
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
-
-    @tool("Arithmetic Tool")
-    def arithmetic_tool(question: str, a: float, b: float) -> float:
-        """Given a question and two numbers, perform the calculation to answer the question.
-    
-           Args:
-               question (str): Question to answer
-               a (float): First number
-               b (float): Second number
-                
-           Returns:
-               float: Result number
-                
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-            
-            response = client.models.generate_content(
-                model=ARITHMETIC_MODEL,
-                contents=question,
-                config=types.GenerateContentConfig(
-                    tools=[add, subtract, multiply, divide, modulus]
-                )
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+
+        return client.models.generate_content(
+            model=YOUTUBE_ANALYSIS_MODEL,
+            contents=types.Content(
+                parts=[types.Part(file_data=types.FileData(file_uri=url)),
+                       types.Part(text=question)]
             )
+        )
+    except Exception as e:
+        raise RuntimeError(f"Processing failed: {str(e)}")
 
-            return response.text
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
+@tool("Document Analysis Tool")
+def document_analysis_tool(question: str, file_path: str) -> str:
+    """Given a question and document file, analyze the document to answer the question.
 
-    @tool("Code Execution Tool")
-    def code_execution_tool(question: str, file_path: str) -> str:
-        """Given a question and Python file, execute the file to answer the question.
-    
-           Args:
-               question (str): Question to answer
-               file_path (str): The Python file path
-                
-           Returns:
-               str: Answer to the question
-                
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+       Args:
+           question (str): Question about a document file
+           file_path (str): The document file path
+            
+       Returns:
+           str: Answer to the question about the document file
+            
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 
+        contents = []
+        
+        if is_ext(file_path, ".docx"):
+            text_data = read_docx_text(file_path)
+            contents = [f"{question}\n{text_data}"]
+            print(f"=> Text data:\n{text_data}")
+        elif is_ext(file_path, ".pptx"):
+            text_data = read_pptx_text(file_path)
+            contents = [f"{question}\n{text_data}"]
+            print(f"=> Text data:\n{text_data}")
+        else:
             file = client.files.upload(file=file_path)
-
-            response = client.models.generate_content(
-                model=CODE_EXECUTION_MODEL,
-                contents=[file, question],
-                config=types.GenerateContentConfig(
-                    tools=[types.Tool(code_execution=types.ToolCodeExecution)]
-                ),
-            )
+            contents = [file, question]
+        
+        response = client.models.generate_content(
+            model=DOCUMENT_ANALYSIS_MODEL,
+            contents=contents
+        )
+      
+        return response.text
+    except Exception as e:
+        raise RuntimeError(f"Processing failed: {str(e)}")
+
+@tool("Arithmetic Tool")
+def arithmetic_tool(question: str, a: float, b: float) -> float:
+    """Given a question and two numbers, perform the calculation to answer the question.
+
+       Args:
+           question (str): Question to answer
+           a (float): First number
+           b (float): Second number
             
-            for part in response.candidates[0].content.parts:
-                if part.code_execution_result is not None:
-                    return part.code_execution_result.output
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
-
-    @tool("Code Generation Tool")
-    def code_generation_tool(question: str, json_data: str) -> str:
-        """Given a question and JSON data, generate and execute code to answer the question.
-
-           Args:
-               question (str): Question to answer
-                file_path (str): The JSON data
-
-           Returns:
-               str: Answer to the question
-               
-           Raises:
-               RuntimeError: If processing fails"""
-        try:
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-                    
-            response = client.models.generate_content(
-                model=CODE_GENERATION_MODEL,
-                contents=[f"{question}\n{json_data}"],
-                config=types.GenerateContentConfig(
-                    tools=[types.Tool(code_execution=types.ToolCodeExecution)]
-                ),
+       Returns:
+           float: Result number
+            
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+        
+        response = client.models.generate_content(
+            model=ARITHMETIC_MODEL,
+            contents=question,
+            config=types.GenerateContentConfig(
+                tools=[add, subtract, multiply, divide, modulus]
             )
+        )
+
+        return response.text
+    except Exception as e:
+        raise RuntimeError(f"Processing failed: {str(e)}")
+
+@tool("Code Execution Tool")
+def code_execution_tool(question: str, file_path: str) -> str:
+    """Given a question and Python file, execute the file to answer the question.
+
+       Args:
+           question (str): Question to answer
+           file_path (str): The Python file path
+            
+       Returns:
+           str: Answer to the question
             
-            for part in response.candidates[0].content.parts:
-                if part.code_execution_result is not None:
-                    return part.code_execution_result.output
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+
+        file = client.files.upload(file=file_path)
+
+        response = client.models.generate_content(
+            model=CODE_EXECUTION_MODEL,
+            contents=[file, question],
+            config=types.GenerateContentConfig(
+                tools=[types.Tool(code_execution=types.ToolCodeExecution)]
+            ),
+        )
+        
+        for part in response.candidates[0].content.parts:
+            if part.code_execution_result is not None:
+                return part.code_execution_result.output
+    except Exception as e:
+        raise RuntimeError(f"Processing failed: {str(e)}")
+
+@tool("Code Generation Tool")
+def code_generation_tool(question: str, json_data: str) -> str:
+    """Given a question and JSON data, generate and execute code to answer the question.
+
+       Args:
+           question (str): Question to answer
+            file_path (str): The JSON data
+
+       Returns:
+           str: Answer to the question
+           
+       Raises:
+           RuntimeError: If processing fails"""
+    try:
+        client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+                
+        response = client.models.generate_content(
+            model=CODE_GENERATION_MODEL,
+            contents=[f"{question}\n{json_data}"],
+            config=types.GenerateContentConfig(
+                tools=[types.Tool(code_execution=types.ToolCodeExecution)]
+            ),
+        )
+        
+        for part in response.candidates[0].content.parts:
+            if part.code_execution_result is not None:
+                return part.code_execution_result.output
+    except Exception as e:
+        raise RuntimeError(f"Processing failed: {str(e)}")
+            
+@CrewBase
+class GAIACrew():
+	agents: List[BaseAgent]
+	tasks: List[Task]
                         
     # Agents