Update app.py

app.py

@@ -48,9 +48,14 @@ 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:
-        # Execute the code
+        # Log the executed code for debugging
+        st.write("Executing Code:", code)
+
+        # Prepare execution environment
         exec_globals = {"plt": plt}
         exec_locals = {}
+
+        # Execute the provided code
         exec(code, exec_globals, exec_locals)
 
         # Save the generated plot to a buffer
@@ -62,9 +67,11 @@ def python_repl(code: Annotated[str, "The python code to execute to generate you
         plt.clf()
         plt.close()
 
-        # Encode image as base64
+        # Encode the 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)}
 
@@ -84,21 +91,27 @@ 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) 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]}
+    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')}"
+                    }
+                ]
+            }
 
 # Define router
 def router(state):
@@ -192,4 +205,4 @@ if st.button("Run Workflow"):
                             # Display any text content
                             st.write(message["content"])
         except Exception as e:
-            st.error(f"An error occurred: {e}")
+            st.error(f"An error occurred: {e}")