added graph viz for fun

app/main.py

@@ -2,6 +2,8 @@ import gradio as gr
 from gradio_client import Client
 from langgraph.graph import StateGraph, START, END
 from typing import TypedDict, Optional
+import io
+from PIL import Image
 
 #OPEN QUESTION: SHOULD WE PASS ALL PARAMS FROM THE ORCHESTRATOR TO THE NODES INSTEAD OF SETTING IN EACH MODULE?
 
@@ -99,40 +101,71 @@ ui = gr.Interface(
     flagging_mode="never"
 )
 
+# Add a function to generate the graph visualization
+def get_graph_visualization():
+    """Generate and return the LangGraph workflow visualization as a PIL Image."""
+    # Generate the graph as PNG bytes
+    graph_png_bytes = graph.get_graph().draw_mermaid_png()
+    
+    # Convert bytes to PIL Image for Gradio display
+    graph_image = Image.open(io.BytesIO(graph_png_bytes))
+    return graph_image
+
+
 # Guidance for ChatUI - can be removed later. Questionable whether front end even necessary. Maybe nice to show the graph.
 with gr.Blocks(title="ChatFed Orchestrator") as demo:
     gr.Markdown("# ChatFed Orchestrator")
-    gr.Markdown("This LangGraph server exposes MCP endpoints for the ChatUI module to call.")
-    gr.Markdown("**Available MCP Tools:**")
-    gr.Markdown("- `process_query`: accepts query with optional filters")  
+    gr.Markdown("This LangGraph server exposes MCP endpoints for the ChatUI module to call (which triggers the graph).")
     
-    with gr.Accordion("MCP Endpoint Information", open=True):
-        gr.Markdown(f"""
-        **MCP Server Endpoint:** https://giz-chatfed-orchestrator.hf.space/gradio_api/mcp/sse
-        
-        **For ChatUI Integration:**
-        ```python
-        from gradio_client import Client
-        
-        # Connect to orchestrator
-        orchestrator_client = Client("https://giz-chatfed-orchestrator.hf.space")
-        
-        # Basic usage (no filters)
-        response = orchestrator_client.predict(
-            query="query",
-            api_name="/process_query"
-        )
+    with gr.Row():
+        # Left column - Graph visualization
+        with gr.Column(scale=1):
+            gr.Markdown("**Workflow Visualization**")
+            graph_display = gr.Image(
+                value=get_graph_visualization(), 
+                label="LangGraph Workflow",
+                interactive=False,
+                height=300
+            )
+            
+            # Add a refresh button for the graph
+            refresh_graph_btn = gr.Button("🔄 Refresh Graph", size="sm")
+            refresh_graph_btn.click(
+                fn=get_graph_visualization,
+                outputs=graph_display
+            )
         
-        # Advanced usage with any combination of filters
-        response = orchestrator_client.predict(
-            query="query",
-            reports_filter="annual_reports",
-            sources_filter="internal", 
-            year_filter="2024",
-            api_name="/process_query"
-        )
-        ```
-        """)
+        # Right column - Interface and documentation  
+        with gr.Column(scale=2):
+            gr.Markdown("**Available MCP Tools:**")
+            
+            with gr.Accordion("MCP Endpoint Information", open=True):
+                gr.Markdown(f"""
+                **MCP Server Endpoint:** https://giz-chatfed-orchestrator.hf.space/gradio_api/mcp/sse
+                
+                **For ChatUI Integration:**
+                ```python
+                from gradio_client import Client
+                
+                # Connect to orchestrator
+                orchestrator_client = Client("https://giz-chatfed-orchestrator.hf.space")
+                
+                # Basic usage (no filters)
+                response = orchestrator_client.predict(
+                    query="query",
+                    api_name="/process_query"
+                )
+                
+                # Advanced usage with any combination of filters
+                response = orchestrator_client.predict(
+                    query="query",
+                    reports_filter="annual_reports",
+                    sources_filter="internal", 
+                    year_filter="2024",
+                    api_name="/process_query"
+                )
+                ```
+                """)
 
     with gr.Accordion("Quick Testing Interface", open=True):
         ui.render()

requirements.txt

@@ -1,4 +1,5 @@
 gradio[mcp]
 gradio_client>=1.0.0
 langgraph>=0.2.0
+Pillow>=9.0.0