hfchat / app.py
tianlong12's picture
Update app.py
64816e4 verified
raw
history blame
1.44 kB
from hugchat import hugchat
from hugchat.login import Login
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn
import os
# 从环境变量中读取登录信息
EMAIL = os.getenv("EMAIL")
PASSWD = os.getenv("PASSWD")
if not EMAIL or not PASSWD:
raise Exception("EMAIL or PASSWD environment variables are not set")
# 使用 /tmp/cookies 目录
cookie_path_dir = "/tmp/cookies/"
if not os.path.exists(cookie_path_dir):
os.makedirs(cookie_path_dir)
sign = Login(EMAIL, PASSWD)
try:
print("Attempting to login...")
cookies = sign.login(cookie_dir_path=cookie_path_dir, save_cookies=True)
print("Login successful!")
except Exception as e:
print(f"Login failed: {e}")
raise
# 创建 ChatBot
chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
class RequestBody(BaseModel):
prompt: str
max_tokens: int
class Choice(BaseModel):
text: str
class CompletionResponse(BaseModel):
choices: list[Choice]
app = FastAPI()
@app.post("/v1/chat/completions", response_model=CompletionResponse)
async def completions(body: RequestBody):
try:
response = chatbot.chat(body.prompt, max_length=body.max_tokens).wait_until_done()
return CompletionResponse(choices=[Choice(text=response)])
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)