Update app.py

app.py

@@ -63,23 +63,22 @@ layer_agent_config_rec = {
 
 
 def stream_response(messages: Iterable[ResponseChunk]):
-    layer_outputs = {}  # Store outputs per layer
-
+    layer_outputs = {}
     for message in messages:
         if message['response_type'] == 'intermediate':
             layer = message['metadata']['layer']
             if layer not in layer_outputs:
                 layer_outputs[layer] = []
             layer_outputs[layer].append(message['delta'])
-
         else:
             # Display accumulated layer outputs
             for layer, outputs in layer_outputs.items():
-                st.markdown(f"### Layer {layer}")
+                st.write(f"Layer {layer}")
+                cols = st.columns(len(outputs))
                 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
-
+                    with cols[i]:
+                        st.expander(label=f"Agent {i+1}", expanded=False).write(output)
+            
             # Clear layer outputs for the next iteration
             layer_outputs = {}