Update app.py

app.py

@@ -77,19 +77,20 @@ async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], Asyn
                 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)
+                        st.expander(label=f"Agent {i + 1}", expanded=False).write(output)
 
             layer_outputs.clear()
             yield message['delta']
 
-    if isinstance(messages, AsyncIterable):
-        # Process asynchronous messages
+    # Distinguish between async and sync iterables
+    if hasattr(messages, "__aiter__"):  # Asynchronous iterable
         async for message in messages:
-            await process_message(message)
-    else:
-        # Process synchronous messages
+            yield await process_message(message)
+    elif hasattr(messages, "__iter__"):  # Synchronous iterable
         for message in messages:
-            await process_message(message)
+            yield await process_message(message)
+    else:
+        raise TypeError("'messages' must be an Iterable or AsyncIterable.")
 
 # Set up the MOAgent
 def set_moa_agent(
@@ -242,23 +243,33 @@ for message in st.session_state.messages:
 
 if query := st.chat_input("Ask a question"):
     async def handle_query():
+        # Append the user's message to the session state
         st.session_state.messages.append({"role": "user", "content": query})
         with st.chat_message("user"):
             st.write(query)
 
+        # Access the MOAgent instance
         moa_agent: MOAgent = st.session_state.moa_agent
 
         with st.chat_message("assistant"):
             message_placeholder = st.empty()
+
+            # Get messages from MOAgent (supports async streaming)
             messages = moa_agent.chat(query, output_format='json')
-            async for response in stream_or_async_response(messages):
-                message_placeholder.markdown(response)
+            
+            try:
+                # Process each response chunk
+                async for response in stream_or_async_response(messages):
+                    message_placeholder.markdown(response)
 
-        st.session_state.messages.append({"role": "assistant", "content": response})
+                # Save the final response to session state
+                st.session_state.messages.append({"role": "assistant", "content": 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("""