From 071302d10fffa7c6f1730603d1e3673dc1b606ad Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Mon, 24 Mar 2025 05:37:03 +0800
Subject: [PATCH] Moved update status logic to document routes.

- Removed update status from health check endpoint
- Added update_status field to PipelineStatusResponse
---
 lightrag/api/lightrag_server.py         |  6 +-----
 lightrag/api/routers/document_routes.py | 10 +++++++++-
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py
index 2ebf0ca4..42bccd75 100644
--- a/lightrag/api/lightrag_server.py
+++ b/lightrag/api/lightrag_server.py
@@ -423,9 +423,6 @@ def create_app(args):
     @app.get("/health", dependencies=[Depends(optional_api_key_dependency)])
     async def get_status():
         """Get current system status"""
-        # Get update flags status for all namespaces
-        update_status = await get_all_update_flags_status()
-
         username = os.getenv("AUTH_USERNAME")
         password = os.getenv("AUTH_PASSWORD")
         if not (username and password):
@@ -453,11 +450,10 @@ def create_app(args):
                 "vector_storage": args.vector_storage,
                 "enable_llm_cache_for_extract": args.enable_llm_cache_for_extract,
             },
-            "update_status": update_status,
             "core_version": core_version,
             "api_version": __api_version__,
             "auth_mode": auth_mode,
-        }
+            }
 
     # Custom StaticFiles class to prevent caching of HTML files
     class NoCacheStaticFiles(StaticFiles):
diff --git a/lightrag/api/routers/document_routes.py b/lightrag/api/routers/document_routes.py
index 9de306be..f1fd694b 100644
--- a/lightrag/api/routers/document_routes.py
+++ b/lightrag/api/routers/document_routes.py
@@ -111,6 +111,7 @@ class PipelineStatusResponse(BaseModel):
         request_pending: Flag for pending request for processing
         latest_message: Latest message from pipeline processing
         history_messages: List of history messages
+        update_status: Status of update flags for all namespaces
     """
 
     autoscanned: bool = False
@@ -123,6 +124,7 @@ class PipelineStatusResponse(BaseModel):
     request_pending: bool = False
     latest_message: str = ""
     history_messages: Optional[List[str]] = None
+    update_status: Optional[dict] = None
 
     class Config:
         extra = "allow"  # Allow additional fields from the pipeline status
@@ -796,12 +798,18 @@ def create_document_routes(
             HTTPException: If an error occurs while retrieving pipeline status (500)
         """
         try:
-            from lightrag.kg.shared_storage import get_namespace_data
+            from lightrag.kg.shared_storage import get_namespace_data, get_all_update_flags_status
 
             pipeline_status = await get_namespace_data("pipeline_status")
+            
+            # Get update flags status for all namespaces
+            update_status = await get_all_update_flags_status()
 
             # Convert to regular dict if it's a Manager.dict
             status_dict = dict(pipeline_status)
+            
+            # Add update_status to the status dictionary
+            status_dict["update_status"] = update_status
 
             # Convert history_messages to a regular list if it's a Manager.list
             if "history_messages" in status_dict: