File size: 5,607 Bytes
08ca036 |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import asyncio
import types
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from fastapi import Request, HTTPException, Form, Security
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.security import HTTPAuthorizationCredentials
from starlette.background import BackgroundTask
import utils.globals as globals
from app import app, templates, security_scheme
from chatgpt.ChatService import ChatService
from chatgpt.authorization import refresh_all_tokens
from utils.Logger import logger
from utils.configs import api_prefix, scheduled_refresh
from utils.retry import async_retry
scheduler = AsyncIOScheduler()
@app.on_event("startup")
async def app_start():
if scheduled_refresh:
scheduler.add_job(id='refresh', func=refresh_all_tokens, trigger='cron', hour=3, minute=0, day='*/2',
kwargs={'force_refresh': True})
scheduler.start()
asyncio.get_event_loop().call_later(0, lambda: asyncio.create_task(refresh_all_tokens(force_refresh=False)))
async def to_send_conversation(request_data, req_token):
chat_service = ChatService(req_token)
try:
await chat_service.set_dynamic_data(request_data)
await chat_service.get_chat_requirements()
return chat_service
except HTTPException as e:
await chat_service.close_client()
raise HTTPException(status_code=e.status_code, detail=e.detail)
except Exception as e:
await chat_service.close_client()
logger.error(f"Server error, {str(e)}")
raise HTTPException(status_code=500, detail="Server error")
async def process(request_data, req_token):
chat_service = await to_send_conversation(request_data, req_token)
await chat_service.prepare_send_conversation()
res = await chat_service.send_conversation()
return chat_service, res
@app.post(f"/{api_prefix}/v1/chat/completions" if api_prefix else "/v1/chat/completions")
async def send_conversation(request: Request, credentials: HTTPAuthorizationCredentials = Security(security_scheme)):
req_token = credentials.credentials
try:
request_data = await request.json()
except Exception:
raise HTTPException(status_code=400, detail={"error": "Invalid JSON body"})
chat_service, res = await async_retry(process, request_data, req_token)
try:
if isinstance(res, types.AsyncGeneratorType):
background = BackgroundTask(chat_service.close_client)
return StreamingResponse(res, media_type="text/event-stream", background=background)
else:
background = BackgroundTask(chat_service.close_client)
return JSONResponse(res, media_type="application/json", background=background)
except HTTPException as e:
await chat_service.close_client()
if e.status_code == 500:
logger.error(f"Server error, {str(e)}")
raise HTTPException(status_code=500, detail="Server error")
raise HTTPException(status_code=e.status_code, detail=e.detail)
except Exception as e:
await chat_service.close_client()
logger.error(f"Server error, {str(e)}")
raise HTTPException(status_code=500, detail="Server error")
@app.get(f"/{api_prefix}/tokens" if api_prefix else "/tokens", response_class=HTMLResponse)
async def upload_html(request: Request):
tokens_count = len(set(globals.token_list) - set(globals.error_token_list))
return templates.TemplateResponse("tokens.html",
{"request": request, "api_prefix": api_prefix, "tokens_count": tokens_count})
@app.post(f"/{api_prefix}/tokens/upload" if api_prefix else "/tokens/upload")
async def upload_post(text: str = Form(...)):
lines = text.split("\n")
for line in lines:
if line.strip() and not line.startswith("#"):
globals.token_list.append(line.strip())
with open(globals.TOKENS_FILE, "a", encoding="utf-8") as f:
f.write(line.strip() + "\n")
logger.info(f"Token count: {len(globals.token_list)}, Error token count: {len(globals.error_token_list)}")
tokens_count = len(set(globals.token_list) - set(globals.error_token_list))
return {"status": "success", "tokens_count": tokens_count}
@app.post(f"/{api_prefix}/tokens/clear" if api_prefix else "/tokens/clear")
async def upload_post():
globals.token_list.clear()
globals.error_token_list.clear()
with open(globals.TOKENS_FILE, "w", encoding="utf-8") as f:
pass
logger.info(f"Token count: {len(globals.token_list)}, Error token count: {len(globals.error_token_list)}")
tokens_count = len(set(globals.token_list) - set(globals.error_token_list))
return {"status": "success", "tokens_count": tokens_count}
@app.post(f"/{api_prefix}/tokens/error" if api_prefix else "/tokens/error")
async def error_tokens():
error_tokens_list = list(set(globals.error_token_list))
return {"status": "success", "error_tokens": error_tokens_list}
@app.get(f"/{api_prefix}/tokens/add/{{token}}" if api_prefix else "/tokens/add/{token}")
async def add_token(token: str):
if token.strip() and not token.startswith("#"):
globals.token_list.append(token.strip())
with open(globals.TOKENS_FILE, "a", encoding="utf-8") as f:
f.write(token.strip() + "\n")
logger.info(f"Token count: {len(globals.token_list)}, Error token count: {len(globals.error_token_list)}")
tokens_count = len(set(globals.token_list) - set(globals.error_token_list))
return {"status": "success", "tokens_count": tokens_count}
|