Update app.py

app.py

@@ -19,25 +19,41 @@ def initialize_pipeline():
     )
     return D2GLLMPipeline("d2g_pipeline", model_storage=ms, filling_llm="my_filling_model")
 
-def load_dialog_data(json_file: str, custom_dialog_json: Optional[str] = None) -> List[Dict[str, str]]:
-    """Load dialog data from JSON file or custom JSON string"""
+def load_dialog_data(json_file: str, custom_dialog_json: Optional[str] = None) -> List[List[Dict[str, str]]]:
+    """Load dialog data from JSON file or custom JSON string and ensure it's in list[list[dict]] format"""
     if json_file == "custom" and custom_dialog_json:
         try:
-            return json.loads(custom_dialog_json)
+            data = json.loads(custom_dialog_json)
         except json.JSONDecodeError as e:
             gr.Error(f"Invalid JSON format in custom dialog: {str(e)}")
             return []
+    else:
+        file_path = f"{json_file}.json"
+        try:
+            with open(file_path, 'r') as f:
+                data = json.load(f)
+        except FileNotFoundError:
+            gr.Error(f"File {file_path} not found!")
+            return []
+        except json.JSONDecodeError:
+            gr.Error(f"Invalid JSON format in {file_path}!")
+            return []
     
-    file_path = f"{json_file}.json"
-    try:
-        with open(file_path, 'r') as f:
-            return json.load(f)
-    except FileNotFoundError:
-        gr.Error(f"File {file_path} not found!")
-        return []
-    except json.JSONDecodeError:
-        gr.Error(f"Invalid JSON format in {file_path}!")
+    # Convert to list[list[dict]] format if needed
+    if not data:
         return []
+    
+    # Check if it's already list[list[dict]]
+    if isinstance(data, list) and len(data) > 0:
+        if isinstance(data[0], list):
+            # Already in list[list[dict]] format
+            return data
+        elif isinstance(data[0], dict):
+            # Convert list[dict] to list[list[dict]]
+            return [data]
+    
+    # If it's something else, wrap it in double list
+    return [data] if isinstance(data, list) else [[data]]
 
 def create_network_visualization(graph: nx.Graph) -> go.Figure:
     """Create a Plotly network visualization from NetworkX graph"""
@@ -149,9 +165,12 @@ def create_network_visualization(graph: nx.Graph) -> go.Figure:
     
     return fig
 
-def create_chat_visualization(dialog_data: List[Dict[str, str]]) -> str:
-    """Create a chat-like visualization of the dialog"""
-    chat_html = """
+def create_chat_visualization(dialog_data: List[Dict[str, str]], dialog_index: int = 0, total_dialogs: int = 1) -> str:
+    """Create a chat-like visualization of the dialog with navigation info"""
+    chat_html = f"""
+    <div style="margin-bottom: 10px; text-align: center; font-weight: bold; color: #666;">
+        Dialog {dialog_index + 1} of {total_dialogs}
+    </div>
     <div style="max-height: 500px; overflow-y: auto; border: 1px solid #ddd; border-radius: 10px; padding: 20px; background-color: #f9f9f9;">
     """
     
@@ -183,46 +202,57 @@ def create_chat_visualization(dialog_data: List[Dict[str, str]]) -> str:
     chat_html += "</div>"
     return chat_html
 
-def process_dialog_and_visualize(dialog_choice: str, custom_dialog: str = "") -> tuple:
+def process_dialog_and_visualize(dialog_choice: str, custom_dialog: str = "", current_dialog_index: int = 0) -> tuple:
     """Process the selected dialog and create visualization"""
     try:
         # Load the selected dialog data
-        dialog_data = load_dialog_data(dialog_choice, custom_dialog if dialog_choice == "custom" else None)
+        dialog_data_list = load_dialog_data(dialog_choice, custom_dialog if dialog_choice == "custom" else None)
+        
+        if not dialog_data_list:
+            return None, "Failed to load dialog data", "", 0, len(dialog_data_list), gr.update(visible=False), gr.update(visible=False)
         
-        if not dialog_data:
-            return None, "Failed to load dialog data", ""
+        # Ensure current_dialog_index is within bounds
+        current_dialog_index = max(0, min(current_dialog_index, len(dialog_data_list) - 1))
         
         # Initialize pipeline
         pipe = initialize_pipeline()
         
-        # Process the data
-        data = PipelineRawDataType(dialogs=dialog_data)
+        # Process the data (use all dialogs for graph generation)
+        data = PipelineRawDataType(dialogs=dialog_data_list)
         graph, report = pipe.invoke(data)
         
         # Create visualization
         fig = create_network_visualization(graph.graph)
         
-        # Create chat visualization
-        chat_viz = create_chat_visualization(dialog_data)
+        # Create chat visualization for the current dialog
+        current_dialog = dialog_data_list[current_dialog_index]
+        chat_viz = create_chat_visualization(current_dialog, current_dialog_index, len(dialog_data_list))
         
         # Create summary information
         num_nodes = graph.graph.number_of_nodes()
         num_edges = graph.graph.number_of_edges()
+        total_turns = sum(len(dialog) for dialog in dialog_data_list)
         
         summary = f"""
         ## Graph Summary
         - **Number of nodes**: {num_nodes}
         - **Number of edges**: {num_edges}
