Update app.py

app.py

@@ -8,19 +8,78 @@ import os
 app = Flask(__name__)
 
 # Configuration
-ONYX_API_URL = "https://cloud.onyx.app/api/chat/send-chat-message"
+ONYX_BASE_URL = os.environ.get("ONYX_BASE_URL", "https://cloud.onyx.app")
 ONYX_API_TOKEN = os.environ.get("ONYX_SECRET", "<your-token-here>")
 
-# Store chat sessions for context
-chat_sessions = {}
+# Store chat sessions
+chat_sessions_cache = {}
 
 
-def create_chat_session_id():
-    """Generate a new chat session UUID"""
-    return str(uuid.uuid4())
+def get_headers():
+    """Get authorization headers"""
+    return {
+        "Authorization": f"Bearer {ONYX_API_TOKEN}",
+        "Content-Type": "application/json"
+    }
+
+
+def create_chat_session(persona_id=0):
+    """Create a new chat session in Onyx"""
+    url = f"{ONYX_BASE_URL}/api/chat/create-chat-session"
+    
+    payload = {
+        "persona_id": persona_id,
+        "description": "OpenAI Compatible API Session"
+    }
+    
+    try:
+        response = requests.post(url, json=payload, headers=get_headers(), timeout=30)
+        
+        if response.status_code == 200:
+            data = response.json()
+            # The response might have chat_session_id directly or nested
+            session_id = data.get('chat_session_id') or data.get('id') or data
+            print(f"Created chat session: {session_id}")
+            return str(session_id)
+        else:
+            print(f"Failed to create chat session: {response.status_code} - {response.text}")
+            return None
+    except Exception as e:
+        print(f"Error creating chat session: {e}")
+        return None
+
+
+def get_or_create_session(session_key="default", persona_id=0):
+    """Get existing session or create new one"""
+    if session_key not in chat_sessions_cache:
+        session_id = create_chat_session(persona_id)
+        if session_id:
+            chat_sessions_cache[session_key] = {
+                "session_id": session_id,
+                "parent_message_id": None
+            }
+    return chat_sessions_cache.get(session_key)
 
 
-def build_onyx_payload(messages, model_provider, model_version, temperature, chat_session_id=None):
+def parse_model_string(model):
+    """
+    Parse model string in format 'provider/model_version'
+    Examples: 
+        - 'openai/gpt-4' -> ('openai', 'gpt-4')
+        - 'anthropic/claude-3-opus' -> ('anthropic', 'claude-3-opus')
+        - 'gpt-4' -> ('openai', 'gpt-4')
+    """
+    if '/' in model:
+        parts = model.split('/', 1)
+        return parts[0], parts[1]
+    elif ':' in model:
+        parts = model.split(':', 1)
+        return parts[0], parts[1]
+    else:
+        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"""
     
     # Extract the last user message
@@ -29,7 +88,6 @@ def build_onyx_payload(messages, model_provider, model_version, temperature, cha
         if msg.get('role') == 'user':
             content = msg.get('content', '')
             if isinstance(content, list):
-                # Handle content array format (for vision models, etc.)
                 text_parts = [p.get('text', '') for p in content if p.get('type') == 'text']
                 last_user_message = ' '.join(text_parts)
             else:
@@ -45,85 +103,61 @@ def build_onyx_payload(messages, model_provider, model_version, temperature, cha
     # Prepend system prompt to message if exists
     full_message = last_user_message
     if system_prompt:
-        full_message = f"[System Instructions: {system_prompt.strip()}]\n\n{last_user_message}"
+        full_message = f"[System: {system_prompt.strip()}]\n\n{last_user_message}"
     
     payload = {
         "message": full_message,
+        "chat_session_id": chat_session_id,
+        "parent_message_id": parent_message_id if parent_message_id else None,
+        "stream": True,
         "llm_override": {
             "model_provider": model_provider,
             "model_version": model_version,
             "temperature": temperature
         },
-        "allowed_tool_ids": [],
         "file_descriptors": [],
-        "internal_search_filters": {
-            "source_type": [],
-            "document_set": [],
-            "tags": []
-        },
-        "deep_research": False,
-        "origin": "unset",
-        "parent_message_id": -1,
-        "chat_session_id": chat_session_id or create_chat_session_id(),
-        "chat_session_info": {
-            "persona_id": 0,
-            "description": "OpenAI Compatible API Session",
-            "project_id": 0
-        },
-        "stream": True,  # Always stream internally
         "include_citations": False
     }
     
+    # Remove None values
+    payload = {k: v for k, v in payload.items() if v is not None}
+    
     return payload
 
 
-def parse_model_string(model):
-    """
-    Parse model string in format 'provider/model_version' or 'provider:model_version'
-    Examples: 
-        - 'openai/gpt-4' -> ('openai', 'gpt-4')
-        - 'anthropic/claude-3-opus' -> ('anthropic', 'claude-3-opus')
-        - 'gpt-4' -> ('openai', 'gpt-4')  # default provider
-    """
-    if '/' in model:
-        parts = model.split('/', 1)
-        return parts[0], parts[1]
-    elif ':' in model:
-        parts = model.split(':', 1)
-        return parts[0], parts[1]
-    else:
-        # Default provider
-        return "openai", model
-
-
 def parse_onyx_stream_chunk(chunk_text):
     """Parse a chunk from Onyx stream and extract the text content"""
+    if not chunk_text or not chunk_text.strip():
+        return None, None
+    
     try:
         data = json.loads(chunk_text)
         
-        # Handle different response types from Onyx
         if isinstance(data, dict):
-            # Check for answer/text content
+            # Extract message ID for tracking conversation
+            message_id = data.get('message_id')
+            
+            # Check for different content fields
             if 'answer_piece' in data:
-                return data['answer_piece']
+                return data['answer_piece'], message_id
             elif 'text' in data:
-                return data['text']
+                return data['text'], message_id
             elif 'content' in data:
-                return data['content']
-            elif 'message' in data:
-                return data['message']
-            # Check for error
+                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']}]"
+                return f"[Error: {data['error']}]", message_id
+                
         elif isinstance(data, str):
-            return data
+            return data, None
             
     except json.JSONDecodeError:
-        # If not JSON, return as-is if it looks like content
-        if chunk_text.strip() and not chunk_text.startswith('{'):
-            return chunk_text
+        # Not JSON, might be raw text
+        if chunk_text.strip() and not chunk_text.strip().startswith('{'):
+            return chunk_text.strip(), None
     
-    return None
+    return None, None
 
 
 def generate_openai_stream_chunk(content, model, chunk_id, finish_reason=None):
@@ -142,9 +176,16 @@ def generate_openai_stream_chunk(content, model, chunk_id, finish_reason=None):
     return f"data: {json.dumps(chunk)}\n\n"
 
 
-def stream_onyx_response(payload, headers, model):
+def stream_onyx_response(payload, model, session_key):
     """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"
+    
+    # Try alternate endpoints if needed
+    endpoints = [
+        f"{ONYX_BASE_URL}/api/chat/send-message",
+        f"{ONYX_BASE_URL}/api/chat/send-chat-message",
+    ]
     
     # Send initial chunk with role
     initial_chunk = {
@@ -160,121 +201,177 @@ def stream_onyx_response(payload, headers, model):
     }
     yield f"data: {json.dumps(initial_chunk)}\n\n"
     
-    try:
-        with requests.post(
-            ONYX_API_URL, 
-            json=payload, 
-            headers=headers, 
-            stream=True,
-            timeout=120
-        ) as response:
-            
-            if response.status_code != 200:
-                error_content = f"Error from Onyx API: {response.status_code} - {response.text}"
-                yield generate_openai_stream_chunk(error_content, model, chunk_id)
-                yield generate_openai_stream_chunk("", model, chunk_id, "stop")
-                yield "data: [DONE]\n\n"
-                return
+    last_message_id = None
+    
+    for url in endpoints:
+        try:
+            print(f"Trying endpoint: {url}")
+            print(f"Payload: {json.dumps(payload, indent=2)}")
             
-            buffer = ""
-            for chunk in response.iter_content(chunk_size=None, decode_unicode=True):
-                if chunk:
-                    buffer += chunk
-                    
-                    # Process complete lines
-                    while '\n' in buffer:
-                        line, buffer = buffer.split('\n', 1)
-                        line = line.strip()
+            with requests.post(
+                url, 
+                json=payload, 
+                headers=get_headers(), 
+                stream=True,
+                timeout=120
+            ) as response:
+                
+                print(f"Response status: {response.status_code}")
+                
+                if response.status_code == 404:
+                    continue  # Try next endpoint
+                
+                if response.status_code != 200:
+                    error_text = response.text
+                    print(f"Error response: {error_text}")
+                    yield generate_openai_stream_chunk(
+                        f"Error {response.status_code}: {error_text}", 
+                        model, chunk_id
+                    )
+                    yield generate_openai_stream_chunk("", model, chunk_id, "stop")
+                    yield "data: [DONE]\n\n"
+                    return
+                
+                buffer = ""
+                for chunk in response.iter_content(chunk_size=None, decode_unicode=True):
+                    if chunk:
+                        buffer += chunk
                         
