Update crew.py

crew.py

@@ -180,16 +180,47 @@ def run_crew(question, file_path):
         except Exception as e:
             raise RuntimeError(f"Processing failed: {str(e)}")
 
+    @tool("Code Generation Tool")
+    def code_generation_tool(question: str) -> str:
+        """Generate and execute Python code to answer a 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"])
+
+            file = client.files.upload(file=file_path)
+
+            response = client.models.generate_content(
+                model=CODE_MODEL,
+                contents=[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 Execution Tool")
     def code_execution_tool(question: str, file_path: str) -> str:
-        """Execute a Python code file and answer a question about it.
+        """Execute a Python code file to answer a question.
     
            Args:
-               question (str): Question about a Pyton code file
+               question (str): Question to answer
                file_path (str): The Python code file path
                 
            Returns:
-               str: Answer to the question about the Python code file
+               str: Answer to the question
                 
            Raises:
                RuntimeError: If processing fails"""
@@ -279,6 +310,17 @@ def run_crew(question, file_path):
         tools=[document_analysis_tool],
         verbose=False
     )
+
+    code_generation_agent = Agent(
+        role="Code Generation Agent",
+        goal="Generate Python code and execute it to help answer question \"{question}\"",
+        backstory="As an expert Python code generation assistant, you generate and execute code to help answer the question.",
+        allow_delegation=False,
+        llm=AGENT_MODEL,
+        max_iter=3,
+        tools=[code_execution_tool],
+        verbose=False
+    )
     
     code_execution_agent = Agent(
         role="Code Execution Agent",
@@ -319,7 +361,8 @@ def run_crew(question, file_path):
                 audio_analysis_agent, 
                 video_analysis_agent, 
                 youtube_analysis_agent, 
-                #document_analysis_agent, 
+                document_analysis_agent,
+                code_generation_agent,
                 code_execution_agent],
         manager_agent=manager_agent,
         tasks=[manager_task],