Update app.py

app.py

@@ -89,17 +89,20 @@ async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], Asyn
             layer_outputs.clear()
             return message['delta']
 
-    # Determine if messages is async or sync and iterate accordingly
+    # Check if the input is an async or sync iterable
     if hasattr(messages, "__aiter__"):  # Asynchronous iterable
         async for message in messages:
-            yield await process_message(message)
+            resolved_message = await process_message(message)
+            yield resolved_message
     elif hasattr(messages, "__iter__"):  # Synchronous iterable
         for message in messages:
-            yield process_message(message)  # Do not `await` sync messages
+            resolved_message = process_message(message)  # Do not `await` sync messages
+            yield resolved_message
     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'],
@@ -269,7 +272,12 @@ if query := st.chat_input("Ask a question"):
                 # Stream and display responses from `stream_or_async_response`
                 final_response = ""
                 async for response in stream_or_async_response(messages):
-                    final_response += response  # Accumulate full response
+                    # Await the response if it's a coroutine
+                    if asyncio.iscoroutine(response):
+                        response = await response
+                    
+                    # Accumulate and display the response
+                    final_response += response
                     message_placeholder.markdown(final_response)
 
                 # Save the final response to session state
@@ -281,6 +289,7 @@ if query := st.chat_input("Ask a question"):
     asyncio.run(handle_query())
 
 
+
 # Add acknowledgment at the bottom
 st.markdown("---")
 st.markdown("""