Update app.py

app.py

@@ -64,35 +64,28 @@ layer_agent_config_rec = {
 
 def stream_response(messages: Iterable[ResponseChunk]):
     layer_outputs = {}  # Store outputs per layer
-    progress_bar = st.progress(0)  # Initialize progress bar
-    current_step = 0  # Start step counter
 
     for message in messages:
-        current_step += 1
-        progress_bar.progress(current_step % 100)  # Use a modular percentage for progress
-
         if message['response_type'] == 'intermediate':
-            layer = message['metadata']['layer']  # Identify the layer
+            layer = message['metadata']['layer']
             if layer not in layer_outputs:
-                layer_outputs[layer] = []  # Initialize if layer not seen yet
-            layer_outputs[layer].append(message['delta'])  # Add intermediate delta
-            
-            # Render intermediate output in real-time
-            with st.container():
-                st.markdown(f"### Layer {layer} (In Progress)")
-                for output in layer_outputs[layer]:
-                    st.markdown(f"- {output}")
+                layer_outputs[layer] = []
+            layer_outputs[layer].append(message['delta'])
+
         else:
-            # Final rendering of outputs
+            # Display accumulated layer outputs
             for layer, outputs in layer_outputs.items():
-                st.markdown(f"### Layer {layer} Final Output")
-                for output in outputs:
-                    st.markdown(f"- {output}")
+                st.markdown(f"### Layer {layer}")
+                for i, output in enumerate(outputs):
+                    with st.expander(label=f"Agent {i+1} Output", expanded=False):
+                        st.write(output)  # Display output inside the dropdown
 
-            layer_outputs = {}  # Reset layer outputs for next iteration
-            yield message['delta']  # Yield final processed result
+            # Clear layer outputs for the next iteration
+            layer_outputs = {}
+            
+            # Yield the main agent's output
+            yield message['delta']
 
-    progress_bar.empty()  # Clear progress bar after all processing