Merge branch 'add-diagnostic-tool'

lightrag/tools/check_initialization.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+"""
+Diagnostic tool to check LightRAG initialization status.
+
+This tool helps developers verify that their LightRAG instance is properly
+initialized before use, preventing common initialization errors.
+
+Usage:
+    python -m lightrag.tools.check_initialization
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+# Add parent directory to path for imports
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+from lightrag import LightRAG
+from lightrag.base import StoragesStatus
+
+
+async def check_lightrag_setup(rag_instance: LightRAG, verbose: bool = False) -> bool:
+    """
+    Check if a LightRAG instance is properly initialized.
+
+    Args:
+        rag_instance: The LightRAG instance to check
+        verbose: If True, print detailed diagnostic information
+
+    Returns:
+        True if properly initialized, False otherwise
+    """
+    issues = []
+    warnings = []
+
+    print("🔍 Checking LightRAG initialization status...\n")
+
+    # Check storage initialization status
+    if not hasattr(rag_instance, "_storages_status"):
+        issues.append("LightRAG instance missing _storages_status attribute")
+    elif rag_instance._storages_status != StoragesStatus.INITIALIZED:
+        issues.append(
+            f"Storages not initialized (status: {rag_instance._storages_status.name})"
+        )
+    else:
+        print("✅ Storage status: INITIALIZED")
+
+    # Check individual storage components
+    storage_components = [
+        ("full_docs", "Document storage"),
+        ("text_chunks", "Text chunks storage"),
+        ("entities_vdb", "Entity vector database"),
+        ("relationships_vdb", "Relationship vector database"),
+        ("chunks_vdb", "Chunks vector database"),
+        ("doc_status", "Document status tracker"),
+        ("llm_response_cache", "LLM response cache"),
+        ("full_entities", "Entity storage"),
+        ("full_relations", "Relation storage"),
+        ("chunk_entity_relation_graph", "Graph storage"),
+    ]
+
+    if verbose:
+        print("\n📦 Storage Components:")
+
+    for component, description in storage_components:
+        if not hasattr(rag_instance, component):
+            issues.append(f"Missing storage component: {component} ({description})")
+        else:
+            storage = getattr(rag_instance, component)
+            if storage is None:
+                warnings.append(f"Storage {component} is None (might be optional)")
+            elif hasattr(storage, "_storage_lock"):
+                if storage._storage_lock is None:
+                    issues.append(f"Storage {component} not initialized (lock is None)")
+                elif verbose:
+                    print(f"  ✅ {description}: Ready")
+            elif verbose:
+                print(f"  ✅ {description}: Ready")
+
+    # Check pipeline status
+    try:
+        from lightrag.kg.shared_storage import get_namespace_data
+
+        get_namespace_data("pipeline_status")
+        print("✅ Pipeline status: INITIALIZED")
+    except KeyError:
+        issues.append(
+            "Pipeline status not initialized - call initialize_pipeline_status()"
+        )
+    except Exception as e:
+        issues.append(f"Error checking pipeline status: {str(e)}")
+
+    # Print results
+    print("\n" + "=" * 50)
+
+    if issues:
+        print("❌ Issues found:\n")
+        for issue in issues:
+            print(f"  • {issue}")
+
+        print("\n📝 To fix, run this initialization sequence:\n")
+        print("  await rag.initialize_storages()")
+        print("  from lightrag.kg.shared_storage import initialize_pipeline_status")
+        print("  await initialize_pipeline_status()")
+        print(
+            "\n📚 Documentation: https://github.com/HKUDS/LightRAG#important-initialization-requirements"
+        )
+
+        if warnings and verbose:
+            print("\n⚠️  Warnings (might be normal):")
+            for warning in warnings:
+                print(f"  • {warning}")
+
+        return False
+    else:
+        print("✅ LightRAG is properly initialized and ready to use!")
+
+        if warnings and verbose:
+            print("\n⚠️  Warnings (might be normal):")
+            for warning in warnings:
+                print(f"  • {warning}")
+
+        return True
+
+
+async def demo():
+    """Demonstrate the diagnostic tool with a test instance."""
+    from lightrag.llm.openai import openai_embed, gpt_4o_mini_complete
+    from lightrag.kg.shared_storage import initialize_pipeline_status
+
+    print("=" * 50)
+    print("LightRAG Initialization Diagnostic Tool")
+    print("=" * 50)
+
+    # Create test instance
+    rag = LightRAG(
+        working_dir="./test_diagnostic",
+        embedding_func=openai_embed,
+        llm_model_func=gpt_4o_mini_complete,
+    )
+
+    print("\n🔴 BEFORE initialization:\n")
+    await check_lightrag_setup(rag, verbose=True)
+
+    print("\n" + "=" * 50)
+    print("\n🔄 Initializing...\n")
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    print("\n🟢 AFTER initialization:\n")
+    await check_lightrag_setup(rag, verbose=True)
+
+    # Cleanup
+    import shutil
+
+    shutil.rmtree("./test_diagnostic", ignore_errors=True)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Check LightRAG initialization status")
+    parser.add_argument(
+        "--demo", action="store_true", help="Run a demonstration with a test instance"
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Show detailed diagnostic information",
+    )
+
+    args = parser.parse_args()
+
+    if args.demo:
+        asyncio.run(demo())
+    else:
+        print("Run with --demo to see the diagnostic tool in action")
+        print("Or import this module and use check_lightrag_setup() with your instance")