print image to path

agent.py

@@ -17,8 +17,8 @@ from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingF
 from langchain_anthropic import ChatAnthropic
 
 ##!other libraries 
-from IPython.display import Image, display
-
+from PIL import Image
+import graphviz
 
 load_dotenv()
 
@@ -189,10 +189,12 @@ def build_graph(provider: str = "huggingface"):
     compiled_graph = builder.compile() # This line should already be there or be the next line
 
     try:
-        display(Image(graph.get_graph().draw_mermaid_png()))
+        dot_path = "graph"
+        compiled_graph.get_graph().render(dot_path, format="png", cleanup=True)
     except Exception:
         print('cant print the graph')
         # This requires some extra dependencies and is optional
         pass
 
-    return compiled_graph 
+    return compiled_graph 
+

app.py

@@ -4,6 +4,7 @@ from huggingface_hub import InferenceClient
 from langchain_core.messages import HumanMessage
 from langchain.agents import AgentExecutor
 from agent import build_graph
+from PIL import Image
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
@@ -35,18 +36,42 @@ try:
 except Exception as e:
         print(f"Error instantiating agent: {e}")
 
+def show_graph():
+    if not os.path.exists("graph.png"):
+        return None
+    return Image.open("graph.png")
+
+
 def run_langgraph_agent(user_input: str):
     graph = build_graph()
     result = graph.invoke({"messages": [HumanMessage(content=user_input)]})
     return result["messages"][-1].content if "messages" in result else result
 
 
-demo = gr.Interface(
-    fn=run_langgraph_agent,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your message..."),
-    outputs="text",
-    title="LangGraph Agent Chat",
-)
+# 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")
+
+    # Wire up the buttons to functions
+    submit_btn.click(run_langgraph_agent, inputs=user_input, outputs=chat_output)
+    graph_btn.click(show_graph, inputs=None, outputs=graph_image)
 
 if __name__ == "__main__":
     demo.launch()