Spaces:
Running
Running
Update main.py
Browse files
main.py
CHANGED
|
@@ -1,10 +1,15 @@
|
|
| 1 |
from fastapi import FastAPI, HTTPException, Request
|
| 2 |
-
from fastapi.responses import StreamingResponse
|
| 3 |
from pydantic import BaseModel
|
| 4 |
import requests
|
|
|
|
|
|
|
| 5 |
|
| 6 |
app = FastAPI()
|
| 7 |
|
|
|
|
|
|
|
|
|
|
| 8 |
# Valid API keys
|
| 9 |
valid_api_keys = ['PARTH-SADARIA-NI-API-CHAWI', 'HEET-NI-CHEESEWADI-API-KEY']
|
| 10 |
|
|
@@ -12,19 +17,46 @@ class Payload(BaseModel):
|
|
| 12 |
model: str
|
| 13 |
messages: list
|
| 14 |
|
| 15 |
-
@app.get("/")
|
| 16 |
async def root():
|
| 17 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
|
| 19 |
-
|
| 20 |
-
|
|
|
|
| 21 |
try:
|
| 22 |
-
response =
|
| 23 |
response.raise_for_status()
|
| 24 |
return response.json()
|
| 25 |
except requests.exceptions.RequestException as e:
|
| 26 |
raise HTTPException(status_code=500, detail=f"Request failed: {e}")
|
| 27 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 28 |
@app.post("/v1/chat/completions")
|
| 29 |
async def get_completion(payload: Payload, request: Request):
|
| 30 |
api_key = request.headers.get("Authorization")
|
|
@@ -42,7 +74,7 @@ async def get_completion(payload: Payload, request: Request):
|
|
| 42 |
# Define a generator to stream the response
|
| 43 |
def stream_generator():
|
| 44 |
try:
|
| 45 |
-
with
|
| 46 |
response.raise_for_status()
|
| 47 |
for chunk in response.iter_content(chunk_size=1024):
|
| 48 |
if chunk: # Only yield non-empty chunks
|
|
|
|
| 1 |
from fastapi import FastAPI, HTTPException, Request
|
| 2 |
+
from fastapi.responses import StreamingResponse, HTMLResponse
|
| 3 |
from pydantic import BaseModel
|
| 4 |
import requests
|
| 5 |
+
import time
|
| 6 |
+
from functools import lru_cache
|
| 7 |
|
| 8 |
app = FastAPI()
|
| 9 |
|
| 10 |
+
# Create a session for reusing the HTTP connection
|
| 11 |
+
session = requests.Session()
|
| 12 |
+
|
| 13 |
# Valid API keys
|
| 14 |
valid_api_keys = ['PARTH-SADARIA-NI-API-CHAWI', 'HEET-NI-CHEESEWADI-API-KEY']
|
| 15 |
|
|
|
|
| 17 |
model: str
|
| 18 |
messages: list
|
| 19 |
|
| 20 |
+
@app.get("/", response_class=HTMLResponse)
|
| 21 |
async def root():
|
| 22 |
+
html_content = """
|
| 23 |
+
<!DOCTYPE html>
|
| 24 |
+
<html lang="en">
|
| 25 |
+
<head>
|
| 26 |
+
<meta charset="UTF-8">
|
| 27 |
+
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
| 28 |
+
<title>Loki.AI API</title>
|
| 29 |
+
<style>
|
| 30 |
+
body { font-family: Arial, sans-serif; text-align: center; margin-top: 50px; }
|
| 31 |
+
h1 { color: #4CAF50; }
|
| 32 |
+
a { color: #007BFF; text-decoration: none; }
|
| 33 |
+
a:hover { text-decoration: underline; }
|
| 34 |
+
</style>
|
| 35 |
+
</head>
|
| 36 |
+
<body>
|
| 37 |
+
<h1>Welcome to Loki.AI API!</h1>
|
| 38 |
+
<p>Created by Parth Sadaria</p>
|
| 39 |
+
<p>Check out the GitHub for more projects:</p>
|
| 40 |
+
<a href="https://github.com/ParthSadaria" target="_blank">github.com/ParthSadaria</a>
|
| 41 |
+
</body>
|
| 42 |
+
</html>
|
| 43 |
+
"""
|
| 44 |
+
return HTMLResponse(content=html_content)
|
| 45 |
|
| 46 |
+
# Cache function with lru_cache
|
| 47 |
+
@lru_cache(maxsize=1)
|
| 48 |
+
def get_cached_models():
|
| 49 |
try:
|
| 50 |
+
response = session.get("https://lokiai.vercel.app/api/v1/models", timeout=3)
|
| 51 |
response.raise_for_status()
|
| 52 |
return response.json()
|
| 53 |
except requests.exceptions.RequestException as e:
|
| 54 |
raise HTTPException(status_code=500, detail=f"Request failed: {e}")
|
| 55 |
|
| 56 |
+
@app.get("/models")
|
| 57 |
+
async def get_models():
|
| 58 |
+
return get_cached_models()
|
| 59 |
+
|
| 60 |
@app.post("/v1/chat/completions")
|
| 61 |
async def get_completion(payload: Payload, request: Request):
|
| 62 |
api_key = request.headers.get("Authorization")
|
|
|
|
| 74 |
# Define a generator to stream the response
|
| 75 |
def stream_generator():
|
| 76 |
try:
|
| 77 |
+
with session.post(url, json=payload_dict, stream=True, timeout=5) as response:
|
| 78 |
response.raise_for_status()
|
| 79 |
for chunk in response.iter_content(chunk_size=1024):
|
| 80 |
if chunk: # Only yield non-empty chunks
|