main.py

import httpx
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
import json
import random
import logging

# Configure logging to output to stdout
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

app = FastAPI()

# List of API URLs to be randomized
API_URLS = [
    "https://api.deepinfra.com/v1/openai/chat/completions",
    "https://stage.api.deepinfra.com/v1/openai/chat/completions",
]

@app.post("/v1/openai/chat/completions")
async def proxy_deepinfra(request: Request):
    """
    Proxies chat completion requests to the DeepInfra API.
    It randomizes the order of API URLs and uses the next as a fallback.
    """
    try:
        body = await request.json()
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid JSON in request body")

    headers = {
        'sec-ch-ua-platform': request.headers.get('sec-ch-ua-platform', '"Windows"'),
        'Referer': request.headers.get('Referer', 'https://deepinfra.com/'),
        'sec-ch-ua': request.headers.get('sec-ch-ua', '"Chromium";v="140", "Not=A?Brand";v="24", "Google Chrome";v="140"'),
        'sec-ch-ua-mobile': request.headers.get('sec-ch-ua-mobile', '?0'),
        'User-Agent': request.headers.get('User-Agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36'),
        'accept': request.headers.get('accept', 'text-event-stream'),
        'X-Deepinfra-Source': request.headers.get('X-Deepinfra-Source', 'web-embed'),
        'Content-Type': request.headers.get('Content-Type', 'application/json'),
    }

    shuffled_urls = random.sample(API_URLS, len(API_URLS))

    async def stream_generator():
        last_error = None
        for url in shuffled_urls:
            logging.info(f"Attempting to connect to: {url}")
            try:
                async with httpx.AsyncClient() as client:
                    async with client.stream("POST", url, headers=headers, json=body, timeout=None) as response:
                        response.raise_for_status()
                        logging.info(f"Successfully connected. Streaming from: {url}")
                        async for chunk in response.aiter_bytes():
                            yield chunk
                        return
            except (httpx.RequestError, httpx.HTTPStatusError) as e:
                last_error = e
                logging.warning(f"Failed to connect to {url}: {e}. Trying next URL.")
                continue

        if last_error:
            logging.error(f"All API endpoints failed. Last error: {last_error}")
            # In a streaming response, we can't easily raise an HTTPException after starting.
            # The connection will simply close, which the client will see as a failed request.

    return StreamingResponse(stream_generator(), media_type="text-event-stream")

# The `if __name__ == "__main__":` block is removed as Gunicorn will run the app.