feat: add update flags status to API health endpoint

This commit is contained in:
yangdx
2025-03-01 14:58:26 +08:00
parent 41eff2ca2f
commit 40e9e26edb
2 changed files with 34 additions and 7 deletions

View File

@@ -39,6 +39,12 @@ from .routers.graph_routes import create_graph_routes
from .routers.ollama_api import OllamaAPI from .routers.ollama_api import OllamaAPI
from lightrag.utils import logger, set_verbose_debug from lightrag.utils import logger, set_verbose_debug
from lightrag.kg.shared_storage import (
get_namespace_data,
get_pipeline_status_lock,
initialize_pipeline_status,
get_all_update_flags_status,
)
# Load environment variables # Load environment variables
load_dotenv(override=True) load_dotenv(override=True)
@@ -134,13 +140,6 @@ def create_app(args):
try: try:
# Initialize database connections # Initialize database connections
await rag.initialize_storages() await rag.initialize_storages()
# Import necessary functions from shared_storage
from lightrag.kg.shared_storage import (
get_namespace_data,
get_pipeline_status_lock,
initialize_pipeline_status,
)
await initialize_pipeline_status() await initialize_pipeline_status()
# Auto scan documents if enabled # Auto scan documents if enabled
@@ -376,6 +375,9 @@ def create_app(args):
@app.get("/health", dependencies=[Depends(optional_api_key)]) @app.get("/health", dependencies=[Depends(optional_api_key)])
async def get_status(): async def get_status():
"""Get current system status""" """Get current system status"""
# Get update flags status for all namespaces
update_status = await get_all_update_flags_status()
return { return {
"status": "healthy", "status": "healthy",
"working_directory": str(args.working_dir), "working_directory": str(args.working_dir),
@@ -395,6 +397,7 @@ def create_app(args):
"graph_storage": args.graph_storage, "graph_storage": args.graph_storage,
"vector_storage": args.vector_storage, "vector_storage": args.vector_storage,
}, },
"update_status": update_status,
} }
# Webui mount webui/index.html # Webui mount webui/index.html

View File

@@ -222,6 +222,30 @@ async def set_all_update_flags(namespace: str):
_update_flags[namespace][i] = True _update_flags[namespace][i] = True
async def get_all_update_flags_status() -> Dict[str, list]:
"""
Get update flags status for all namespaces.
Returns:
Dict[str, list]: A dictionary mapping namespace names to lists of update flag statuses
"""
if _update_flags is None:
return {}
result = {}
async with get_internal_lock():
for namespace, flags in _update_flags.items():
worker_statuses = []
for flag in flags:
if is_multiprocess:
worker_statuses.append(flag.value)
else:
worker_statuses.append(flag)
result[namespace] = worker_statuses
return result
def try_initialize_namespace(namespace: str) -> bool: def try_initialize_namespace(namespace: str) -> bool:
""" """
Try to initialize a namespace. Returns True if the current process gets initialization permission. Try to initialize a namespace. Returns True if the current process gets initialization permission.