-        - **Dialog turns**: {len(dialog_data)}
+        - **Total dialogs**: {len(dialog_data_list)}
+        - **Total dialog turns**: {total_turns}
+        - **Currently viewing**: Dialog {current_dialog_index + 1} ({len(current_dialog)} turns)
         
         ## Processing Report
-        Generated graph from {len(dialog_data)} dialog turns with {num_nodes} nodes and {num_edges} edges.
+        Generated graph from {len(dialog_data_list)} dialog(s) with {total_turns} total turns resulting in {num_nodes} nodes and {num_edges} edges.
         """
         
-        return fig, summary, chat_viz
+        # Show navigation buttons only if there are multiple dialogs
+        show_nav = len(dialog_data_list) > 1
+        
+        return (fig, summary, chat_viz, current_dialog_index, len(dialog_data_list), 
+                gr.update(visible=show_nav), gr.update(visible=show_nav))
         
     except Exception as e:
-        return None, f"Error processing dialog: {str(e)}", ""
+        return None, f"Error processing dialog: {str(e)}", "", 0, 0, gr.update(visible=False), gr.update(visible=False)
 
 # Create the Gradio interface
 def create_gradio_app():
@@ -230,10 +260,14 @@ def create_gradio_app():
         gr.Markdown("# Dialog2Graph Interactive Visualizer")
         gr.Markdown("Select a dialog dataset to process and visualize as a graph network using Plotly.")
         
+        # State variables for dialog navigation
+        current_dialog_index_state = gr.State(0)
+        total_dialogs_state = gr.State(0)
+        
         with gr.Row():
             with gr.Column(scale=1):
                 dialog_selector = gr.Radio(
-                    choices=["dialog1", "dialog2", "dialog3", "custom"],
+                    choices=["dialog1", "dialog2", "dialog3", "dialog4", "custom"],
                     label="Select Dialog Dataset",
                     value="dialog1",
                     info="Choose one of the available dialog datasets or use custom JSON"
@@ -258,6 +292,7 @@ def create_gradio_app():
                     - **dialog1**: Hotel booking conversation
                     - **dialog2**: Food delivery conversation  
                     - **dialog3**: Technical support conversation
+                    - **dialog4**: Multiple dialogs (calendar, support, subscription)
                     - **custom**: Provide your own dialog as JSON
                     """)
             
@@ -270,8 +305,46 @@ def create_gradio_app():
             
             with gr.Column(scale=1):
                 gr.Markdown("### Dialog Conversation")
