54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
# 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) |