add image

app.py

@@ -45,35 +45,18 @@ def show_graph():
 def run_langgraph_agent(user_input: str):
     graph = build_graph()
     result = graph.invoke({"messages": [HumanMessage(content=user_input)]})
-    if "messages" in result and len(result["messages"]) > 0:
-        answer = result["messages"][-1].content
-        return str(answer)
-    return "No answer generated"
-
-
-# demo = gr.Interface(
-#     fn=run_langgraph_agent,
-#     inputs=gr.Textbox(lines=2, placeholder="Enter your message..."),
-#     outputs="text",
-#     title="LangGraph Agent Chat",
-# )
-
-# Create Gradio Blocks for flexible layout
-with gr.Blocks() as demo:
-    gr.Markdown("# LangGraph Agent Chat + Graph Viewer")
-
-    with gr.Row():
-        with gr.Column(scale=2):
-            user_input = gr.Textbox(lines=2, placeholder="Enter your message...")
-            submit_btn = gr.Button("Send")
-            chat_output = gr.Textbox(label="Agent Reply")
-
-        with gr.Column(scale=1):
-            graph_btn = gr.Button("Show Graph")
-            graph_image = gr.Image(label="LangGraph Visualization")
-
-    submit_btn.click(run_langgraph_agent, inputs=user_input, outputs=chat_output)
-    graph_btn.click(show_graph, inputs=None, outputs=graph_image)
+    text_output = result["messages"][-1].content if "messages" in result else result
+
+    image = show_graph()  # return PIL.Image or None
+    return text_output, image
+
+
+demo = gr.Interface(
+    fn=run_langgraph_agent,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your message..."),
+    outputs=[gr.Textbox(), gr.Image(type="pil")],
+    title="LangGraph Agent Chat",
+)
 
 if __name__ == "__main__":
     demo.launch()