Update app.py

app.py

@@ -64,27 +64,39 @@ layer_agent_config_rec = {
 
 def stream_response(messages: Iterable[ResponseChunk]):
     layer_outputs = {}
+    progress_bar = st.progress(0)
+    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)
+        
         if message['response_type'] == 'intermediate':
             layer = message['metadata']['layer']
             if layer not in layer_outputs:
                 layer_outputs[layer] = []
             layer_outputs[layer].append(message['delta'])
+            
+            # Real-time rendering for intermediate outputs
+            with st.container():
+                st.markdown(f"**Layer {layer} (In Progress)**")
+                for output in layer_outputs[layer]:
+                    st.markdown(f"- {output}")
+
         else:
-            # Display accumulated layer outputs
+            # Finalize and display accumulated layer outputs
             for layer, outputs in layer_outputs.items():
-                st.write(f"Layer {layer}")
-                cols = st.columns(len(outputs))
-                for i, output in enumerate(outputs):
-                    with cols[i]:
-                        st.expander(label=f"Agent {i+1}", expanded=False).write(output)
-            
-            # Clear layer outputs for the next iteration
-            layer_outputs = {}
+                st.markdown(f"### Layer {layer} Final Output")
+                for output in outputs:
+                    st.write(output)
+            layer_outputs = {}  # Reset for next layers
             
             # Yield the main agent's output
             yield message['delta']
 
+    progress_bar.empty()  # Clear progress bar once done
+
 def set_moa_agent(
     main_model: str = default_config['main_model'],
     cycles: int = default_config['cycles'],