+                
+                # Navigation controls for multiple dialogs
+                with gr.Row(visible=False) as nav_row:
+                    prev_btn = gr.Button("← Previous Dialog", size="sm")
+                    next_btn = gr.Button("Next Dialog →", size="sm")
+                
                 chat_output = gr.HTML(label="Chat Visualization")
         
+        # Navigation functions
+        def navigate_dialog(direction: int, current_index: int, total_dialogs: int, dialog_choice: str, custom_dialog: str):
+            if total_dialogs <= 1:
+                return current_index, "", ""
+                
+            new_index = current_index + direction
+            new_index = max(0, min(new_index, total_dialogs - 1))
+            
+            try:
+                dialog_data_list = load_dialog_data(dialog_choice, custom_dialog if dialog_choice == "custom" else None)
+                if dialog_data_list and new_index < len(dialog_data_list):
+                    current_dialog = dialog_data_list[new_index]
+                    chat_viz = create_chat_visualization(current_dialog, new_index, len(dialog_data_list))
+
+                    total_turns = sum(len(dialog) for dialog in dialog_data_list)
+                    
+                    summary = f"""
+                    ## Graph Summary
+                    - **Total dialogs**: {len(dialog_data_list)}
+                    - **Total dialog turns**: {total_turns}
+                    - **Currently viewing**: Dialog {new_index + 1} ({len(current_dialog)} turns)
+                    
+                    ## Processing Report
+                    Navigate between dialogs to view different conversations.
+                    """
+                    
+                    return new_index, summary, chat_viz
+            except Exception:
+                pass
+                
+            return current_index, "", ""
+        
         # Event handlers
         def toggle_custom_input(choice):
             return gr.update(visible=(choice == "custom"))
@@ -284,21 +357,34 @@ def create_gradio_app():
         
         process_btn.click(
             fn=process_dialog_and_visualize,
-            inputs=[dialog_selector, custom_dialog_input],
-            outputs=[plot_output, summary_output, chat_output]
+            inputs=[dialog_selector, custom_dialog_input, current_dialog_index_state],
+            outputs=[plot_output, summary_output, chat_output, current_dialog_index_state, total_dialogs_state, nav_row, nav_row]
+        )
+        
+        # Navigation button handlers
+        prev_btn.click(
+            fn=lambda curr_idx, total, choice, custom: navigate_dialog(-1, curr_idx, total, choice, custom),
+            inputs=[current_dialog_index_state, total_dialogs_state, dialog_selector, custom_dialog_input],
+            outputs=[current_dialog_index_state, summary_output, chat_output]
+        )
+        
+        next_btn.click(
+            fn=lambda curr_idx, total, choice, custom: navigate_dialog(1, curr_idx, total, choice, custom),
+            inputs=[current_dialog_index_state, total_dialogs_state, dialog_selector, custom_dialog_input],
+            outputs=[current_dialog_index_state, summary_output, chat_output]
         )
         
         # Auto-process on selection change (but not for custom to avoid premature processing)
-        def auto_process(choice, custom_text):
+        def auto_process(choice, custom_text, curr_idx):
             if choice != "custom":
-                return process_dialog_and_visualize(choice, custom_text)
+                return process_dialog_and_visualize(choice, custom_text, curr_idx)
             else:
-                return None, "Select 'Process Dialog & Generate Graph' to process custom dialog", ""
+                return None, "Select 'Process Dialog & Generate Graph' to process custom dialog", "", 0, 0, gr.update(visible=False), gr.update(visible=False)
         
         dialog_selector.change(
             fn=auto_process,
-            inputs=[dialog_selector, custom_dialog_input],
-            outputs=[plot_output, summary_output, chat_output]
+            inputs=[dialog_selector, custom_dialog_input, current_dialog_index_state],
+            outputs=[plot_output, summary_output, chat_output, current_dialog_index_state, total_dialogs_state, nav_row, nav_row]
         )
     
     return app