requests_service.py

import requests
import os
import sys

# Fetch the key from HF Secrets
ONYX_API_KEY = os.getenv("PR_KEY")
ONYX_BASE_URL = "https://cloud.onyx.app/api"

def get_headers():
    # If your key already starts with 'Bearer', this logic will break. 
    # This assumes ONYX_API_KEY is JUST the token.
    return {
        "Authorization": f"Bearer {ONYX_API_KEY}",
        "Content-Type": "application/json"
    }

def create_onyx_session():
    url = f"{ONYX_BASE_URL}/chat/create-chat-session"
    # Ensure persona_id is 0 (default) or a valid ID from your Onyx dashboard
    payload = {"persona_id": 0, "description": "Hugging Face Bridge"}
    
    try:
        print(f"DEBUG: Attempting session create with key ending in ...{ONYX_API_KEY[-5:] if ONYX_API_KEY else 'NONE'}")
        resp = requests.post(url, json=payload, headers=get_headers(), timeout=10)
        
        if resp.status_code == 200:
            return resp.json().get("id")
        
        # This will print to your Hugging Face Logs
        print(f"ONYX AUTH REJECTION: {resp.status_code} - {resp.text}", file=sys.stderr)
        return None
    except Exception as e:
        print(f"CONNECTION ERROR: {str(e)}", file=sys.stderr)
        return None

def send_onyx_message(message, model_name, session_id, provider, version):
    url = f"{ONYX_BASE_URL}/chat/send-chat-message"
    
    # EXACT payload structure you requested
    payload = {
        "message": message,
        "llm_override": {
            "model_provider": provider,
            "model_version": version,
            "temperature": 0.7
        },
        "allowed_tool_ids": [],
        "file_descriptors": [],
        "internal_search_filters": {"source_type": ["ingestion_api"]},
        "deep_research": False,
        "origin": "unset",
        "parent_message_id": None,
        "chat_session_id": session_id,
        "stream": True,
        "include_citations": True
    }
    
    return requests.post(url, json=payload, headers=get_headers(), stream=True)