rag_tools/main.py

import time
from flask import Flask, request, jsonify
import logging
import atexit
import os
import sys
import signal
from reader import Reader

class FlaskAppManager:
    def __init__(self):
        # 创建应用
        self.app = Flask(__name__)
        # 创建读取器(进程间独立，子线程间共享)
        self.reader = Reader()
        self.shutting_down = False
        self.setup_handlers()
        self.upload_dir = "./uploads"

    def setup_handlers(self):
        """设置所有退出处理程序"""
        # 信号处理
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)

        # atexit 处理
        atexit.register(self.cleanup)

    def signal_handler(self, signum, frame):
        """信号处理函数"""
        if self.shutting_down:
            return  # 已经在关闭中，忽略重复信号

        self.shutting_down = True
        signal_name = "SIGINT" if signum == signal.SIGINT else "SIGTERM"
        print(f"\n🚨 接收到 {signal_name} 信号，开始优雅关闭...")

        self.cleanup()
        sys.exit(0)  # 正常退出


    def cleanup(self):
        """退出Flask清理资源"""
        if hasattr(self, '_cleaned'):
            return

        self._cleaned = True
        print("🔧 清理资源中...")


        # 停止 Reader
        if hasattr(self.reader, 'stop'):
            print("停止文件读取器...")
            self.reader.stop()

        print("✅ 资源清理完成")

    def run(self,port=5000, debug=True, threaded=True):
        """运行应用"""
        # 启动Reader
        self.reader.start()
        print("🚀 启动 RAG Flask 应用...")
        print("💡 提示: 按 Ctrl+C 可以优雅关闭应用")
        # 创建文件目录
        if not os.path.exists(self.upload_dir):
            os.makedirs(self.upload_dir)

        try:
            self.app.run(port=port, debug=debug, threaded=threaded)
        except Exception as e:
            print(f"❌ 应用异常: {e}")
        finally:
            self.cleanup()

# 创建任务管理器
app_manager = FlaskAppManager()
@app_manager.app.route('/api/v1/rag_upload', methods=["POST"])
def rag_upload():
    response = {"message": "None"}
    # 接收文件
    if 'file' not in request.files:
        return {"message": "未上传文件"}, 400
    file = request.files['file']
    if file.filename == '':
        return {"message": "空文件名"}, 400
    # 提取文件类型
    #
    # 获取文件名（不包括路径）
    filename = os.path.basename(file.filename)
    file_path = os.path.join(app_manager.upload_dir, filename)
    # 保存文件
    file.save(file_path)
    # 向读取器添加任务
    if app_manager.reader.add_task(file_path):
        response = ({"message":f"文件【{file.filename}】上传成功"}, 200)
    else:
        response = ({"message":f"文件【{file.filename}】上传失败"}, 400)
    print('检查任务队列', app_manager.reader.task_list)
    return response

if __name__ == '__main__':
    app_manager.run(port=5000, debug=False, threaded=True)