Update app.py

app.py

@@ -71,14 +71,20 @@ async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], Asyn
     layer_outputs = {}
 
     async def process_message(message):
-        # Store intermediate messages in `layer_outputs`
-        if message['response_type'] == 'intermediate':
-            layer = message['metadata']['layer']
+        """
+        Process an individual message and return the content for streaming.
+        Returns an empty string if the message type is intermediate or invalid.
+        """
+        # Debug: Log the incoming message
+        st.write(f"Processing message: {message}")
+
+        if message.get('response_type') == 'intermediate':
+            layer = message['metadata'].get('layer', "Unknown")
             if layer not in layer_outputs:
                 layer_outputs[layer] = []
-            layer_outputs[layer].append(message['delta'])
+            layer_outputs[layer].append(message.get('delta', ""))
             return ""  # Intermediate messages don't contribute to the final response
-        elif message['response_type'] == 'final':
+        elif message.get('response_type') == 'final':
             # Final message processing
             for layer, outputs in layer_outputs.items():
                 st.write(f"Layer {layer}")
@@ -87,22 +93,25 @@ 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'] or ""  # Ensure no None value is returned
+            return message.get('delta', "")  # Ensure no None value is returned
         else:
-            return ""  # Default to empty string if no valid response type
+            # Debug: Log unexpected message types
+            st.write(f"Unexpected message type: {message}")
+            return ""  # Default to empty string for invalid types
 
     # Check if the input is an async or sync iterable
     if hasattr(messages, "__aiter__"):  # Asynchronous iterable
         async for message in messages:
             resolved_message = await process_message(message)
-            yield resolved_message
+            yield resolved_message or ""  # Yield empty string if None
     elif hasattr(messages, "__iter__"):  # Synchronous iterable
         for message in messages:
             resolved_message = process_message(message)  # Do not `await` sync messages
-            yield resolved_message
+            yield resolved_message or ""  # Yield empty string if None
     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'],
@@ -267,20 +276,28 @@ if query := st.chat_input("Ask a question"):
 
             # Get messages from MOAgent (supports async streaming)
             messages = moa_agent.chat(query, output_format='json')
-            
+
             try:
                 # Stream and display responses from `stream_or_async_response`
                 final_response = ""
                 async for response in stream_or_async_response(messages):
                     # Skip None or empty responses
-                    if response:
+                    if response is not None and response.strip():
                         final_response += response
                         message_placeholder.markdown(final_response)
+                    else:
+                        # Debug: Log skipped empty/None responses
+                        st.write(f"Skipped empty or None response: {response}")
 
                 # Save the final response to session state
-                st.session_state.messages.append({"role": "assistant", "content": final_response})
+                if final_response.strip():  # Only save non-empty responses
+                    st.session_state.messages.append({"role": "assistant", "content": final_response})
+                else:
+                    st.error("Received no valid response from the assistant.")
             except Exception as e:
                 st.error(f"Error processing response: {e}")
+                # Debug: Log the exception
+                st.write(f"Exception details: {e}")
 
     # Run the asynchronous handle_query function
     asyncio.run(handle_query())
@@ -288,6 +305,7 @@ if query := st.chat_input("Ask a question"):
 
 
 
+
 # Add acknowledgment at the bottom
 st.markdown("---")
 st.markdown("""