app.py

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import os
import httpx # A modern, async HTTP client

app = FastAPI()

# Configuration from environment variables/secrets
OMNI_API_KEY = os.getenv("OMNI_API_KEY")
OMNI_BASE_URL = "https://user669-omniapi.hf.space" # The base URL for the external API
CHAT_COMPLETIONS_PATH = "/chat/completions" # The specific endpoint path
DEFAULT_MODEL = "gpt-4o" # The default model you want to use

# Basic input model for the chat completions request
class ChatRequest(BaseModel):
    messages: list[dict]
    model: str = DEFAULT_MODEL
    # You can add other parameters here if your API supports them, e.g., temperature, max_tokens
    # temperature: float = 0.7
    # max_tokens: int = 150

@app.on_event("startup")
async def startup_event():
    # Check if API key is loaded
    if not OMNI_API_KEY:
        print("CRITICAL ERROR: OMNI_API_KEY not found in environment variables.")
        # In a production app, you might want to stop the service or log more aggressively
        # For Spaces, it will just start but requests to /chat will fail.

@app.post("/chat")
async def chat_proxy(request: ChatRequest):
    if not OMNI_API_KEY:
        raise HTTPException(status_code=500, detail="API key not configured.")

    headers = {
        "Authorization": f"Bearer {OMNI_API_KEY}",
        "Content-Type": "application/json"
    }

    payload = {
        "model": request.model,
        "messages": request.messages,
        # Add any other parameters from ChatRequest here if they are part of the payload
        # "temperature": request.temperature,
        # "max_tokens": request.max_tokens,
    }

    full_url = f"{OMNI_BASE_URL}{CHAT_COMPLETIONS_PATH}"

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(full_url, json=payload, headers=headers, timeout=60.0) # Add a timeout
            response.raise_for_status() # Raise an exception for HTTP errors (4xx or 5xx)
            return JSONResponse(content=response.json())
        except httpx.RequestError as exc:
            raise HTTPException(status_code=500, detail=f"An error occurred while requesting: {exc.request.url!r}.")
        except httpx.HTTPStatusError as exc:
            raise HTTPException(status_code=exc.response.status_code, detail=f"Error response {exc.response.status_code} from {exc.request.url!r}: {exc.response.text}")
        except Exception as e:
            # Catch any other unexpected errors
            raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")

# Health check endpoint (optional but good practice)
@app.get("/health")
def health_check():
    return {"status": "ok"}

# Ensure the app listens on port 7860 for Hugging Face Spaces
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)