Update app.py

app.py

@@ -63,35 +63,43 @@ layer_agent_config_rec = {
 
 # Unified streaming function to handle async and sync responses
 async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], AsyncIterable[ResponseChunk]]):
+    """
+    Handles both synchronous and asynchronous message streams.
+    Processes each response chunk and manages intermediate outputs.
+
+    :param messages: Union[Iterable, AsyncIterable]
+    """
     layer_outputs = {}
 
     async def process_message(message):
+        # Store intermediate messages in `layer_outputs`
         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:
+            # Final message processing
             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)
-
             layer_outputs.clear()
-            yield message['delta']
+            return message['delta']
 
-    # Distinguish between async and sync iterables
+    # Determine if messages is async or sync and iterate accordingly
     if hasattr(messages, "__aiter__"):  # Asynchronous iterable
         async for message in messages:
             yield await process_message(message)
     elif hasattr(messages, "__iter__"):  # Synchronous iterable
         for message in messages:
-            yield await process_message(message)
+            yield process_message(message)  # Do not `await` sync messages
     else:
         raise TypeError("'messages' must be an Iterable or AsyncIterable.")
 
+
 # Set up the MOAgent
 def set_moa_agent(
     main_model: str = default_config['main_model'],
@@ -258,18 +266,21 @@ if query := st.chat_input("Ask a question"):
             messages = moa_agent.chat(query, output_format='json')
             
             try:
-                # Process each response chunk
+                # Stream and display responses from `stream_or_async_response`
+                final_response = ""
                 async for response in stream_or_async_response(messages):
-                    message_placeholder.markdown(response)
+                    final_response += response  # Accumulate full response
+                    message_placeholder.markdown(final_response)
 
                 # Save the final response to session state
-                st.session_state.messages.append({"role": "assistant", "content": response})
+                st.session_state.messages.append({"role": "assistant", "content": final_response})
             except Exception as e:
                 st.error(f"Error processing response: {e}")
 
     # Run the asynchronous handle_query function
     asyncio.run(handle_query())
 
+
 # Add acknowledgment at the bottom
 st.markdown("---")
 st.markdown("""