Updated app.py

app.py

@@ -79,8 +79,41 @@ def parse_model_string(model):
         return "openai", model
 
 
-def build_onyx_payload(messages, model_provider, model_version, temperature, chat_session_id, parent_message_id=None):
-    """Convert OpenAI format to Onyx payload"""
+# Known provider name mappings
+# Update these based on what's configured in Onyx Cloud admin panel
+PROVIDER_ALIASES = {
+    "openai": "openai",
+    "anthropic": "anthropic",
+    "google": "google",
+    "azure": "azure",
+    "bedrock": "bedrock",
+    "cohere": "cohere",
+    "mistral": "mistral",
+    # Add more aliases as needed
+}
+
+
+def normalize_provider_name(provider):
+    """
+    Normalize provider name to match Onyx configuration.
+    Handles case sensitivity and common aliases.
+    """
+    provider_lower = provider.lower().strip()
+    return PROVIDER_ALIASES.get(provider_lower, provider_lower)
+
+
+def build_onyx_payload(messages, model_provider, model_version, temperature, chat_session_id, parent_message_id=None, stream=True):
+    """Convert OpenAI format to Onyx payload
+    
+    Args:
+        messages: List of OpenAI format messages
+        model_provider: LLM provider name (e.g., 'openai', 'anthropic')
+        model_version: Model version (e.g., 'gpt-4', 'claude-3-opus-20240229')
+        temperature: Temperature setting for generation
+        chat_session_id: Onyx chat session ID
+        parent_message_id: Optional parent message ID for threading
+        stream: Whether to stream the response (default True)
+    """
     
     # Extract the last user message
     last_user_message = ""
