Update app.py

app.py

@@ -64,22 +64,16 @@ layer_agent_config_rec = {
 }
 
 
-def stream_response(messages: Iterable[ResponseChunk]):
+async def async_stream_response(messages: Union[Iterable[ResponseChunk], AsyncIterable[ResponseChunk]]):
     layer_outputs = {}
-    lock = threading.Lock()
-    
-    def safe_append(layer, delta):
-        with lock:
-            if layer not in layer_outputs:
-                layer_outputs[layer] = []
-            layer_outputs[layer].append(delta)
 
-    for message in messages:
+    async def process_message(message):
         if message['response_type'] == 'intermediate':
             layer = message['metadata']['layer']
-            safe_append(layer, message['delta'])
+            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.write(f"Layer {layer}")
                 cols = st.columns(len(outputs))
@@ -87,12 +81,18 @@ def stream_response(messages: Iterable[ResponseChunk]):
                     with cols[i]:
                         st.expander(label=f"Agent {i+1}", expanded=False).write(output)
             
-            # Clear layer outputs for the next iteration
-            layer_outputs = {}
-            
-            # Yield the main agent's output
+            layer_outputs.clear()
             yield message['delta']
 
+    if isinstance(messages, AsyncIterable):
+        # Process asynchronous messages
+        async for message in messages:
+            await process_message(message)
+    else:
+        # Process synchronous messages
+        for message in messages:
+            await process_message(message)
+
 def set_moa_agent(
     main_model: str = default_config['main_model'],
     cycles: int = default_config['cycles'],