import json
import time

from flask import Flask, request
from flask import stream_with_context, Response
from dataclasses import dataclass, asdict
from ko_chatbot import KoChatBot

##程序开始
app = Flask(__name__)
app.config['PERMANENT_SESSION_LIFETIME'] = 600  # 单位为秒

chatbot = KoChatBot().init_ko_chatbot(['vector_storage/wiki_test'])

@app.route('/chat', methods=['GET', 'POST'])
def chat():
    """
    直接与大模型进行对话
    :return:
    """
    request_params = request.get_json()
    if request_params is None or request_params == []:
        request_params = request.form

    # 获取所有的参数, 如果没有则设置默认值
    prompt = request_params.get("prompt")
    prompt = prompt if prompt else "你好！"
    history = request_params.get("history")
    history = json.loads(history) if history else []
    history = [(e['question'], e['answer']) for e in history]
    max_length = request_params.get("max_length")
    max_length = max_length if max_length else 8192
    top_p = request_params.get("top_p")
    top_p = top_p if top_p else 0.9
    temperature = request_params.get("temperature")
    temperature = temperature if temperature else 0.1
    stream = request_params.get("stream")
    stream = stream if stream else False

    res = {
        "response": '',
        "stream": False,
        "time": None,
        "status": 500,
        "errorMessage": "发生错误"
    }

    def chat():
        response = chatbot.query2llm(prompt)
        res["response"] = response
        res["time"] = time.time()
        res["errorMessage"] = ""
        res['history'] = []
        return json.dumps(res) + "\n"

    return Response(chat(), mimetype='application/json')


@app.route('/chat2kb', methods=['GET', 'POST'])
def chat2kb():
    request_params = request.get_json()
    if request_params is None or request_params == []:
        request_params = request.form

    # 获取所有的参数, 如果没有则设置默认值
    prompt = request_params.get("prompt")
    prompt = prompt if prompt else "你好！"
    history = request_params.get("history")
    history = json.loads(history) if history else []
    history = [(e['question'], e['answer']) for e in history]
    max_length = request_params.get("max_length")
    max_length = max_length if max_length else 8192
    top_p = request_params.get("top_p")
    top_p = top_p if top_p else 0.9
    temperature = request_params.get("temperature")
    temperature = temperature if temperature else 0.1
    stream = request_params.get("stream")
    stream = stream if stream else False
    res = {
        "response": '',
        "stream": False,
        "time": None,
        "status": 500,
        "errorMessage": "发生错误"
    }

    def chat(query):
        response, docs = chatbot.query2kb(query)
        res["status"] = 200
        res["response"] = response
        res["time"] = time.time()
        res["errorMessage"] = ""
        res['history'] = []
        res['docs'] = docs
        return json.dumps(res) + "\n"
    return Response(chat(prompt), mimetype='application/json')


@app.route('/chat2url', methods=['GET', 'POST'])
def chat2url():
    request_params = request.get_json()
    if request_params is None or request_params == []:
        request_params = request.form

    # 获取所有的参数, 如果没有则设置默认值
    prompt = request_params.get("prompt")
    prompt = prompt if prompt else "你好！"
    history = request_params.get("history")
    history = json.loads(history) if history else []
    history = [(e['question'], e['answer']) for e in history]
    max_length = request_params.get("max_length")
    max_length = max_length if max_length else 8192
    top_p = request_params.get("top_p")
    top_p = top_p if top_p else 0.9
    temperature = request_params.get("temperature")
    temperature = temperature if temperature else 0.1
    stream = request_params.get("stream")
    stream = stream if stream else False
    url = request_params.get('url')
    res = {
        "response": '',
        "stream": False,
        "time": None,
        "status": 500,
        "errorMessage": "发生错误"
    }

    def chat(url, query):
        response = chatbot.query2url(url, query)
        res["status"] = 200
        res["response"] = response
        res["time"] = time.time()
        res["errorMessage"] = ""
        res['history'] = []
        return json.dumps(res) + "\n"

    return Response(chat(url,prompt), mimetype='application/json')


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8989)
    print('a')