Update tools/ai_tools.py

tools/ai_tools.py

@@ -223,27 +223,23 @@ class AITools():
         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.
-    
+    @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 Python file path
-                
+                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"])
-    
-            file = client.files.upload(file=file_path)
-    
+                    
             response = client.models.generate_content(
-                model=CODE_EXECUTION_MODEL,
-                contents=[file, question],
+                model=CODE_GENERATION_MODEL,
+                contents=[f"{question}\n{json_data}"],
                 config=types.GenerateContentConfig(
                     tools=[types.Tool(code_execution=types.ToolCodeExecution)]
                 ),
@@ -254,24 +250,28 @@ class AITools():
                     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:
+        """Given a question and Python file, execute the file to answer the question.
     
-    @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
+               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"])
-                    
+    
+            file = client.files.upload(file=file_path)
+    
             response = client.models.generate_content(
-                model=CODE_GENERATION_MODEL,
-                contents=[f"{question}\n{json_data}"],
+                model=CODE_EXECUTION_MODEL,
+                contents=[file, question],
                 config=types.GenerateContentConfig(
                     tools=[types.Tool(code_execution=types.ToolCodeExecution)]
                 ),