Updated again

app.py

@@ -376,11 +376,14 @@ def collect_full_response(payload, model, session_key):
             print(f"Trying endpoint: {url}")
             print(f"Payload: {json.dumps(payload, indent=2)}")
             
+            # Check if we are requesting a stream or simple JSON
+            is_streaming_request = payload.get('stream', False)
+            
             with requests.post(
                 url, 
                 json=payload, 
                 headers=get_headers(), 
-                stream=True,
+                stream=is_streaming_request, # Use stream=True only if requested
                 timeout=120
             ) as response:
                 
@@ -400,52 +403,69 @@ def collect_full_response(payload, model, session_key):
                         }
                     }, response.status_code
                 
-                buffer = ""
-                debug_captured_chunks = []
-                
-                for chunk in response.iter_content(chunk_size=None, decode_unicode=True):
-                    if chunk:
-                        # Capture first few chunks for debugging
-                        if len(debug_captured_chunks) < 3:
-                            debug_captured_chunks.append(repr(chunk))
-                            
-                        buffer += chunk
+                # CASE 1: Non-Streaming Response (JSON)
+                if not is_streaming_request:
+                    try:
+                        data = response.json()
+                        # Extract content - logs show 'answer' field is used
+                        full_content = data.get('answer') or data.get('message') or data.get('content') or ""
                         
-                        while '\n' in buffer:
-                            line, buffer = buffer.split('\n', 1)
-                            line = line.strip()
-                            
-                            if not line:
-                                continue
-                            
-                            if line.startswith('data: '):
-                                line = line[6:]
-                            
-                            if line == '[DONE]':
-                                continue
+                        # Extract Message ID if present to update session
+                        msg_id = data.get('message_id')
+                        if session_key in chat_sessions_cache and msg_id:
+                            chat_sessions_cache[session_key]['parent_message_id'] = msg_id
                             
-                            content, msg_id, packet_type = parse_onyx_stream_chunk(line)
-                            if msg_id:
-                                last_message_id = msg_id
-                            if packet_type == 'stop':
-                                break
-                            if content and packet_type in ['content', 'legacy', 'raw', 'error']:
-                                full_content += content
-                
-                if buffer.strip():
-                    if buffer.strip().startswith('data: '):
-                        buffer = buffer.strip()[6:]
-                    content, msg_id, packet_type = parse_onyx_stream_chunk(buffer.strip())
-                    if msg_id:
-                        last_message_id = msg_id
-                    if content and packet_type in ['content', 'legacy', 'raw', 'error']:
-                        full_content += content
-                
-                # Update session (only if we got a valid 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
+                        
+                    except json.JSONDecodeError:
+                        print("Failed to decode JSON response")
+                         # Fallback to text if JSON fails
+                        full_content = response.text
+                        break
                 
-                break
+                # CASE 2: Streaming Response (Original Logic)
+                else:
+                    buffer = ""
+                    for chunk in response.iter_content(chunk_size=None, decode_unicode=True):
+                        if chunk:
+                            buffer += chunk
+                            
+                            while '\n' in buffer:
+                                line, buffer = buffer.split('\n', 1)
+                                line = line.strip()
+                                
+                                if not line:
+                                    continue
+                                
+                                if line.startswith('data: '):
+                                    line = line[6:]
+                                
+                                if line == '[DONE]':
+                                    continue
+                                
+                                content, msg_id, packet_type = parse_onyx_stream_chunk(line)
+                                if msg_id:
+                                    last_message_id = msg_id
+                                if packet_type == 'stop':
+                                    break
+                                if content and packet_type in ['content', 'legacy', 'raw', 'error']:
+                                    full_content += content
+                    
+                    if buffer.strip():
+                        if buffer.strip().startswith('data: '):
+                            buffer = buffer.strip()[6:]
+                        content, msg_id, packet_type = parse_onyx_stream_chunk(buffer.strip())
+                        if msg_id:
+                            last_message_id = msg_id
+                        if content and packet_type in ['content', 'legacy', 'raw', 'error']:
+                            full_content += content
+                    
+                    # Update session (only if we got a valid 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
+
                 
         except requests.exceptions.RequestException as e:
             print(f"Request error: {e}")