Add DAG visualization functionality for robot task planning

app.py

@@ -3,6 +3,7 @@ import spaces  # Import spaces module for ZeroGPU
 from huggingface_hub import login
 import os
 from json_processor import JsonProcessor
+from dag_visualizer import DAGVisualizer
 import json
 
 # 1) Read Secrets
@@ -213,6 +214,31 @@ def generate_response_gpu(prompt, max_tokens=512, selected_model=DEFAULT_MODEL):
     except Exception as generation_error:
         return f"❌ Generation Error: {str(generation_error)}"
 
+def create_dag_visualization(task_json_str):
+    """Create DAG visualization from task JSON"""
+    try:
+        if not task_json_str.strip():
+            return None, "Please provide task JSON data"
+        
+        # Parse JSON
+        task_data = json.loads(task_json_str)
+        
+        # Create DAG visualizer
+        dag_visualizer = DAGVisualizer()
+        
+        # Generate visualization
+        image_path = dag_visualizer.create_dag_visualization(task_data)
+        
+        if image_path:
+            return image_path, "✅ DAG visualization created successfully!"
+        else:
+            return None, "❌ Failed to create DAG visualization"
+            
+    except json.JSONDecodeError as e:
+        return None, f"❌ JSON Parse Error: {str(e)}"
+    except Exception as e:
+        return None, f"❌ DAG Creation Error: {str(e)}"
+
 def chat_interface(message, history, max_tokens, selected_model):
     """Chat interface - runs on CPU, calls GPU functions"""
     if not message.strip():
