File size: 1,440 Bytes
ef5946c
 
 
 
 
3c30c75
f6190d4
3c30c75
 
 
c8bf377
 
 
 
64816e4
 
f6190d4
 
 
ef5946c
c8bf377
 
b331f16
c8bf377
6958029
c8bf377
 
 
ef5946c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
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)