@@ -109,7 +142,7 @@ def build_onyx_payload(messages, model_provider, model_version, temperature, cha
         "message": full_message,
         "chat_session_id": chat_session_id,
         "parent_message_id": parent_message_id if parent_message_id else None,
-        "stream": True,
+        "stream": stream,  # Now respects caller's preference
         "llm_override": {
             "model_provider": model_provider,
             "model_version": model_version,
@@ -126,38 +159,83 @@ def build_onyx_payload(messages, model_provider, model_version, temperature, cha
 
 
 def parse_onyx_stream_chunk(chunk_text):
-    """Parse a chunk from Onyx stream and extract the text content"""
+    """Parse a chunk from Onyx stream and extract the text content.
+    
+    New Onyx API uses packet-based format:
+    - First packet: {"user_message_id": int, "reserved_assistant_message_id": int}
+    - Content packets: {"ind": int, "obj": {"type": "message_delta", "content": "..."}}
+    - Stop packet: {"ind": int, "obj": {"type": "stop"}}
+    
+    Returns:
+        tuple: (content, message_id, packet_type)
+    """
     if not chunk_text or not chunk_text.strip():
-        return None, None
+        return None, None, None
     
     try:
         data = json.loads(chunk_text)
         
-        if isinstance(data, dict):
-            # Extract message ID for tracking conversation
-            message_id = data.get('message_id')
+        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:
+            obj = data['obj']
+            packet_type = obj.get('type', '')
             
-            # Check for different content fields
-            if 'answer_piece' in data:
-                return data['answer_piece'], message_id
-            elif 'text' in data:
-                return data['text'], message_id
-            elif 'content' in data:
-                return data['content'], message_id
-            elif 'message' in data and isinstance(data['message'], str):
-                return data['message'], message_id
-            elif 'error' in data:
-                return f"[Error: {data['error']}]", message_id
-                
-        elif isinstance(data, str):
-            return data, None
+            if packet_type == 'message_delta':
+                # This is the actual content!
+                content = obj.get('content', '')
+                return content, None, 'content'
+            
+            elif packet_type == 'message_start':
+                # Contains final_documents, not content
+                return None, None, 'message_start'
+            
+            elif packet_type == 'stop':
+                # End of stream
+                return None, None, 'stop'
+            
+            elif packet_type == 'error':
+                error_msg = obj.get('message', obj.get('error', 'Unknown error'))
+                return f"[Error: {error_msg}]", None, 'error'
+            
+            elif packet_type == 'citation_delta':
+                # Citation info, not content
+                return None, None, 'citation'
+            
+            elif packet_type in ['reasoning_start', 'reasoning_delta', 'reasoning_done']:
+                # Reasoning packets
+                return None, None, 'reasoning'
             
+            else:
+                # Other packet types (search, tools, etc.)
+                return None, None, packet_type
+        
+        # FALLBACK: Old format support (for backward compatibility)
+        message_id = data.get('message_id')
+        
+        if 'answer_piece' in data:
+            return data['answer_piece'], message_id, 'legacy'
+        elif 'text' in data:
+            return data['text'], message_id, 'legacy'
+        elif 'content' in data and isinstance(data['content'], str):
+            return data['content'], message_id, 'legacy'
+        elif 'error' in data:
+            return f"[Error: {data['error']}]", message_id, 'error'
+        
+        return None, None, None
+        
     except json.JSONDecodeError:
         # Not JSON, might be raw text
         if chunk_text.strip() and not chunk_text.strip().startswith('{'):
-            return chunk_text.strip(), None
+            return chunk_text.strip(), None, 'raw'
     
-    return None, None
+    return None, None, None
 
 
 def generate_openai_stream_chunk(content, model, chunk_id, finish_reason=None):
@@ -183,8 +261,8 @@ def stream_onyx_response(payload, model, session_key):
     
     # Try alternate endpoints if needed
     endpoints = [
-        f"{ONYX_BASE_URL}/api/chat/send-message",
-        f"{ONYX_BASE_URL}/api/chat/send-chat-message",
+        f"{ONYX_BASE_URL}/api/chat/send-chat-message",  # Primary (new)
+        f"{ONYX_BASE_URL}/api/chat/send-message",       # Fallback (deprecated)
     ]
     
     # Send initial chunk with role
@@ -250,20 +328,22 @@ def stream_onyx_response(payload, model, session_key):
                             if line == '[DONE]':
                                 continue
                             
-                            content, msg_id = parse_onyx_stream_chunk(line)
+                            content, msg_id, packet_type = parse_onyx_stream_chunk(line)
                             if msg_id:
                                 last_message_id = msg_id
-                            if content:
+                            if packet_type == 'stop':
+                                break
+                            if content and packet_type in ['content', 'legacy', 'raw', 'error']:
                                 yield generate_openai_stream_chunk(content, model, chunk_id)
                 
                 # Process remaining buffer
                 if buffer.strip():
                     if buffer.strip().startswith('data: '):
                         buffer = buffer.strip()[6:]
-                    content, msg_id = parse_onyx_stream_chunk(buffer.strip())
+                    content, msg_id, packet_type = parse_onyx_stream_chunk(buffer.strip())
                     if msg_id:
                         last_message_id = msg_id
-                    if content:
+                    if content and packet_type in ['content', 'legacy', 'raw', 'error']:
                         yield generate_openai_stream_chunk(content, model, chunk_id)
                 
                 # Update session with last message ID
@@ -287,8 +367,8 @@ def collect_full_response(payload, model, session_key):
     last_message_id = None
     
     endpoints = [
-        f"{ONYX_BASE_URL}/api/chat/send-message",
-        f"{ONYX_BASE_URL}/api/chat/send-chat-message",
+        f"{ONYX_BASE_URL}/api/chat/send-chat-message",  # Primary (new)
+        f"{ONYX_BASE_URL}/api/chat/send-message",       # Fallback (deprecated)
     ]
     
     for url in endpoints:
@@ -338,19 +418,21 @@ def collect_full_response(payload, model, session_key):
                             if line == '[DONE]':
                                 continue
                             
-                            content, msg_id = parse_onyx_stream_chunk(line)
+                            content, msg_id, packet_type = parse_onyx_stream_chunk(line)
                             if msg_id:
                                 last_message_id = msg_id
-                            if content:
+                            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 = parse_onyx_stream_chunk(buffer.strip())
+                    content, msg_id, packet_type = parse_onyx_stream_chunk(buffer.strip())
                     if msg_id:
                         last_message_id = msg_id
-                    if content:
+                    if content and packet_type in ['content', 'legacy', 'raw', 'error']:
                         full_content += content
                 
                 # Update session
@@ -424,8 +506,9 @@ def chat_completions():
             }
         }), 400
     
-    # Parse model string
+    # Parse model string and normalize provider name
     model_provider, model_version = parse_model_string(model)
+    model_provider = normalize_provider_name(model_provider)
     print(f"Model provider: {model_provider}, version: {model_version}")
     
     # Get or create chat session
@@ -446,7 +529,8 @@ def chat_completions():
         model_version=model_version,
         temperature=temperature,
         chat_session_id=session_info['session_id'],
-        parent_message_id=session_info.get('parent_message_id')
+        parent_message_id=session_info.get('parent_message_id'),
+        stream=stream  # Pass client's streaming preference
     )
     
     if stream: