from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from fastapi.responses import JSONResponse
import requests
import json


class Text(BaseModel):
    content: str = ""


app = FastAPI()

headers = {
    'Content-Type': 'application/json',
    'Authorization': 'Bearer ' + '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'
}


@app.get("/")
def home():
    html_content = open('index.html').read()
    return HTMLResponse(content=html_content, status_code=200)


@app.post("/qa_maker")
def sentiment_analysis_ep(content: Text = None):
    url = 'https://api.openai.com/v1/chat/completions'
    prompt = '根据下面的文章，生成的“问题和回答”QA对，大于5个，以一行一个json格式（{“question”:"xxx","answer":"xxx"}）生成：\n'
    messages = [{"role": "user", "content": prompt + content.content}]
    data = {
        "model": "gpt-3.5-turbo",
        "messages": messages
    }
    print("messages = \n", messages)
    result = requests.post(url=url,
                           data=json.dumps(data),
                           headers=headers
                           )
    res = str(result.json()['choices'][0]['message']['content']).strip()
    print('res:', res)
    res = {'content': res}
    return JSONResponse(content=res)


@app.post("/chatpdf")
def chat_pdf_ep(content: Text = None):
    url = 'https://api.openai.com/v1/chat/completions'
    messages = [
        {
            "role": "system",
            "content": "你是一个有用的助手，可以使用文章内容准确地回答问题。使用提供的文章来形成你的答案，但避免逐字复制文章。尽可能使用自己的话。准确、有用、简洁、清晰。"
        }
    ]
    obj = json.loads(content.content)
    messages.append({"role": "system", "content": "文章内容：\n"+obj['doc']})
    history = obj['history']
    for his in history:
        messages.append({"role": "user", "content": his[0]})
        messages.append({"role": "assistant", "content": his[1]})
    messages.append({"role": "user", "content": obj['question']})
    data = {
        "model": "gpt-3.5-turbo",
        "messages": messages
    }
    print("messages = \n", messages)
    result = requests.post(url=url,
                           data=json.dumps(data),
                           headers=headers
                           )
    res = str(result.json()['choices'][0]['message']['content']).strip()
    content = {'content': res}
    print('content:', content)
    return JSONResponse(content=content)


@app.post("/embeddings")
def embeddings_ep(content: Text = None):
    url = 'https://api.openai.com/v1/embeddings'
    data = {
        "model": "text-embedding-ada-002",
        "input": content.content
    }
    result = requests.post(url=url,
                           data=json.dumps(data),
                           headers=headers
                           )
    return JSONResponse(content=result.json())