From 92ecb0da970432661fe708fbd6793afe3a83abb3 Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Thu, 27 Feb 2025 16:07:00 +0800
Subject: [PATCH] Refactor document scanning progress share variable
 initialization

---
 lightrag/api/routers/document_routes.py | 19 +++----------------
 lightrag/lightrag.py                    | 14 +++++++++++++-
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/lightrag/api/routers/document_routes.py b/lightrag/api/routers/document_routes.py
index 1f591750..5d90fb83 100644
--- a/lightrag/api/routers/document_routes.py
+++ b/lightrag/api/routers/document_routes.py
@@ -20,7 +20,7 @@ from lightrag import LightRAG
 from lightrag.base import DocProcessingStatus, DocStatus
 from ..utils_api import get_api_key_dependency
 from lightrag.kg.shared_storage import (
-    get_scan_progress,
+    get_namespace_data,
     get_scan_lock,
 )
 
@@ -376,21 +376,8 @@ async def save_temp_file(input_dir: Path, file: UploadFile = File(...)) -> Path:
 
 async def run_scanning_process(rag: LightRAG, doc_manager: DocumentManager):
     """Background task to scan and index documents"""
-    scan_progress = get_scan_progress()
+    scan_progress = get_namespace_data("scan_progress")
     scan_lock = get_scan_lock()
-
-    # Initialize scan_progress if not already initialized
-    if not scan_progress:
-        scan_progress.update(
-            {
-                "is_scanning": False,
-                "current_file": "",
-                "indexed_count": 0,
-                "total_files": 0,
-                "progress": 0,
-            }
-        )
-
     with scan_lock:
         if scan_progress.get("is_scanning", False):
             ASCIIColors.info("Skip document scanning(another scanning is active)")
@@ -491,7 +478,7 @@ def create_document_routes(
                 - total_files: Total number of files to process
                 - progress: Percentage of completion
         """
-        return dict(get_scan_progress())
+        return dict(get_namespace_data("scan_progress"))
 
     @router.post("/upload", dependencies=[Depends(optional_api_key)])
     async def upload_to_input_dir(
diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py
index 08ca202f..ec0accc3 100644
--- a/lightrag/lightrag.py
+++ b/lightrag/lightrag.py
@@ -267,8 +267,20 @@ class LightRAG:
     _storages_status: StoragesStatus = field(default=StoragesStatus.NOT_CREATED)
 
     def __post_init__(self):
-        from lightrag.kg.shared_storage import initialize_share_data
+        from lightrag.kg.shared_storage import initialize_share_data, try_initialize_namespace, get_namespace_data
         initialize_share_data()
+        need_init = try_initialize_namespace("scan_progress")
+        scan_progress = get_namespace_data("scan_progress")
+        if need_init:
+            scan_progress.update(
+                {
+                    "is_scanning": False,
+                    "current_file": "",
+                    "indexed_count": 0,
+                    "total_files": 0,
+                    "progress": 0,
+                }
+            )
 
         os.makedirs(os.path.dirname(self.log_file_path), exist_ok=True)
         set_logger(self.log_file_path, self.log_level)