from hugchat import hugchat
from hugchat.login import Login
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

# 登录 Hugging Face 并授权 huggingchat
EMAIL = "your email"
PASSWD = "your password"
cookie_path_dir = "./cookies/"
sign = Login(EMAIL, PASSWD)
cookies = sign.login(cookie_dir_path=cookie_path_dir, save_cookies=True)

# 创建 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)