from fastapi import FastAPI
import openai

import requests
import bot


app = FastAPI()

# 初始化
openai.api_key = "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"
openai.api_base = "https://openai.naraka.top/v1"




@app.get("/generate/")
async def generate(prompt: str, types: str = 'content', model: str = "gpt-3.5-turbo-0613"):
    messages = [
        {"role": "user", "content": prompt}
    ]
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0.3,
    )

    if types == 'content':
        content = response.choices[0].message.get("content", '')
    else:
        content = json.loads(response.choices[0].message.get("content", ''))

    usage = response.usage

    # 保存到数据库
    post = {"prompt": prompt, "usage": usage, "content": content}
    collection.insert_one(post)

    return {"content": content}