-                        if not line:
-                            continue
-                        
-                        # Handle SSE format from Onyx
-                        if line.startswith('data: '):
-                            line = line[6:]
-                        
-                        if line == '[DONE]':
-                            continue
-                        
-                        content = parse_onyx_stream_chunk(line)
-                        if content:
-                            yield generate_openai_stream_chunk(content, model, chunk_id)
-            
-            # Process any remaining buffer
-            if buffer.strip():
-                content = parse_onyx_stream_chunk(buffer.strip())
-                if content:
-                    yield generate_openai_stream_chunk(content, model, chunk_id)
-    
-    except requests.exceptions.RequestException as e:
-        error_content = f"Request error: {str(e)}"
-        yield generate_openai_stream_chunk(error_content, model, chunk_id)
+                        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 = parse_onyx_stream_chunk(line)
+                            if msg_id:
+                                last_message_id = msg_id
+                            if content:
+                                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())
+                    if msg_id:
+                        last_message_id = msg_id
+                    if content:
+                        yield generate_openai_stream_chunk(content, model, chunk_id)
+                
+                # 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
+                
+        except requests.exceptions.RequestException as e:
+            print(f"Request error for {url}: {e}")
+            continue
     
     # Send final chunk
     yield generate_openai_stream_chunk("", model, chunk_id, "stop")
     yield "data: [DONE]\n\n"
 
 
-def collect_full_response(payload, headers, model):
+def collect_full_response(payload, model, session_key):
     """Collect full streaming response and return as complete OpenAI response"""
     full_content = ""
+    last_message_id = None
     
-    try:
-        with requests.post(
-            ONYX_API_URL, 
-            json=payload, 
-            headers=headers, 
-            stream=True,
-            timeout=120
-        ) as response:
-            
-            if response.status_code != 200:
-                return {
-                    "error": {
-                        "message": f"Error from Onyx API: {response.status_code} - {response.text}",
-                        "type": "api_error",
-                        "code": response.status_code
-                    }
-                }, response.status_code
+    endpoints = [
+        f"{ONYX_BASE_URL}/api/chat/send-message",
+        f"{ONYX_BASE_URL}/api/chat/send-chat-message",
+    ]
+    
+    for url in endpoints:
+        try:
+            print(f"Trying endpoint: {url}")
+            print(f"Payload: {json.dumps(payload, indent=2)}")
             
-            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
+            with requests.post(
+                url, 
+                json=payload, 
+                headers=get_headers(), 
+                stream=True,
+                timeout=120
+            ) as response:
+                
+                print(f"Response status: {response.status_code}")
+                
+                if response.status_code == 404:
+                    continue
+                
+                if response.status_code != 200:
+                    error_text = response.text
+                    print(f"Error response: {error_text}")
+                    return {
+                        "error": {
+                            "message": f"Onyx API error {response.status_code}: {error_text}",
+                            "type": "api_error",
+                            "code": response.status_code
+                        }
+                    }, response.status_code
+                
+                buffer = ""
+                for chunk in response.iter_content(chunk_size=None, decode_unicode=True):
+                    if chunk:
+                        buffer += chunk
                         
