Update app.py

app.py

@@ -49,8 +49,28 @@ def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
     return AgentExecutor(agent=agent, tools=tools)
 
 def agent_node(state, agent, name):
+    # Run the agent and get its output
     result = agent.invoke(state)
-    return {"messages": [HumanMessage(content=result["output"], name=name)]}
+    output_content = result["output"]
+    
+    # Check if the output contains Python code that generates a graph
+    if "matplotlib" in output_content or "plt." in output_content:
+        exec_locals = {}
+        try:
+            exec(output_content, {}, exec_locals)  # Safely execute the code
+            fig = plt.gcf()  # Get the current matplotlib figure
+            
+            # Save the figure to a buffer
+            buf = io.BytesIO()
+            fig.savefig(buf, format="png")
+            buf.seek(0)
+            
+            # Add image to session state for display
+            st.session_state.graph_image = buf
+        except Exception as e:
+            output_content += f"\nError: {str(e)}"
+    
+    return {"messages": [HumanMessage(content=output_content, name=name)]}
 
 @tool
 def RAG(state):