Update app.py

app.py

@@ -178,9 +178,7 @@ def parse_onyx_stream_chunk(chunk_text):
         if not isinstance(data, dict):
             return None, None, None
         
-        # Handle first packet (message IDs)
-        if 'user_message_id' in data or 'reserved_assistant_message_id' in data:
-            return None, data.get('reserved_assistant_message_id'), 'message_ids'
+    
         
         # Handle new packet-based format
         if 'obj' in data:
@@ -255,6 +253,7 @@ def generate_openai_stream_chunk(content, model, chunk_id, finish_reason=None):
 
 
 def stream_onyx_response(payload, model, session_key):
+    final_message_id = None
     """Stream response from Onyx API in OpenAI SSE format"""
     chunk_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
     url = f"{ONYX_BASE_URL}/api/chat/send-message"
@@ -332,7 +331,8 @@ def stream_onyx_response(payload, model, session_key):
                             if msg_id:
                                 last_message_id = msg_id
                             if packet_type == 'stop':
-                                break
+    final_message_id = last_message_id
+    break
                             if content and packet_type in ['content', 'legacy', 'raw', 'error']:
                                 yield generate_openai_stream_chunk(content, model, chunk_id)
                 
@@ -348,7 +348,7 @@ def stream_onyx_response(payload, model, session_key):
                 
                 # Update session with last message ID
                 if session_key in chat_sessions_cache and last_message_id:
-                    chat_sessions_cache[session_key]['parent_message_id'] = last_message_id
+                    
                 
                 break  # Success, exit loop
                 
@@ -359,6 +359,8 @@ def stream_onyx_response(payload, model, session_key):
     # Send final chunk
     yield generate_openai_stream_chunk("", model, chunk_id, "stop")
     yield "data: [DONE]\n\n"
+if final_message_id:
+    chat_sessions_cache[session_key]['parent_message_id'] = final_message_id
 
 
 def collect_full_response(payload, model, session_key):