main.py

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

app = FastAPI()

# Model and API Key validation
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):
    api_key = payload.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
    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"
    
    # Send the API request to the appropriate service
    response = requests.post(url, json={**payload.dict(), "model": full_model_name})
    
    if response.status_code != 200:
        raise HTTPException(status_code=response.status_code, detail="Failed to fetch data")
    
    # Return the response from the model API
    return response.json()

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