Update app.py

app.py

@@ -63,39 +63,39 @@ layer_agent_config_rec = {
 
 
 def stream_response(messages: Iterable[ResponseChunk]):
-    layer_outputs = {}
-    progress_bar = st.progress(0)
+    layer_outputs = {}  # Store outputs per layer
+    progress_bar = st.progress(0)  # Initialize progress bar
     total_steps = len(messages)  # Estimate total messages for progress tracking
     current_step = 0
 
     for message in messages:
         current_step += 1
-        progress_bar.progress(current_step / total_steps)
-        
+        progress_bar.progress(current_step / total_steps)  # Update progress bar
+
         if message['response_type'] == 'intermediate':
-            layer = message['metadata']['layer']
+            layer = message['metadata']['layer']  # Identify the layer
             if layer not in layer_outputs:
-                layer_outputs[layer] = []
-            layer_outputs[layer].append(message['delta'])
+                layer_outputs[layer] = []  # Initialize if layer not seen yet
+            layer_outputs[layer].append(message['delta'])  # Add intermediate delta
             
-            # Real-time rendering for intermediate outputs
-            with st.container():
-                st.markdown(f"**Layer {layer} (In Progress)**")
+            # Render intermediate output in real-time
+            with st.container():  # Create a dynamic container for rendering
+                st.markdown(f"### Layer {layer} (In Progress)")
                 for output in layer_outputs[layer]:
-                    st.markdown(f"- {output}")
+                    st.markdown(f"- {output}")  # Render each intermediate result
 
         else:
-            # Finalize and display accumulated layer outputs
+            # Final rendering of outputs after processing completes for this layer
             for layer, outputs in layer_outputs.items():
                 st.markdown(f"### Layer {layer} Final Output")
                 for output in outputs:
-                    st.write(output)
-            layer_outputs = {}  # Reset for next layers
+                    st.markdown(f"- {output}")
             
-            # Yield the main agent's output
-            yield message['delta']
+            layer_outputs = {}  # Reset layer outputs for the next iteration
+            yield message['delta']  # Yield the final processed result
+
+    progress_bar.empty()  # Clear progress bar after all processing is complete
 
-    progress_bar.empty()  # Clear progress bar once done
 
 def set_moa_agent(
     main_model: str = default_config['main_model'],