main.py

from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
import requests
import json

app = FastAPI()

valid_api_keys = ['PARTH-SADARIA-NI-API-CHAWI', 'HEET-NI-CHEESEWADI-API-KEY']
model_aliases = {
    'Llama-3.1-Nemotron-70B-Instruct': 'nvidia/Llama-3.1-Nemotron-70B-Instruct',
    'Meta-Llama-3.1-8B-Instruct': 'meta-llama/Meta-Llama-3.1-8B-Instruct',
    'Meta-Llama-3.1-70B-Instruct': 'meta-llama/Meta-Llama-3.1-70B-Instruct',
    'Meta-Llama-3.1-70B-Instruct-Turbo': 'meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo',
    'Meta-Llama-3.1-405B-Instruct': 'meta-llama/Meta-Llama-3.1-405B-Instruct',
    'Llama-3.2-11B-Vision-Instruct': 'meta-llama/Llama-3.2-11B-Vision-Instruct',
    'Meta-Llama-3-8B-Instruct': 'meta-llama/Meta-Llama-3-8B-Instruct',
}

class Payload(BaseModel):
    model: str
    messages: list

@app.post("/api/v1/chat/completions")
async def get_completion(payload: Payload, request: Request):
    api_key = request.headers.get("Authorization")
    
    # API Key validation
    if api_key not in valid_api_keys:
        raise HTTPException(status_code=403, detail="Forbidden: Invalid API key")
    
    # Model alias resolution
    user_model = payload.model
    if user_model not in model_aliases:
        raise HTTPException(status_code=400, detail="Invalid model name")
    
    full_model_name = model_aliases.get(user_model, user_model)
    
    is_deepinfra_model = full_model_name in model_aliases.values()

    # Determine the URL to send the request to
    url = "https://api.deepinfra.com/v1/openai/chat/completions" if is_deepinfra_model else "https://gpt.tiptopuni.com/api/openai/v1/chat/completions"
    
    try:
        response = requests.post(url, json={**payload.dict(), "model": full_model_name})
        response.raise_for_status()  # Raises HTTPError if the response status code is 4xx/5xx
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"Request failed: {e}")
    
    try:
        return response.json()
    except ValueError:
        raise HTTPException(status_code=500, detail="Non-JSON response received")

# Run the server with Uvicorn using the 'main' module
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)