Update app.py

app.py

@@ -48,14 +48,9 @@ repl = PythonREPL()
 def python_repl(code: Annotated[str, "The python code to execute to generate your chart."]):
     """Executes Python code to generate a chart and returns the chart as a base64-encoded image."""
     try:
-        # Log the executed code for debugging
-        st.write("Executing Code:", code)
-
-        # Prepare execution environment
+        # Execute the code
         exec_globals = {"plt": plt}
         exec_locals = {}
-
-        # Execute the provided code
         exec(code, exec_globals, exec_locals)
 
         # Save the generated plot to a buffer
@@ -67,11 +62,9 @@ def python_repl(code: Annotated[str, "The python code to execute to generate you
         plt.clf()
         plt.close()
 
-        # Encode the image as base64
+        # Encode image as base64
         encoded_image = base64.b64encode(buf.getvalue()).decode("utf-8")
         return {"status": "success", "image": encoded_image}
-    except FileNotFoundError as e:
-        return {"status": "failed", "error": f"FileNotFoundError: {e}"}
     except Exception as e:
         return {"status": "failed", "error": repr(e)}
 
@@ -91,27 +84,21 @@ def tool_node(state):
     tool_name = last_message.additional_kwargs["function_call"]["name"]
     action = ToolInvocation(tool=tool_name, tool_input=tool_input)
     response = tool_executor.invoke(action)
-    if isinstance(response, dict):
-        if response.get("status") == "success" and "image" in response:
-            return {
-                "messages": [
-                    {
-                        "role": "assistant",
-                        "content": "Image generated successfully.",
-                        "image": response["image"],
-                    }
-                ]
-            }
-        else:
-            # Handle errors more informatively
-            return {
-                "messages": [
-                    {
-                        "role": "assistant",
-                        "content": f"Error in {tool_name}: {response.get('error', 'Unknown error')}"
-                    }
-                ]
-            }
+    if isinstance(response, dict) and response.get("status") == "success" and "image" in response:
+        return {
+            "messages": [
+                {
+                    "role": "assistant",
+                    "content": "Image generated successfully.",
+                    "image": response["image"],
+                }
+            ]
+        }
+    else:
+        function_message = FunctionMessage(
+            content=f"{tool_name} response: {str(response)}", name=action.tool
+        )
+        return {"messages": [function_message]}
 
 # Define router
 def router(state):