From a98f48315cbce19db868740426c63ad3aa12b7d1 Mon Sep 17 00:00:00 2001
From: ochafik <ochafik@google.com>
Date: Wed, 10 Apr 2024 09:27:25 +0100
Subject: [PATCH] agent: python tool: return errors

---
 examples/agent/tools/unsafe_python_tools.py | 33 ++++++++++++---------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/examples/agent/tools/unsafe_python_tools.py b/examples/agent/tools/unsafe_python_tools.py
index 3473a9254..b187f219e 100644
--- a/examples/agent/tools/unsafe_python_tools.py
+++ b/examples/agent/tools/unsafe_python_tools.py
@@ -22,18 +22,23 @@ def execute_python(source: str) -> Union[Dict, str]:
         Returns:
             dict | str: A dictionary containing variables declared, or an error message if an exception occurred.
     """
-    namespace = {}
-    sys.stderr.write(f"Executing Python program:\n{source}\n")
-    exec(source, namespace)
-    results = {
-        k: v
-        for k, v in namespace.items()
-        if not k.startswith('_') \
-            and not isinstance(v, type) \
-            and not isinstance(v, types.ModuleType) \
-            and not callable(v) \
-            and _is_serializable(v)
-    }
-    sys.stderr.write(f"Results: {json.dumps(results, indent=2)}\n")
+    try:
+        namespace = {}
+        sys.stderr.write(f"Executing Python program:\n{source}\n")
+        exec(source, namespace)
+        results = {
+            k: v
+            for k, v in namespace.items()
+            if not k.startswith('_') \
+                and not isinstance(v, type) \
+                and not isinstance(v, types.ModuleType) \
+                and not callable(v) \
+                and _is_serializable(v)
+        }
+        sys.stderr.write(f"Results: {json.dumps(results, indent=2)}\n")
+        return results
+    except Exception as e:
+        msg = f"Error: {sys.exc_info()[1]}"
+        sys.stderr.write(f"{msg}\n")
+        return msg
 
-    return results