Add Python code tool

system_prompt.txt

@@ -1 +1,5 @@
-You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. Even if you're uncertain, provide the best guess in this format.
+YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
+If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
+If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
+If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.

tools/file_tools.py

@@ -1,6 +1,9 @@
 import pandas as pd
+import io
+import contextlib
 from langchain_core.tools import tool
 
+
 @tool
 def analyze_excel_file(file_path: str, query: str) -> str:
     """
@@ -26,4 +29,29 @@ def analyze_excel_file(file_path: str, query: str) -> str:
         return result
 
     except Exception as e:
-        return f"Error analyzing Excel file: {str(e)}"
+        return f"Error analyzing Excel file: {str(e)}"
+
+
+@tool
+def execute_python_code(code: str) -> str:
+    """
+    Executes a Python code string and returns the output or error.
+    Args:
+        code (str): The Python code to execute.
+    Returns:
+        str: The output or error message.
+    """
+    local_vars = {}
+    stdout = io.StringIO()
+    try:
+        with contextlib.redirect_stdout(stdout):
+            exec(code, {}, local_vars)
+        output = stdout.getvalue()
+        if output.strip():
+            return output.strip()
+        # If code defines a variable named 'result', return its value
+        if "result" in local_vars:
+            return str(local_vars["result"])
+        return "Code executed successfully, but produced no output."
+    except Exception as e:
+        return f"Error executing code: {e}"