ST-Raptor/gradio_app.py at master · OpenDataBox/ST-Raptor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
import gradio as gr
import os
import sys
import time
import atexit
import json
import fastapi
import uvicorn
from fastapi import Request

from loguru import logger
from embedding import EmbeddingModel
from utils.api_utils import vlm_generate, llm_generate, embedding_generate
from query.primitive_pipeline import *
from table2tree.extract_excel import *
from table2tree.feature_tree import *
from utils.constants import DELIMITER, LOG_DIR

# 导入其他子文件中的函数
from config import load_api_config, save_api_config, api_config
from new_interface import create_interface

# 初始化时加载配置
load_api_config()

# 从core_functions.py导入核心功能函数
from core_functions import (
    answer_question,
    process_table_for_tree,
    process_question_only,
    clear_all,
    read_all_logs,
    get_llm_generate,
    get_vlm_generate,
    get_embedding_generate,
    rebuild_feature_tree_from_json,
)


def main():
    # 启动时仅清理日志文件（保留 cache/temp 等），避免残留旧日志干扰
    def clean_logs(log_dir="log"):
        try:
            if os.path.exists(log_dir):
                for root, dirs, files in os.walk(log_dir):
                    for fname in files:
                        fpath = os.path.join(root, fname)
                        try:
                            os.remove(fpath)
                        except Exception as e:
                            print(f"[WARN] 无法删除日志文件 {fpath}: {e}")
        except Exception as e:
            print(f"[WARN] 清理日志失败: {e}")

    clean_logs("log")

    print("🚀 启动 ST-Raptor Gradio 界面...")
    print("📋 访问地址: http://localhost:7860")
    print("⏹️  按 Ctrl+C 停止服务")

    def cleanup():
        try:
            clear_all()
        except Exception as e:
            print(f"[WARN] 清理失败: {e}")

    # 注册退出时的清理函数
    atexit.register(cleanup)

    demo = create_interface()
    demo.queue()  # 启用队列模式

    app = fastapi.FastAPI()

    @app.post("/save_tree")
    async def save_tree(request: Request):
        try:
            data = await request.json()
        except Exception as e:
            return {"status": "fail", "msg": f"invalid json: {e}"}
        logger.info(f"[save_tree] received save request, len={len(str(data))}")
        # 前端为选中逻辑添加的 id 不落盘，先剥离
        def strip_ids(obj):
            if isinstance(obj, list):
                return [strip_ids(o) for o in obj]
            if isinstance(obj, dict):
                return {k: strip_ids(v) for k, v in obj.items() if k != "id"}
            return obj

        cleaned = strip_ids(data)
        logger.info("[save_tree] saving cleaned tree and rebuilding feature views")
        ok, msg = rebuild_feature_tree_from_json(cleaned)
        if ok:
            return {"status": "ok"}
        return {"status": "fail", "msg": msg}

    # 将 gradio app 挂载到 fastapi
    app = gr.mount_gradio_app(app, demo, path="/")

    try:
        uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
    except (KeyboardInterrupt, Exception) as e:
        if not isinstance(e, KeyboardInterrupt):
            print(f"服务运行出错: {e}")
    finally:
        print("\n🛑 正在关闭界面并清理环境...")
        try:
            demo.close()  # 显式关闭 Gradio 界面
        except:
            pass
        # 稍微等一下让异步任务取消
        time.sleep(0.5)
        cleanup()
        print("✅ 服务已安全停止。")

if __name__ == "__main__":
    main()