add node name to response

src/agent.py

@@ -74,7 +74,7 @@ class AudioAgent:
             stream_mode="messages"
         ):
             if msg.content:
-                yield msg.content
+                yield msg.content, metadata["langgraph_node"]
 
 async def main():
     agent = AudioAgent()
@@ -83,7 +83,7 @@ async def main():
     print("→", reply)
 
     # streaming example
-    async for msg in agent.stream_chat("Explain how audio normalization works."):
+    async for msg, node in agent.stream_chat("Explain how audio normalization works."):
         print(msg, end="", flush=True)
 
 if __name__ == "__main__":

src/ui.py

@@ -19,7 +19,7 @@ def user_input(user_message, history):
 
 async def bot_response(history):
     """
-    Generate bot response with streaming
+    Generate bot response with streaming, organizing content by graph nodes
     """
     if not history or history[-1]["role"] != "user":
         return
@@ -35,17 +35,39 @@ async def bot_response(history):
         history.append({"role": "assistant", "content": ""})
         yield history
         
+        # Track current node and organize content by nodes
+        current_content = ""
+        current_node = None
+        nodes_content = {}
+        
         # Stream the response
-        async for chunk in agent.stream_chat(user_message):
-            history[-1]["content"] += chunk
+        async for chunk, node_name in agent.stream_chat(user_message):
+            # If we encounter a new node, update the display structure
+            if node_name != current_node:
+                current_node = node_name
+                if node_name not in nodes_content:
+                    nodes_content[node_name] = ""
+            
+            # Add chunk to the current node's content
+            if chunk:
+                nodes_content[node_name] += chunk
+            
+            # Build the formatted content with node headers
+            formatted_content = ""
+            for node, content in nodes_content.items():
+                if content.strip():  # Only show nodes that have content
+                    formatted_content += f"**🔧 {node}**\n\n{content}\n\n"
+            
+            # Update the chat history
+            history[-1]["content"] = formatted_content.rstrip()
             yield history
             
     except Exception as e:
         # Update the last message with error or add error message
         if history and history[-1]["role"] == "assistant":
-            history[-1]["content"] = f"Error: {str(e)}"
+            history[-1]["content"] = f"❌ **Error**: {str(e)}"
         else:
-            history.append({"role": "assistant", "content": f"Error: {str(e)}"})
+            history.append({"role": "assistant", "content": f"❌ **Error**: {str(e)}"})
         yield history
 
 def bot_response_sync(history):