Spaces:
Paused
Paused
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() | |
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) | |