@@ -255,7 +281,7 @@ with gr.Blocks(
     - **⚖️ Dart-llm-model-3B**: Balanced performance and quality  
     - **🎯 Dart-llm-model-8B**: Best quality output, higher latency
     
-    **Capabilities**: Convert natural language robot commands into structured task sequences for excavators, dump trucks, and other construction robots.
+    **Capabilities**: Convert natural language robot commands into structured task sequences for excavators, dump trucks, and other construction robots. **Now with DAG Visualization!**
     
     **Models**: 
     - [YongdongWang/llama-3.2-1b-lora-qlora-dart-llm](https://huggingface.co/YongdongWang/llama-3.2-1b-lora-qlora-dart-llm) (Default)
@@ -265,59 +291,100 @@ with gr.Blocks(
     ⚡ **Using ZeroGPU**: This Space uses dynamic GPU allocation (Nvidia H200). First generation might take a bit longer.
     """)
     
-    with gr.Row():
-        with gr.Column(scale=3):
-            chatbot = gr.Chatbot(
-                label="Task Planning Results",
-                height=500,
-                show_label=True,
-                container=True,
-                bubble_full_width=False,
-                show_copy_button=True
-            )
-            
-            msg = gr.Textbox(
-                label="Robot Command",
-                placeholder="Enter robot task command (e.g., 'Deploy Excavator 1 to Soil Area 1')...",
-                lines=2,
-                max_lines=5,
-                show_label=True,
-                container=True
-            )
-            
+    with gr.Tabs():
+        with gr.Tab("💬 Task Planning"):
             with gr.Row():
-                send_btn = gr.Button("🚀 Generate Tasks", variant="primary", size="sm")
-                clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="sm")
+                with gr.Column(scale=3):
+                    chatbot = gr.Chatbot(
+                        label="Task Planning Results",
+                        height=500,
+                        show_label=True,
+                        container=True,
+                        bubble_full_width=False,
+                        show_copy_button=True
+                    )
+                    
+                    msg = gr.Textbox(
+                        label="Robot Command",
+                        placeholder="Enter robot task command (e.g., 'Deploy Excavator 1 to Soil Area 1')...",
+                        lines=2,
+                        max_lines=5,
+                        show_label=True,
+                        container=True
+                    )
+                    
+                    with gr.Row():
+                        send_btn = gr.Button("🚀 Generate Tasks", variant="primary", size="sm")
+                        clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="sm")
+                
+                with gr.Column(scale=1):
+                    gr.Markdown("### ⚙️ Generation Settings")
+                    
+                    model_selector = gr.Dropdown(
+                        choices=[(config["name"], key) for key, config in MODEL_CONFIGS.items()],
+                        value=DEFAULT_MODEL,
+                        label="Model Size",
+                        info="Select model size (1B = fastest, 8B = best quality)",
+                        interactive=True
+                    )
+                    
+                    max_tokens = gr.Slider(
+                        minimum=50,
+                        maximum=5000,
+                        value=512,
+                        step=10,
+                        label="Max Tokens",
+                        info="Maximum number of tokens to generate"
+                    )
+                    
+                    gr.Markdown("""
+                    ### 📊 Model Status
+                    - **Hardware**: ZeroGPU (Dynamic Nvidia H200)
+                    - **Status**: Ready
+                    - **Note**: First generation allocates GPU resources
+                    - **Dart-llm-model-1B**: Fastest inference (Default)
+                    - **Dart-llm-model-3B**: Balanced speed/quality
+                    - **Dart-llm-model-8B**: Best quality, slower
+                    """)
         
-        with gr.Column(scale=1):
-            gr.Markdown("### ⚙️ Generation Settings")
-            
-            model_selector = gr.Dropdown(
-                choices=[(config["name"], key) for key, config in MODEL_CONFIGS.items()],
-                value=DEFAULT_MODEL,
-                label="Model Size",
-                info="Select model size (1B = fastest, 8B = best quality)",
-                interactive=True
-            )
-            
-            max_tokens = gr.Slider(
-                minimum=50,
-                maximum=5000,
-                value=512,
-                step=10,
-                label="Max Tokens",
-                info="Maximum number of tokens to generate"
-            )
-            
-            gr.Markdown("""
-            ### 📊 Model Status
-            - **Hardware**: ZeroGPU (Dynamic Nvidia H200)
-            - **Status**: Ready
-            - **Note**: First generation allocates GPU resources
-            - **Dart-llm-model-1B**: Fastest inference (Default)
-            - **Dart-llm-model-3B**: Balanced speed/quality
-            - **Dart-llm-model-8B**: Best quality, slower
-            """)
+        with gr.Tab("📊 DAG Visualization"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    json_input = gr.Textbox(
+                        label="Task JSON Data",
+                        placeholder="Paste the generated task JSON here to create a DAG visualization...",
+                        lines=15,
+                        max_lines=25,
+                        show_label=True,
+                        container=True
+                    )
+                    
+                    with gr.Row():
+                        dag_btn = gr.Button("🎨 Generate DAG", variant="primary", size="sm")
+                        dag_clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="sm")
+                    
+                    dag_status = gr.Textbox(
+                        label="Status",
+                        value="Ready to generate DAG visualization",
+                        interactive=False,
+                        show_label=True
+                    )
+                
+                with gr.Column(scale=3):
+                    dag_output = gr.Image(
+                        label="Task Dependency Graph",
+                        show_label=True,
+                        container=True,
+                        height=600
+                    )
+                    
+                    gr.Markdown("""
+                    ### 📈 DAG Features
+                    - **Node Colors**: Red (Start), Orange (Intermediate), Purple (End)
+                    - **Arrows**: Show task dependencies
+                    - **Layout**: Hierarchical based on dependencies
+                    - **Details**: Task info boxes with robots and objects
+                    """)
     
     # Example conversations
     gr.Examples(
@@ -350,6 +417,18 @@ with gr.Blocks(
         lambda: ([], ""),
         outputs=[chatbot, msg]
     )
+    
+    # DAG visualization event handlers
+    dag_btn.click(
+        create_dag_visualization,
+        inputs=[json_input],
+        outputs=[dag_output, dag_status]
+    )
+    
+    dag_clear_btn.click(
+        lambda: ("", None, "Ready to generate DAG visualization"),
+        outputs=[json_input, dag_output, dag_status]
+    )
 
 if __name__ == "__main__":
     app.launch(

dag_visualizer.py

@@ -0,0 +1,334 @@
+import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.use('Agg')  # Use non-interactive backend for server environments
+import networkx as nx
+import json
+import numpy as np
+from loguru import logger
+import os
+import tempfile
+from datetime import datetime
+
+class DAGVisualizer:
+    def __init__(self):
+        # Configure Matplotlib to use IEEE-style parameters
+        plt.rcParams.update({
+            'font.family': 'DejaVu Sans',  # Use available font instead of Times New Roman
+            'font.size': 10,
+            'axes.linewidth': 1.2,
+            'axes.labelsize': 12,
+            'xtick.labelsize': 10,
+            'ytick.labelsize': 10,
+            'legend.fontsize': 10,
+            'figure.titlesize': 14
+        })
+    
+    def create_dag_from_tasks(self, task_data):
+        """
+        Create a directed graph from task data.
+        
+        Args:
+            task_data: Dictionary containing tasks with structure like:
+                {
+                    "tasks": [
+                        {
+                            "task": "task_name",
+                            "instruction_function": {
+                                "name": "function_name",
+                                "robot_ids": ["robot1", "robot2"],
+                                "dependencies": ["dependency_task"],
+                                "object_keywords": ["object1", "object2"]
+                            }
+                        }
+                    ]
+                }
+        
+        Returns:
+            NetworkX DiGraph object
+        """
+        if not task_data or "tasks" not in task_data:
+            logger.warning("Invalid task data structure")
+            return None
+        
+        # Create a directed graph
+        G = nx.DiGraph()
+        
+        # Add nodes and store mapping from task name to ID
+        task_mapping = {}
+        for i, task in enumerate(task_data["tasks"]):
+            task_id = i + 1
+            task_name = task["task"]
+            task_mapping[task_name] = task_id
+            
+            # Add node with attributes
+            G.add_node(task_id, 
+                      name=task_name,
+                      function=task["instruction_function"]["name"],
+                      robots=task["instruction_function"].get("robot_ids", []),
+                      objects=task["instruction_function"].get("object_keywords", []))
+        
+        # Add dependency edges
+        for i, task in enumerate(task_data["tasks"]):
+            task_id = i + 1
+            dependencies = task["instruction_function"]["dependencies"]
+            for dep in dependencies:
+                if dep in task_mapping:
+                    dep_id = task_mapping[dep]
+                    G.add_edge(dep_id, task_id)
+        
+        return G
+    
+    def calculate_layout(self, G):
+        """
+        Calculate hierarchical layout for the graph based on dependencies.
+        """
+        if not G:
+            return {}
+        
+        # Calculate layers based on dependencies
+        layers = {}
+        
+        def get_layer(node_id, visited=None):
+            if visited is None:
+                visited = set()
+            if node_id in visited:
+                return 0
+            visited.add(node_id)
+            
+            predecessors = list(G.predecessors(node_id))
+            if not predecessors:
+                return 0
+            return max(get_layer(pred, visited.copy()) for pred in predecessors) + 1
+        
+        for node in G.nodes():
+            layer = get_layer(node)
+            layers.setdefault(layer, []).append(node)
+        
+        # Calculate positions by layer
+        pos = {}
+        layer_height = 3.0
+        node_width = 4.0
+        
+        for layer_idx, nodes in layers.items():
+            y = layer_height * (len(layers) - 1 - layer_idx)
+            start_x = -(len(nodes) - 1) * node_width / 2
+            for i, node in enumerate(sorted(nodes)):
+                pos[node] = (start_x + i * node_width, y)
+        
+        return pos
+    
+    def create_dag_visualization(self, task_data, title="Robot Task Dependency Graph"):
+        """
+        Create a DAG visualization from task data and return the image path.
+        
+        Args:
+            task_data: Task data dictionary
+            title: Title for the graph
+        
+        Returns:
+            str: Path to the generated image file
+        """
+        try:
+            # Create graph
+            G = self.create_dag_from_tasks(task_data)
+            if not G or len(G.nodes()) == 0:
+                logger.warning("No tasks found or invalid graph structure")
+                return None
+            
+            # Calculate layout
+            pos = self.calculate_layout(G)
+            
+            # Create figure
+            fig, ax = plt.subplots(1, 1, figsize=(max(12, len(G.nodes()) * 2), 8))
+            
+            # Draw edges with arrows
+            nx.draw_networkx_edges(G, pos, 
+                                  edge_color='#2E86AB',
+                                  arrows=True,
+                                  arrowsize=20,
+                                  arrowstyle='->',
+                                  width=2,
+                                  alpha=0.8,
+                                  connectionstyle="arc3,rad=0.1")
+            
+            # Color nodes based on their position in the graph
+            node_colors = []
+            for node in G.nodes():
+                if G.in_degree(node) == 0:  # Start nodes
+                    node_colors.append('#F24236')
+                elif G.out_degree(node) == 0:  # End nodes
+                    node_colors.append('#A23B72')
+                else:  # Intermediate nodes
+                    node_colors.append('#F18F01')
+            
+            # Draw nodes
+            nx.draw_networkx_nodes(G, pos,
+                                  node_color=node_colors,
+                                  node_size=3500,
+                                  alpha=0.9,
+                                  edgecolors='black',
+                                  linewidths=2)
+            
+            # Label nodes with task IDs
+            node_labels = {node: f"T{node}" for node in G.nodes()}
+            nx.draw_networkx_labels(G, pos, node_labels,
+                                   font_size=18,
+                                   font_weight='bold',
+                                   font_color='white')
+            
+            # Add detailed info text boxes for each task
+            for i, node in enumerate(G.nodes()):
+                x, y = pos[node]
+                function_name = G.nodes[node]['function']
+                robots = G.nodes[node]['robots']
+                objects = G.nodes[node]['objects']
+                
+                # Create info text content
+                info_text = f"Task {node}: {function_name.replace('_', ' ').title()}\n"
+                if robots:
+                    robot_text = ", ".join([r.replace('robot_', '').replace('_', ' ').title() for r in robots])
+                    info_text += f"Robots: {robot_text}\n"
+                if objects:
+                    object_text = ", ".join(objects)
+                    info_text += f"Objects: {object_text}"
+                
+                # Calculate offset based on node position to avoid overlaps
+                offset_x = 2.2 if i % 2 == 0 else -2.2
+                offset_y = 0.5 if i % 4 < 2 else -0.5
+                
+                # Choose alignment based on offset direction
+                h_align = 'left' if offset_x > 0 else 'right'
+                
+                # Draw text box
+                bbox_props = dict(boxstyle="round,pad=0.4", 
+                                 facecolor='white', 
+                                 edgecolor='gray',
+                                 alpha=0.95,
+                                 linewidth=1)
+                
+                ax.text(x + offset_x, y + offset_y, info_text,
+                       bbox=bbox_props,
+                       fontsize=12,
+                       verticalalignment='center',
+                       horizontalalignment=h_align,
+                       weight='bold')
+                
+                # Draw dashed connector line from node to text box
+                ax.plot([x, x + offset_x], [y, y + offset_y], 
+                        linestyle='--', color='gray', alpha=0.6, linewidth=1)
+            
+            # Expand axis limits to fit everything
+            x_vals = [coord[0] for coord in pos.values()]
+            y_vals = [coord[1] for coord in pos.values()]
+            ax.set_xlim(min(x_vals) - 4.0, max(x_vals) + 4.0)
+            ax.set_ylim(min(y_vals) - 2.0, max(y_vals) + 2.0)
+            
+            # Set overall figure properties
+            ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
+            ax.set_aspect('equal')
+            ax.margins(0.2)
+            ax.axis('off')
+            
+            # Add legend for node types - Hidden to avoid covering content
+            # legend_elements = [
+            #     plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='#F24236', 
+            #                markersize=10, label='Start Tasks', markeredgecolor='black'),
+            #     plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='#A23B72', 
+            #                markersize=10, label='End Tasks', markeredgecolor='black'),
+            #     plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='#F18F01', 
+            #                markersize=10, label='Intermediate Tasks', markeredgecolor='black'),
+            #     plt.Line2D([0], [0], color='#2E86AB', linewidth=2, label='Dependencies')
+            # ]
+            # ax.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(1.05, 1.05))
+            
+            # Adjust layout and save
+            plt.tight_layout()
+            
+            # Create temporary file for saving the image
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            temp_dir = tempfile.gettempdir()
+            image_path = os.path.join(temp_dir, f'dag_visualization_{timestamp}.png')
+            
+            plt.savefig(image_path, dpi=400, bbox_inches='tight', 
+                        pad_inches=0.1, facecolor='white', edgecolor='none')
+            plt.close(fig)  # Close figure to free memory
+            
+            logger.info(f"DAG visualization saved to: {image_path}")
+            return image_path
+            
+        except Exception as e:
+            logger.error(f"Error creating DAG visualization: {e}")
+            return None
+    
+    def create_simplified_dag_visualization(self, task_data, title="Robot Task Graph"):
+        """
+        Create a simplified DAG visualization suitable for smaller displays.
+        
+        Args:
+            task_data: Task data dictionary
+            title: Title for the graph
+        
+        Returns:
+            str: Path to the generated image file
+        """
+        try:
+            # Create graph
+            G = self.create_dag_from_tasks(task_data)
+            if not G or len(G.nodes()) == 0:
+                logger.warning("No tasks found or invalid graph structure")
+                return None
+            
+            # Calculate layout
+            pos = self.calculate_layout(G)
+            
+            # Create figure for simplified graph
+            fig, ax = plt.subplots(1, 1, figsize=(10, 6))
+            
+            # Draw edges
+            nx.draw_networkx_edges(G, pos, 
+                                  edge_color='black',
+                                  arrows=True,
+                                  arrowsize=15,
+                                  arrowstyle='->',
+                                  width=1.5)
+            
+            # Draw nodes
+            nx.draw_networkx_nodes(G, pos,
+                                  node_color='lightblue',
+                                  node_size=3000,
+                                  edgecolors='black',
+                                  linewidths=1.5)
+            
+            # Add node labels with simplified names
+            labels = {}
+            for node in G.nodes():
+                function_name = G.nodes[node]['function']
+                simplified_name = function_name.replace('_', ' ').title()
+                if len(simplified_name) > 15:
+                    simplified_name = simplified_name[:12] + "..."
+                labels[node] = f"T{node}\n{simplified_name}"
+            
+            nx.draw_networkx_labels(G, pos, labels,
+                                   font_size=11,
+                                   font_weight='bold')
+            
+            ax.set_title(title, fontsize=14, fontweight='bold')
+            ax.axis('off')
+            
+            # Adjust layout and save
+            plt.tight_layout()
+            
+            # Create temporary file for saving the image
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            temp_dir = tempfile.gettempdir()
+            image_path = os.path.join(temp_dir, f'simple_dag_{timestamp}.png')
+            
+            plt.savefig(image_path, dpi=400, bbox_inches='tight')
+            plt.close(fig)  # Close figure to free memory
+            
+            logger.info(f"Simplified DAG visualization saved to: {image_path}")
+            return image_path
+            
+        except Exception as e:
+            logger.error(f"Error creating simplified DAG visualization: {e}")
+            return None

json_processor.py

@@ -1,46 +1,46 @@
-import json
-import re
-import ast
-from loguru import logger
-
-class JsonProcessor:
-    def process_response(self, response):
-        try:
-            # Search for JSON string in the response
-            json_str_match = re.search(r'\{.*\}', response, re.DOTALL)
-            if json_str_match:
-                # Get the matched JSON string
-                json_str = json_str_match.group()
-                logger.debug(f"Full JSON string: {json_str}")  
-
-                # Try to parse as Python literal first, then convert to JSON
-                try:
-                    # First try to evaluate as Python literal
-                    python_obj = ast.literal_eval(json_str)
-                    # Convert to proper JSON
-                    response_json = json.loads(json.dumps(python_obj))
-                except (ValueError, SyntaxError):
-                    # Fall back to string replacement method
-                    # Replace escape characters and remove trailing commas
-                    json_str = json_str.replace("\\", "")
-                    json_str = json_str.replace(r'\\_', '_')
-                    json_str = re.sub(r',\s*}', '}', json_str)
-                    json_str = re.sub(r',\s*\]', ']', json_str)
-                    
-                    # Convert Python format to JSON format
-                    json_str = json_str.replace("'", '"')  # Single quotes to double quotes
-                    json_str = json_str.replace('None', 'null')  # Python None to JSON null
-
-                    # Parse the JSON string
-                    response_json = json.loads(json_str)
-                return response_json
-            else:
-                logger.error("No JSON string match found in response.")
-                return None
-
-        except json.JSONDecodeError as e:
-            logger.error(f"JSONDecodeError: {e}")
-        except Exception as e:
-            logger.error(f"Unexpected error: {e}")
-
-        return None
+import json
+import re
+import ast
+from loguru import logger
+
+class JsonProcessor:
+    def process_response(self, response):
+        try:
+            # Search for JSON string in the response
+            json_str_match = re.search(r'\{.*\}', response, re.DOTALL)
+            if json_str_match:
+                # Get the matched JSON string
+                json_str = json_str_match.group()
+                logger.debug(f"Full JSON string: {json_str}")  
+
+                # Try to parse as Python literal first, then convert to JSON
+                try:
+                    # First try to evaluate as Python literal
+                    python_obj = ast.literal_eval(json_str)
+                    # Convert to proper JSON
+                    response_json = json.loads(json.dumps(python_obj))
+                except (ValueError, SyntaxError):
+                    # Fall back to string replacement method
+                    # Replace escape characters and remove trailing commas
+                    json_str = json_str.replace("\\", "")
+                    json_str = json_str.replace(r'\\_', '_')
+                    json_str = re.sub(r',\s*}', '}', json_str)
+                    json_str = re.sub(r',\s*\]', ']', json_str)
+                    
+                    # Convert Python format to JSON format
+                    json_str = json_str.replace("'", '"')  # Single quotes to double quotes
+                    json_str = json_str.replace('None', 'null')  # Python None to JSON null
+
+                    # Parse the JSON string
+                    response_json = json.loads(json_str)
+                return response_json
+            else:
+                logger.error("No JSON string match found in response.")
+                return None
+
+        except json.JSONDecodeError as e:
+            logger.error(f"JSONDecodeError: {e}")
+        except Exception as e:
+            logger.error(f"Unexpected error: {e}")
+
+        return None

requirements.txt

@@ -10,3 +10,6 @@ sentencepiece
 protobuf
 spaces
 loguru
+matplotlib
+networkx
+numpy

test_dag_integration.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+"""
+Test script for DAG integration in DART-LLM-Multi-Model
+"""
+
+from dag_visualizer import DAGVisualizer
+from json_processor import JsonProcessor
+import json
+
+def test_dag_integration():
+    """Test the DAG integration with sample task data"""
+    print("Testing DAG integration...")
+    
+    # Sample response with task data (similar to what the model might generate)
+    sample_response = """
+    Based on your command, here are the robot tasks:
+    
+    {
+        "tasks": [
+            {
+                "task": "move_excavator_to_soil_area",
+                "instruction_function": {
+                    "name": "move_to_position",
+                    "robot_ids": ["robot_excavator_01"],
+                    "dependencies": [],
+                    "object_keywords": ["soil_area_1"]
+                }
+            },
+            {
+                "task": "excavate_soil",
+                "instruction_function": {
+                    "name": "excavate_material",
+                    "robot_ids": ["robot_excavator_01"],
+                    "dependencies": ["move_excavator_to_soil_area"],
+                    "object_keywords": ["soil"]
+                }
+            },
+            {
+                "task": "move_dump_truck",
+                "instruction_function": {
+                    "name": "move_to_position",
+                    "robot_ids": ["robot_dump_truck_01"],
+                    "dependencies": [],
+                    "object_keywords": ["soil_area_1"]
+                }
+            },
+            {
+                "task": "load_soil_to_truck",
+                "instruction_function": {
+                    "name": "transfer_material",
+                    "robot_ids": ["robot_excavator_01", "robot_dump_truck_01"],
+                    "dependencies": ["excavate_soil", "move_dump_truck"],
+                    "object_keywords": ["soil"]
+                }
+            },
+            {
+                "task": "transport_to_dump_site",
+                "instruction_function": {
+                    "name": "move_to_position",
+                    "robot_ids": ["robot_dump_truck_01"],
+                    "dependencies": ["load_soil_to_truck"],
+                    "object_keywords": ["dump_site"]
+                }
+            }
+        ]
+    }
+    """
+    
+    # Test JSON processing
+    processor = JsonProcessor()
+    print("1. Testing JSON processing...")
+    processed_json = processor.process_response(sample_response)
+    
+    if processed_json:
+        print("✓ JSON processing successful")
+        print(f"   Found {len(processed_json['tasks'])} tasks")
+    else:
+        print("✗ JSON processing failed")
+        return False
+    
+    # Test DAG visualization
+    print("2. Testing DAG visualization...")
+    visualizer = DAGVisualizer()
+    
+    try:
+        dag_image_path = visualizer.create_dag_visualization(
+            processed_json, 
+            title="Test Robot Task Dependency Graph"
+        )
+        
+        if dag_image_path:
+            print(f"✓ DAG visualization created: {dag_image_path}")
+            return True
+        else:
+            print("✗ DAG visualization failed")
+            return False
+            
+    except Exception as e:
+        print(f"✗ DAG visualization error: {e}")
+        return False
+
+def test_simplified_dag():
+    """Test simplified DAG visualization"""
+    print("\n3. Testing simplified DAG...")
+    
+    simple_task_data = {
+        "tasks": [
+            {
+                "task": "move_robot",
+                "instruction_function": {
+                    "name": "move_to_position",
+                    "robot_ids": ["robot_01"],
+                    "dependencies": [],
+                    "object_keywords": ["target_area"]
+                }
+            },
+            {
+                "task": "perform_operation",
+                "instruction_function": {
+                    "name": "excavate",
+                    "robot_ids": ["robot_01"],
+                    "dependencies": ["move_robot"],
+                    "object_keywords": ["soil"]
+                }
+            }
+        ]
+    }
+    
+    visualizer = DAGVisualizer()
+    
+    try:
+        dag_image_path = visualizer.create_simplified_dag_visualization(
+            simple_task_data,
+            title="Simplified Test DAG"
+        )
+        
+        if dag_image_path:
+            print(f"✓ Simplified DAG visualization created: {dag_image_path}")
+            return True
+        else:
+            print("✗ Simplified DAG visualization failed")
+            return False
+            
+    except Exception as e:
+        print(f"✗ Simplified DAG visualization error: {e}")
+        return False
+
+def main():
+    """Run all tests"""
+    print("=" * 60)
+    print("DART-LLM-Multi-Model DAG Integration Test")
+    print("=" * 60)
+    
+    success_count = 0
+    total_tests = 2
+    
+    if test_dag_integration():
+        success_count += 1
+    
+    if test_simplified_dag():
+        success_count += 1
+    
+    print("\n" + "=" * 60)
+    print(f"Test Results: {success_count}/{total_tests} passed")
+    
+    if success_count == total_tests:
+        print("🎉 All DAG integration tests passed!")
+        return True
+    else:
+        print("❌ Some tests failed!")
+        return False
+
+if __name__ == "__main__":
+    success = main()
+    exit(0 if success else 1)