Update app.py

app.py

@@ -66,7 +66,6 @@ async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], Asyn
     """
     Handles both synchronous and asynchronous message streams.
     Processes each response chunk and manages intermediate outputs.
-
     :param messages: Union[Iterable, AsyncIterable]
     """
     layer_outputs = {}
@@ -78,7 +77,8 @@ async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], Asyn
             if layer not in layer_outputs:
                 layer_outputs[layer] = []
             layer_outputs[layer].append(message['delta'])
-        else:
+            return ""  # Intermediate messages don't contribute to the final response
+        elif message['response_type'] == 'final':
             # Final message processing
             for layer, outputs in layer_outputs.items():
                 st.write(f"Layer {layer}")
@@ -87,7 +87,9 @@ async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], Asyn
                     with cols[i]:
                         st.expander(label=f"Agent {i + 1}", expanded=False).write(output)
             layer_outputs.clear()
-            return message['delta']
+            return message['delta'] or ""  # Ensure no None value is returned
+        else:
+            return ""  # Default to empty string if no valid response type
 
     # Check if the input is an async or sync iterable
     if hasattr(messages, "__aiter__"):  # Asynchronous iterable
@@ -101,8 +103,6 @@ async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], Asyn
     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'],
@@ -272,13 +272,10 @@ 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):
-                    # 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)
+                    # Skip None or empty responses
+                    if response:
+                        final_response += response
+                        message_placeholder.markdown(final_response)
 
                 # Save the final response to session state
                 st.session_state.messages.append({"role": "assistant", "content": final_response})
@@ -290,6 +287,7 @@ if query := st.chat_input("Ask a question"):
 
 
 
+
 # Add acknowledgment at the bottom
 st.markdown("---")
 st.markdown("""