-                        content = parse_onyx_stream_chunk(line)
-                        if content:
-                            full_content += content
-            
-            # Process remaining buffer
-            if buffer.strip():
-                content = parse_onyx_stream_chunk(buffer.strip())
-                if content:
-                    full_content += content
+                        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 = parse_onyx_stream_chunk(line)
+                            if msg_id:
+                                last_message_id = msg_id
+                            if content:
+                                full_content += content
+                
+                if buffer.strip():
+                    if buffer.strip().startswith('data: '):
+                        buffer = buffer.strip()[6:]
+                    content, msg_id = parse_onyx_stream_chunk(buffer.strip())
+                    if msg_id:
+                        last_message_id = msg_id
+                    if content:
+                        full_content += content
+                
+                # Update session
+                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}")
+            continue
     
-    except requests.exceptions.RequestException as e:
+    if not full_content:
         return {
             "error": {
-                "message": f"Request error: {str(e)}",
+                "message": "No response from Onyx API",
                 "type": "api_error",
                 "code": 500
             }
         }, 500
     
-    # Build OpenAI-compatible response
     response_data = {
         "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
         "object": "chat.completion",
@@ -289,7 +386,7 @@ def collect_full_response(payload, headers, model):
             "finish_reason": "stop"
         }],
         "usage": {
-            "prompt_tokens": -1,  # Not available from Onyx
+            "prompt_tokens": -1,
             "completion_tokens": -1,
             "total_tokens": -1
         }
@@ -306,8 +403,9 @@ def chat_completions():
     
     try:
         data = request.json
+        print(f"Received request: {json.dumps(data, indent=2)}")
     except Exception as e:
-        return jsonify({"error": {"message": "Invalid JSON", "type": "invalid_request_error"}}), 400
+        return jsonify({"error": {"message": f"Invalid JSON: {e}", "type": "invalid_request_error"}}), 400
     
     # Extract parameters
     model = data.get('model', 'openai/gpt-4')
@@ -315,7 +413,9 @@ def chat_completions():
     stream = data.get('stream', False)
     temperature = data.get('temperature', 0.7)
     
