feat(llm): timely execute vid embedding & enhance some HTTP logic

hugegraph-llm/requirements.txt

@@ -15,3 +15,4 @@ pyarrow~=17.0.0 # TODO: a temporary dependency for pandas, figure out why Import
 pandas~=2.2.2
 openpyxl~=3.1.5
 pydantic-settings~=2.6.1
+apscheduler~=3.10.4

hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py

@@ -17,12 +17,16 @@
 
 
 import argparse
+from contextlib import asynccontextmanager
 
 import gradio as gr
 import uvicorn
 from fastapi import FastAPI, Depends, APIRouter
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 
+from apscheduler.schedulers.background import BackgroundScheduler
+from apscheduler.triggers.cron import CronTrigger
+
 from hugegraph_llm.api.admin_api import admin_http_api
 from hugegraph_llm.api.rag_api import rag_http_api
 from hugegraph_llm.config import admin_settings, huge_settings, prompt
@@ -40,10 +44,10 @@
 from hugegraph_llm.demo.rag_demo.vector_graph_block import create_vector_graph_block
 from hugegraph_llm.resources.demo.css import CSS
 from hugegraph_llm.utils.log import log
+from hugegraph_llm.utils.graph_index_utils import fit_vid_index
 
 sec = HTTPBearer()
 
-
 def authenticate(credentials: HTTPAuthorizationCredentials = Depends(sec)):
     correct_token = admin_settings.user_token
     if credentials.credentials != correct_token:
@@ -55,6 +59,33 @@ def authenticate(credentials: HTTPAuthorizationCredentials = Depends(sec)):
             headers={"WWW-Authenticate": "Bearer"},
         )
 
+def schedule_fit_vid_index():
+    try:
+        log.info("Executing fit_vid_index function...")
+        fit_vid_index()
+        log.info("fit_vid_index function executed successfully.")
+    except Exception as e: #pylint: disable=W0718
+        log.error(e)
+
+scheduler = BackgroundScheduler()
+
+scheduler.add_job(
+    schedule_fit_vid_index,
+    CronTrigger(hour=2, minute=0),
+    id="fit_vid_index_job",
+    replace_existing=True,
+)
+
+@asynccontextmanager
+async def lifespan(app: FastAPI): #pylint: disable=W0621
+    if not scheduler.running:
+        scheduler.start()
+        log.info("Scheduler started successfully.")
+
+    yield
+
+    scheduler.shutdown()
+    log.info("Scheduler shut down.")
 
 # pylint: disable=C0301
 def init_rag_ui() -> gr.Interface:
@@ -158,7 +189,7 @@ def refresh_ui_config_prompt() -> tuple:
     parser.add_argument("--host", type=str, default="0.0.0.0", help="host")
     parser.add_argument("--port", type=int, default=8001, help="port")
     args = parser.parse_args()
-    app = FastAPI()
+    app = FastAPI(lifespan=lifespan)
 
     # we don't need to manually check the env now
     # settings.check_env()

hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py

@@ -85,9 +85,9 @@ def create_vector_graph_block():
                 graph_index_btn1 = gr.Button("Clear Graph Data & Index", size="sm")
 
         vector_import_bt = gr.Button("Import into Vector", variant="primary")
-        graph_index_rebuild_bt = gr.Button("Rebuild vid Index")
         graph_extract_bt = gr.Button("Extract Graph Data (1)", variant="primary")
         graph_loading_bt = gr.Button("Load into GraphDB (2)", interactive=True)
+        graph_index_rebuild_bt = gr.Button("Rebuild vid Index")
 
     vector_index_btn0.click(get_vector_index_info, outputs=out).then(
         store_prompt,
@@ -122,6 +122,9 @@ def create_vector_graph_block():
     graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then(
         store_prompt,
         inputs=[input_schema, info_extract_template],
+    ).then(fit_vid_index).then(
+        store_prompt,
+        inputs=[input_schema, info_extract_template],
     )
 
     def on_tab_select(input_f, input_t, evt: gr.SelectData):