# from gevent import monkey
# monkey.patch_all()

from flask import Flask, request, Response
import os
from dotenv import load_dotenv
import json
import time
from gevent.lock import RLock


import datetime
import logging
# 加载环境变量
load_dotenv()
from DifyQueryRetrieval import DifyQueryRetrieval


logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)
logging.getLogger('httpx').setLevel(logging.WARNING)
logging.getLogger('openai').setLevel(logging.WARNING)

logger = logging.getLogger(__name__)

app = Flask(__name__)

dify_query_retrieval = DifyQueryRetrieval(api_key="dataset-skLjmPVonjHo119OWNf3kAmY", base_url="http://10.1.16.39/v1")
@app.route('/retrieve', methods=['POST'])
def retrieve():
    data = request.get_json(force=True)
    original_query_str = data.get('original_query')
    query_list_str = data.get('query_list')
    data_set_list_str = data.get('data_set_list')
    try:
        query_list = query_list_str.split("<sub_query>")
        data_set_list = data_set_list_str.split("<dataset>")
        results = dify_query_retrieval.retrieve_api(original_query_str, query_list, data_set_list)
        return Response(json.dumps(results, ensure_ascii=False), content_type='application/json; charset=utf-8')
    except Exception as e:
        logger.error(f"检索出错: {str(e)}", exc_info=True)
        return Response(json.dumps({"error": str(e)}, ensure_ascii=False), content_type='application/json; charset=utf-8', status=500)


if __name__ == "__main__":
    # 开发环境使用Flask内置服务器
    # 生产环境使用gunicorn支持高并发 uv run gunicorn -w 10 -k gevent -b 0.0.0.0:8001 rag2_0.dify.DifyQueryRetrieval_api:app
    # uv run gunicorn -w 10 -k gevent --preload -b 0.0.0.0:8001 rag2_0.dify.DifyQueryRetrieval_api:app
    app.run(host="0.0.0.0", port=8002, threaded=True)