Update app.py

app.py

@@ -1,10 +1,8 @@
 import json
 import gradio as gr
 import plotly.graph_objects as go
-import plotly.express as px
 import networkx as nx
-from typing import List, Dict, Any
-
+from typing import List, Dict, Optional
 
 from langchain_openai.chat_models import ChatOpenAI
 from dialog2graph.pipelines.model_storage import ModelStorage
@@ -21,8 +19,15 @@ def initialize_pipeline():
     )
     return D2GLLMPipeline("d2g_pipeline", model_storage=ms, filling_llm="my_filling_model")
 
-def load_dialog_data(json_file: str) -> List[Dict[str, str]]:
-    """Load dialog data from JSON file"""
+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"""
+    if json_file == "custom" and custom_dialog_json:
+        try:
+            return json.loads(custom_dialog_json)
+        except json.JSONDecodeError as e:
+            gr.Error(f"Invalid JSON format in custom dialog: {str(e)}")
+            return []
+    
     file_path = f"{json_file}.json"
     try:
         with open(file_path, 'r') as f:
@@ -178,11 +183,11 @@ 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) -> tuple:
+def process_dialog_and_visualize(dialog_choice: str, custom_dialog: str = "") -> tuple:
     """Process the selected dialog and create visualization"""
     try:
         # Load the selected dialog data
-        dialog_data = load_dialog_data(dialog_choice)
+        dialog_data = load_dialog_data(dialog_choice, custom_dialog if dialog_choice == "custom" else None)
         
         if not dialog_data:
             return None, "Failed to load dialog data", ""
@@ -228,10 +233,18 @@ def create_gradio_app():
         with gr.Row():
             with gr.Column(scale=1):
                 dialog_selector = gr.Radio(
-                    choices=["dialog1", "dialog2", "dialog3"],
+                    choices=["dialog1", "dialog2", "dialog3", "custom"],
                     label="Select Dialog Dataset",
                     value="dialog1",
-                    info="Choose one of the available dialog datasets"
+                    info="Choose one of the available dialog datasets or use custom JSON"
+                )
+                
+                custom_dialog_input = gr.Textbox(
+                    label="Custom Dialog JSON",
+                    placeholder='[{"text": "Hello! How can I help?", "participant": "assistant"}, {"text": "I need assistance", "participant": "user"}]',
+                    lines=8,
+                    visible=False,
+                    info="Enter dialog data as JSON array with 'text' and 'participant' fields"
                 )
                 
                 process_btn = gr.Button(
@@ -245,6 +258,7 @@ def create_gradio_app():
                     - **dialog1**: Hotel booking conversation
                     - **dialog2**: Food delivery conversation  
                     - **dialog3**: Technical support conversation
+                    - **custom**: Provide your own dialog as JSON
                     """)
             
             with gr.Column(scale=3):
@@ -259,16 +273,31 @@ def create_gradio_app():
                 chat_output = gr.HTML(label="Chat Visualization")
         
         # Event handlers
+        def toggle_custom_input(choice):
+            return gr.update(visible=(choice == "custom"))
+        
+        dialog_selector.change(
+            fn=toggle_custom_input,
+            inputs=[dialog_selector],
+            outputs=[custom_dialog_input]
+        )
+        
         process_btn.click(
             fn=process_dialog_and_visualize,
-            inputs=[dialog_selector],
+            inputs=[dialog_selector, custom_dialog_input],
             outputs=[plot_output, summary_output, chat_output]
         )
         
-        # Auto-process on selection change
+        # Auto-process on selection change (but not for custom to avoid premature processing)
+        def auto_process(choice, custom_text):
+            if choice != "custom":
+                return process_dialog_and_visualize(choice, custom_text)
+            else:
+                return None, "Select 'Process Dialog & Generate Graph' to process custom dialog", ""
+        
         dialog_selector.change(
-            fn=process_dialog_and_visualize,
-            inputs=[dialog_selector],
+            fn=auto_process,
+            inputs=[dialog_selector, custom_dialog_input],
             outputs=[plot_output, summary_output, chat_output]
         )