merge main.py with new keys system

This commit is contained in:
Game_Time 2023-06-24 23:56:56 +05:00 committed by GitHub
parent 04413a3c6a
commit 3ef5746b1a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -1,6 +1,7 @@
import os import os
import httpx import httpx
import fastapi import fastapi
from keys import Keys
from starlette.requests import Request from starlette.requests import Request
from starlette.responses import StreamingResponse from starlette.responses import StreamingResponse
@ -31,33 +32,43 @@ async def startup_event():
security.enable_proxy() security.enable_proxy()
security.ip_protection_check() security.ip_protection_check()
# Setup key cache
Keys()
@app.get('/') @app.get('/')
async def root(): async def root():
"""Returns the root endpoint.""" """Returns the root endpoint."""
return { return {
'status': 'ok', 'status': 'ok',
'discord': 'https://discord.gg/mX9BYdFeQF', 'discord': 'https://discord.gg/85gdcd57Xr',
'github': 'https://github.com/Luna-OSS' 'github': 'https://github.com/Luna-OSS'
} }
async def _reverse_proxy(request: Request): async def _reverse_proxy(request: Request):
target_url = f'https://api.openai.com/v1/{request.url.path}' target_url = f'https://api.openai.com/v1/{request.url.path}'
key = Keys.get(request.body()['model'])
if not key:
return fastapi.responses.JSONResponse(
status_code=400,
content={
'error': 'No API Key for model given, please try again with a valid model.'
}
)
request_to_api = target_api_client.build_request( request_to_api = target_api_client.build_request(
method=request.method, method=request.method,
url=target_url, url=target_url,
headers={ headers={
'Authorization': 'Bearer ' + os.getenv('OPENAI_KEY'), 'Authorization': 'Bearer ' + key,
'Content-Type': 'application/json' 'Content-Type': 'application/json'
}, },
content=await request.body(), content=await request.body(),
) )
api_response = await target_api_client.send(request_to_api, stream=True) api_response = await target_api_client.send(request_to_api, stream=True)
print(f'[{request.method}] {request.url.path} {api_response.status_code}') print(f'[{request.method}] {request.url.path} {api_response.status_code}')
Keys(key).unlock()
return StreamingResponse( return StreamingResponse(
api_response.aiter_raw(), api_response.aiter_raw(),
status_code=api_response.status_code, status_code=api_response.status_code,