from hugchat import hugchat
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn
import os

# 从环境变量中读取 cookies
cookies_str = os.getenv("COOKIES_STR")
if not cookies_str:
    raise Exception("COOKIES_STR environment variable is not set")

# 打印调试信息
print(f"COOKIES_STR: {cookies_str}")

try:
    print("Attempting to initialize ChatBot with provided cookies...")
    chatbot = hugchat.ChatBot(cookies=cookies_str)
    print("ChatBot initialized successfully!")
except Exception as e:
    print(f"Failed to initialize ChatBot: {e}")
    raise

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)