-    # Validate messages
+    # Use a unique session key per conversation or default
+    session_key = data.get('session_id', 'default')
+    
     if not messages:
         return jsonify({
             "error": {
@@ -326,24 +426,32 @@ def chat_completions():
     
     # Parse model string
     model_provider, model_version = parse_model_string(model)
+    print(f"Model provider: {model_provider}, version: {model_version}")
+    
+    # Get or create chat session
+    session_info = get_or_create_session(session_key)
+    
+    if not session_info:
+        return jsonify({
+            "error": {
+                "message": "Failed to create chat session with Onyx API",
+                "type": "api_error"
+            }
+        }), 500
     
     # Build Onyx payload
     payload = build_onyx_payload(
         messages=messages,
         model_provider=model_provider,
         model_version=model_version,
-        temperature=temperature
+        temperature=temperature,
+        chat_session_id=session_info['session_id'],
+        parent_message_id=session_info.get('parent_message_id')
     )
     
-    headers = {
-        "Authorization": f"Bearer {ONYX_API_TOKEN}",
-        "Content-Type": "application/json"
-    }
-    
     if stream:
-        # Return streaming response
         return Response(
-            stream_onyx_response(payload, headers, model),
+            stream_onyx_response(payload, model, session_key),
             content_type='text/event-stream',
             headers={
                 'Cache-Control': 'no-cache',
@@ -352,21 +460,50 @@ def chat_completions():
             }
         )
     else:
-        # Return complete response
-        response_data, status_code = collect_full_response(payload, headers, model)
+        response_data, status_code = collect_full_response(payload, model, session_key)
         return jsonify(response_data), status_code
 
 
+@app.route('/v1/sessions', methods=['POST'])
+def create_new_session():
+    """Create a new chat session"""
+    data = request.json or {}
+    persona_id = data.get('persona_id', 0)
+    
+    session_id = create_chat_session(persona_id)
+    
+    if session_id:
+        session_key = str(uuid.uuid4())
+        chat_sessions_cache[session_key] = {
+            "session_id": session_id,
+            "parent_message_id": None
+        }
+        return jsonify({
+            "session_key": session_key,
+            "chat_session_id": session_id
+        })
+    else:
+        return jsonify({"error": "Failed to create session"}), 500
+
+
+@app.route('/v1/sessions', methods=['DELETE'])
+def clear_sessions():
+    """Clear all cached sessions"""
+    chat_sessions_cache.clear()
+    return jsonify({"status": "cleared"})
+
+
 @app.route('/v1/models', methods=['GET'])
 def list_models():
     """OpenAI-compatible models listing endpoint"""
     models = [
         {"id": "openai/gpt-4", "object": "model", "owned_by": "openai"},
         {"id": "openai/gpt-4-turbo", "object": "model", "owned_by": "openai"},
+        {"id": "openai/gpt-4o", "object": "model", "owned_by": "openai"},
         {"id": "openai/gpt-3.5-turbo", "object": "model", "owned_by": "openai"},
-        {"id": "anthropic/claude-3-opus", "object": "model", "owned_by": "anthropic"},
-        {"id": "anthropic/claude-3-sonnet", "object": "model", "owned_by": "anthropic"},
-        {"id": "anthropic/claude-3-haiku", "object": "model", "owned_by": "anthropic"},
+        {"id": "anthropic/claude-3-opus-20240229", "object": "model", "owned_by": "anthropic"},
+        {"id": "anthropic/claude-3-sonnet-20240229", "object": "model", "owned_by": "anthropic"},
+        {"id": "anthropic/claude-3-haiku-20240307", "object": "model", "owned_by": "anthropic"},
     ]
     
     return jsonify({
@@ -388,7 +525,26 @@ def get_model(model_id):
 @app.route('/health', methods=['GET'])
 def health_check():
     """Health check endpoint"""
-    return jsonify({"status": "healthy", "timestamp": int(time.time())})
+    return jsonify({
+        "status": "healthy", 
+        "timestamp": int(time.time()),
+        "active_sessions": len(chat_sessions_cache)
+    })
+
+
+@app.route('/debug/test-onyx', methods=['GET'])
+def test_onyx_connection():
+    """Test connection to Onyx API"""
+    results = {}
+    
+    # Test creating a session
+    session_id = create_chat_session()
+    results['create_session'] = {
+        "success": session_id is not None,
+        "session_id": session_id
+    }
+    
+    return jsonify(results)
 
 
 @app.route('/', methods=['GET'])
@@ -400,8 +556,11 @@ def root():
         "endpoints": {
             "chat_completions": "/v1/chat/completions",
             "models": "/v1/models",
-            "health": "/health"
-        }
+            "sessions": "/v1/sessions",
+            "health": "/health",
+            "debug": "/debug/test-onyx"
+        },
+        "model_format": "provider/model_version (e.g., openai/gpt-4)"
     })
 
 
@@ -422,19 +581,20 @@ def not_found(e):
 def server_error(e):
     return jsonify({
         "error": {
-            "message": "Internal server error",
+            "message": f"Internal server error: {str(e)}",
             "type": "server_error",
             "code": 500
         }
     }), 500
 
 
-# ============== Main ==============
-
 if __name__ == '__main__':
-    print("Starting OpenAI-Compatible Onyx API Proxy...")
-    print(f"Onyx API URL: {ONYX_API_URL}")
-    print("="*50)
+    print("="*60)
+    print("OpenAI-Compatible Onyx API Proxy")
+    print("="*60)
+    print(f"Onyx Base URL: {ONYX_BASE_URL}")
+    print(f"Token configured: {'Yes' if ONYX_API_TOKEN != '<your-token-here>' else 'No'}")
+    print("="*60)
     
     app.run(
         host='0.0.0.0',