less comments

examples/lightrag_multi_model_all_modes_demo.py

@@ -3,22 +3,17 @@ import asyncio
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import gpt_4o_mini_complete, gpt_4o_complete, openai_embed
 from lightrag.kg.shared_storage import initialize_pipeline_status
-from lightrag.utils import setup_logger
 
-setup_logger("lightrag", level="INFO")
-
-WORKING_DIR = "./all_modes_demo"
+WORKING_DIR = "./lightrag_demo"
 
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-
 async def initialize_rag():
-    # Initialize LightRAG with a base model (gpt-4o-mini)
     rag = LightRAG(
         working_dir=WORKING_DIR,
         embedding_func=openai_embed,
-        llm_model_func=gpt_4o_mini_complete,  # Default model for most queries
+        llm_model_func=gpt_4o_mini_complete,  # Default model for queries
     )
 
     await rag.initialize_storages()
@@ -26,7 +21,6 @@ async def initialize_rag():
 
     return rag
 
-
 def main():
     # Initialize RAG instance
     rag = asyncio.run(initialize_rag())
@@ -34,60 +28,63 @@ def main():
     # Load the data
     with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
-    
-    # Example query
-    query_text = "What are the main themes in this story?"
-    
-    # Demonstrate using default model (gpt-4o-mini) for all modes
-    print("\n===== Default Model (gpt-4o-mini) =====")
-    
-    for mode in ["local", "global", "hybrid", "naive", "mix"]:
-        print(f"\n--- {mode.upper()} mode with default model ---")
-        response = rag.query(
-            query_text, 
-            param=QueryParam(mode=mode)
+
+    # Query with naive mode (default model)
+    print("--- NAIVE mode ---")
+    print(
+        rag.query(
+            "What are the main themes in this story?", 
+            param=QueryParam(mode="naive")
         )
-        print(response)
-    
-    # Demonstrate using custom model (gpt-4o) for all modes
-    print("\n===== Custom Model (gpt-4o) =====")
-    
-    for mode in ["local", "global", "hybrid", "naive", "mix"]:
-        print(f"\n--- {mode.upper()} mode with custom model ---")
-        response = rag.query(
-            query_text, 
+    )
+
+    # Query with local mode (default model)
+    print("\n--- LOCAL mode ---")
+    print(
+        rag.query(
+            "What are the main themes in this story?", 
+            param=QueryParam(mode="local")
+        )
+    )
+
+    # Query with global mode (default model)
+    print("\n--- GLOBAL mode ---")
+    print(
+        rag.query(
+            "What are the main themes in this story?", 
+            param=QueryParam(mode="global")
+        )
+    )
+
+    # Query with hybrid mode (default model)
+    print("\n--- HYBRID mode ---")
+    print(
+        rag.query(
+            "What are the main themes in this story?", 
+            param=QueryParam(mode="hybrid")
+        )
+    )
+
+    # Query with mix mode (default model)
+    print("\n--- MIX mode ---")
+    print(
+        rag.query(
+            "What are the main themes in this story?", 
+            param=QueryParam(mode="mix")
+        )
+    )
+
+    # Query with a custom model (gpt-4o) for a more complex question
+    print("\n--- Using custom model for complex analysis ---")
+    print(
+        rag.query(
+            "How does the character development reflect Victorian-era attitudes?",
             param=QueryParam(
-                mode=mode,
-                model_func=gpt_4o_complete  # Override with more capable model
+                mode="global",
+                model_func=gpt_4o_complete  # Override default model with more capable one
             )
         )
-        print(response)
-    
-    # Mixed approach - use different models for different modes
-    print("\n===== Strategic Model Selection =====")
-    
-    # Complex analytical question
-    complex_query = "How does the character development in the story reflect Victorian-era social values?"
-    
-    # Use default model for simpler modes
-    print("\n--- NAIVE mode with default model (suitable for simple retrieval) ---")
-    response1 = rag.query(
-        complex_query,
-        param=QueryParam(mode="naive")  # Use default model for basic retrieval
     )
-    print(response1)
-    
-    # Use more capable model for complex modes
-    print("\n--- HYBRID mode with more capable model (for complex analysis) ---")
-    response2 = rag.query(
-        complex_query,
-        param=QueryParam(
-            mode="hybrid",
-            model_func=gpt_4o_complete  # Use more capable model for complex analysis
-        )
-    )
-    print(response2)
-
 
 if __name__ == "__main__":
     main()