Spaces:
Running
Running
File size: 4,197 Bytes
4986fe4 ef215d3 7c60ac5 378f2c3 8834a20 3109050 378f2c3 4986fe4 378f2c3 8e4491b 4986fe4 8e4491b 4986fe4 8e4491b 378f2c3 7c60ac5 7ef5d89 7c60ac5 4986fe4 7c60ac5 4986fe4 7c60ac5 e2e24f9 7c60ac5 7ef5d89 a68045e 8834a20 e2e24f9 8834a20 7ef5d89 3109050 045bd95 ef215d3 378f2c3 8e4491b 378f2c3 045bd95 378f2c3 8e4491b 9d35223 045bd95 3109050 8834a20 8e4491b 8834a20 114ca84 3109050 378f2c3 7ef5d89 045bd95 6a84e5c 8e4491b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import os
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse, HTMLResponse
from pydantic import BaseModel
import httpx
from functools import lru_cache
load_dotenv()
app = FastAPI()
# Get API keys and secret endpoint from environment variables
api_keys_str = os.getenv('API_KEYS')
valid_api_keys = api_keys_str.split(',') if api_keys_str else []
secret_api_endpoint = os.getenv('SECRET_API_ENDPOINT')
secret_api_endpoint_2 = os.getenv('SECRET_API_ENDPOINT_2')
# Validate if the main secret API endpoint is set
if not secret_api_endpoint or not secret_api_endpoint_2:
raise HTTPException(status_code=500, detail="API endpoint(s) are not configured in environment variables.")
# Define models that should use the secondary endpoint
alternate_models = {"gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b"}
class Payload(BaseModel):
model: str
messages: list
@app.get("/", response_class=HTMLResponse)
async def root():
html_content = """
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Loki.AI API</title>
<style>
body { font-family: Arial, sans-serif; text-align: center; margin-top: 50px; background-color: #121212; color: white; }
h1 { color: #4CAF50; }
a { color: #BB86FC; text-decoration: none; }
a:hover { text-decoration: underline; }
</style>
</head>
<body>
<h1>Welcome to Loki.AI API!</h1>
<p>Created by Parth Sadaria</p>
<p>Go to /models for more info</p>
<p>Check out the GitHub for more projects:</p>
<a href="https://github.com/ParthSadaria" target="_blank">github.com/ParthSadaria</a>
</body>
</html>
"""
return HTMLResponse(content=html_content)
async def get_models():
async with httpx.AsyncClient() as client:
try:
response = await client.get(f"{secret_api_endpoint}/v1/models", timeout=3)
response.raise_for_status()
return response.json()
except httpx.RequestError as e:
raise HTTPException(status_code=500, detail=f"Request failed: {e}")
@app.get("/v1/models")
async def fetch_models():
return await get_models()
@app.post("/v1/chat/completions")
async def get_completion(payload: Payload, request: Request):
api_key = request.headers.get("Authorization")
# Validate API key
if api_key not in valid_api_keys:
raise HTTPException(status_code=403, detail="Forbidden: Invalid API key. Join dsc.gg/chadgang and DM @mr_leaderyt on Discord for a free API key :)")
# Determine which endpoint to use based on the model
endpoint = secret_api_endpoint_2 if payload.model in alternate_models else secret_api_endpoint
payload_dict = {**payload.dict(), "stream": True}
async def stream_generator(payload_dict):
async with httpx.AsyncClient() as client:
try:
async with client.stream("POST", f"{endpoint}/v1/chat/completions", json=payload_dict, timeout=10) as response:
response.raise_for_status()
async for line in response.aiter_lines():
if line:
yield f"{line}\n"
except httpx.HTTPStatusError as status_err:
raise HTTPException(status_code=status_err.response.status_code, detail=f"HTTP error: {status_err}")
except httpx.RequestError as req_err:
raise HTTPException(status_code=500, detail=f"Streaming failed: {req_err}")
except Exception as e:
raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
return StreamingResponse(stream_generator(payload_dict), media_type="application/json")
@app.on_event("startup")
async def startup_event():
print("API endpoints:")
print("GET /")
print("GET /models")
print("POST /v1/chat/completions")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
|