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")

sign = Login(EMAIL, PASSWD)

try:
    print("Attempting to login...")
    cookies = sign.login(save_cookies=False)  # 不保存到文件中，直接获取cookie字典
    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)