from typing import Any

from flask import Flask, request
from langchain.callbacks import StreamingStdOutCallbackHandler
from langchain.schema import LLMResult

from zl_chatbot import ZlChatBot
import urllib.parse

# 创建Flask应用
app = Flask(__name__)
app.config['PERMANENT_SESSION_LIFETIME'] = 600  # 单位为秒

# chatbot.init_from_pdf('./data/test.pdf')
chatbot = ZlChatBot().init_chatbot_from_vec_db(['./vector_storage/zl_db/zhongliang_abstract', './vector_storage/zl_db/zhongliang_keyword', './vector_storage/zl_db/zhongliang_query'])

import time
def generate_steam(text: str):
    for c in text:
        yield c
        time.sleep(0.0001)

@app.route('/query2kb_stream')
def stream():
    query = request.args.get('q')
    # 对查询字符串参数进行解码
    query = urllib.parse.unquote(query)
    res = chatbot.get_from_cache(query)
    if res:
        return generate_steam(res+"[EOF]")
    return chatbot.get_stream(query)

# 定义路由和视图函数
@app.route('/query2kb')
def display_string():
    query = request.args.get('q')
    # 对查询字符串参数进行解码
    query = urllib.parse.unquote(query)
    res = chatbot.get_from_cache(query)
    if res:
        return res
    return chatbot.query2kb(query)

# 运行应用
if __name__ == '__main__':
    app.run(host="0.0.0.0")