fix demo

.gitignore

@@ -57,7 +57,7 @@ ignore_this.txt
 *.ignore.*
 
 # Project-specific files
-dickens/
+dickens*/
 book.txt
 lightrag-dev/
 gui/

README.md

@@ -102,33 +102,47 @@ Use the below Python snippet (in a script) to initialize LightRAG and perform qu
 
 ```python
 import os
+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
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir="your/path",
         embedding_func=openai_embed,
         llm_model_func=gpt_4o_mini_complete
-)
+    )
 
-# Insert text
+    await rag.initialize_storages()
-rag.insert("Your text")
+    await initialize_pipeline_status()
 
-# Perform naive search
+    return rag
-mode="naive"
-# Perform local search
-mode="local"
-# Perform global search
-mode="global"
-# Perform hybrid search
-mode="hybrid"
-# Mix mode Integrates knowledge graph and vector retrieval.
-mode="mix"
 
-rag.query(
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+    # Insert text
+    rag.insert("Your text")
+
+    # Perform naive search
+    mode="naive"
+    # Perform local search
+    mode="local"
+    # Perform global search
+    mode="global"
+    # Perform hybrid search
+    mode="hybrid"
+    # Mix mode Integrates knowledge graph and vector retrieval.
+    mode="mix"
+
+    rag.query(
         "What are the top themes in this story?",
         param=QueryParam(mode=mode)
-)
+    )
+
+if __name__ == "__main__":
+    main()
 ```
 
 ### Query Param
@@ -190,7 +204,8 @@ async def embedding_func(texts: list[str]) -> np.ndarray:
         base_url="https://api.upstage.ai/v1/solar"
     )
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         embedding_func=EmbeddingFunc(
@@ -198,7 +213,12 @@ rag = LightRAG(
             max_token_size=8192,
             func=embedding_func
         )
-)
+    )
+
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
 ```
 </details>
 
@@ -210,10 +230,6 @@ rag = LightRAG(
 See `lightrag_hf_demo.py`
 
 ```python
-from lightrag.llm import hf_model_complete, hf_embed
-from transformers import AutoModel, AutoTokenizer
-from lightrag.utils import EmbeddingFunc
-
 # Initialize LightRAG with Hugging Face model
 rag = LightRAG(
     working_dir=WORKING_DIR,
@@ -242,9 +258,6 @@ If you want to use Ollama models, you need to pull model you plan to use and emb
 Then you only need to set LightRAG as follows:
 
 ```python
-from lightrag.llm.ollama import ollama_model_complete, ollama_embed
-from lightrag.utils import EmbeddingFunc
-
 # Initialize LightRAG with Ollama model
 rag = LightRAG(
     working_dir=WORKING_DIR,
@@ -325,12 +338,15 @@ LightRAG supports integration with LlamaIndex.
 
 ```python
 # Using LlamaIndex with direct OpenAI access
+import asyncio
 from lightrag import LightRAG
 from lightrag.llm.llama_index_impl import llama_index_complete_if_cache, llama_index_embed
 from llama_index.embeddings.openai import OpenAIEmbedding
 from llama_index.llms.openai import OpenAI
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir="your/path",
         llm_model_func=llama_index_complete_if_cache,  # LlamaIndex-compatible completion function
         embedding_func=EmbeddingFunc(    # LlamaIndex-compatible embedding function
@@ -338,7 +354,42 @@ rag = LightRAG(
             max_token_size=8192,
             func=lambda texts: llama_index_embed(texts, embed_model=embed_model)
         ),
-)
+    )
+
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
+        rag.insert(f.read())
+
+    # Perform naive search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+    )
+
+    # Perform local search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+    )
+
+    # Perform global search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+    )
+
+    # Perform hybrid search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+    )
+
+if __name__ == "__main__":
+    main()
 ```
 
 #### For detailed documentation and examples, see:
@@ -353,11 +404,6 @@ rag = LightRAG(
 LightRAG now supports multi-turn dialogue through the conversation history feature. Here's how to use it:
 
 ```python
-from lightrag import LightRAG, QueryParam
-
-# Initialize LightRAG
-rag = LightRAG(working_dir=WORKING_DIR)
-
 # Create conversation history
 conversation_history = [
     {"role": "user", "content": "What is the main character's attitude towards Christmas?"},
@@ -387,11 +433,6 @@ response = rag.query(
 LightRAG now supports custom prompts for fine-tuned control over the system's behavior. Here's how to use it:
 
 ```python
-from lightrag import LightRAG, QueryParam
-
-# Initialize LightRAG
-rag = LightRAG(working_dir=WORKING_DIR)
-
 # Create query parameters
 query_param = QueryParam(
     mode="hybrid",  # or other mode: "local", "global", "hybrid", "mix" and "naive"
@@ -456,16 +497,6 @@ rag.query_with_separate_keyword_extraction(
 <summary> <b>Insert Custom KG</b> </summary>
 
 ```python
-rag = LightRAG(
-     working_dir=WORKING_DIR,
-     llm_model_func=llm_model_func,
-     embedding_func=EmbeddingFunc(
-          embedding_dim=embedding_dimension,
-          max_token_size=8192,
-          func=embedding_func,
-     ),
-)
-
 custom_kg = {
     "entities": [
         {
@@ -534,6 +565,7 @@ rag = LightRAG(
         "insert_batch_size": 20  # Process 20 documents per batch
     }
 )
+
 rag.insert(["TEXT1", "TEXT2", "TEXT3", ...])  # Documents will be processed in batches of 20
 ```
 
@@ -560,27 +592,6 @@ rag.insert(["TEXT1", "TEXT2",...], ids=["ID_FOR_TEXT1", "ID_FOR_TEXT2"])
 
 </details>
 
-<details>
-  <summary><b>Incremental Insert</b></summary>
-
-```python
-# Incremental Insert: Insert new documents into an existing LightRAG instance
-rag = LightRAG(
-     working_dir=WORKING_DIR,
-     llm_model_func=llm_model_func,
-     embedding_func=EmbeddingFunc(
-          embedding_dim=embedding_dimension,
-          max_token_size=8192,
-          func=embedding_func,
-     ),
-)
-
-with open("./newText.txt") as f:
-    rag.insert(f.read())
-```
-
-</details>
-
 <details>
   <summary><b>Insert using Pipeline</b></summary>
 
@@ -592,6 +603,7 @@ And using a routine to process news documents.
 
 ```python
 rag = LightRAG(..)
+
 await rag.apipeline_enqueue_documents(input)
 # Your routine in loop
 await rag.apipeline_process_enqueue_documents(input)
@@ -633,8 +645,6 @@ export NEO4J_PASSWORD="password"
 
 # Note: Default settings use NetworkX
 # Initialize LightRAG with Neo4J implementation.
-WORKING_DIR = "./local_neo4jWorkDir"
-
 rag = LightRAG(
     working_dir=WORKING_DIR,
     llm_model_func=gpt_4o_mini_complete,  # Use gpt_4o_mini_complete LLM model
@@ -706,14 +716,14 @@ You can also install `faiss-gpu` if you have GPU support.
 
 - Here we are using `sentence-transformers` but you can also use `OpenAIEmbedding` model with `3072` dimensions.
 
-```
+```python
 async def embedding_func(texts: list[str]) -> np.ndarray:
     model = SentenceTransformer('all-MiniLM-L6-v2')
     embeddings = model.encode(texts, convert_to_numpy=True)
     return embeddings
 
 # Initialize LightRAG with the LLM model function and embedding function
-    rag = LightRAG(
+rag = LightRAG(
     working_dir=WORKING_DIR,
     llm_model_func=llm_model_func,
     embedding_func=EmbeddingFunc(
@@ -725,7 +735,7 @@ async def embedding_func(texts: list[str]) -> np.ndarray:
     vector_db_storage_cls_kwargs={
         "cosine_better_than_threshold": 0.3  # Your desired threshold
     }
-    )
+)
 ```
 
 </details>
@@ -733,17 +743,6 @@ async def embedding_func(texts: list[str]) -> np.ndarray:
 ## Delete
 
 ```python
-
-rag = LightRAG(
-     working_dir=WORKING_DIR,
-     llm_model_func=llm_model_func,
-     embedding_func=EmbeddingFunc(
-          embedding_dim=embedding_dimension,
-          max_token_size=8192,
-          func=embedding_func,
-     ),
-)
-
 #  Delete Entity: Deleting entities by their names
 rag.delete_by_entity("Project Gutenberg")
 

examples/copy_llm_cache_to_another_storage.py

@@ -10,7 +10,7 @@ import os
 from dotenv import load_dotenv
 
 from lightrag.kg.postgres_impl import PostgreSQLDB, PGKVStorage
-from lightrag.storage import JsonKVStorage
+from lightrag.kg.json_kv_impl import JsonKVStorage
 from lightrag.namespace import NameSpace
 
 load_dotenv()

examples/lightrag_api_ollama_demo.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, HTTPException, File, UploadFile
+from contextlib import asynccontextmanager
 from pydantic import BaseModel
 import os
 from lightrag import LightRAG, QueryParam
@@ -8,12 +9,12 @@ from typing import Optional
 import asyncio
 import nest_asyncio
 import aiofiles
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # Apply nest_asyncio to solve event loop issues
 nest_asyncio.apply()
 
 DEFAULT_RAG_DIR = "index_default"
-app = FastAPI(title="LightRAG API", description="API for RAG operations")
 
 DEFAULT_INPUT_FILE = "book.txt"
 INPUT_FILE = os.environ.get("INPUT_FILE", f"{DEFAULT_INPUT_FILE}")
@@ -28,7 +29,8 @@ if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
 
-rag = LightRAG(
+async def init():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=ollama_model_complete,
         llm_model_name="gemma2:9b",
@@ -42,9 +44,26 @@ rag = LightRAG(
                 texts, embed_model="nomic-embed-text", host="http://localhost:11434"
             ),
         ),
+    )
+
+    # Add initialization code
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global rag
+    rag = await init()
+    print("done!")
+    yield
+
+
+app = FastAPI(
+    title="LightRAG API", description="API for RAG operations", lifespan=lifespan
 )
-
-
 # Data models
 class QueryRequest(BaseModel):
     query: str

examples/lightrag_api_openai_compatible_demo.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, HTTPException, File, UploadFile
+from contextlib import asynccontextmanager
 from pydantic import BaseModel
 import os
 from lightrag import LightRAG, QueryParam
@@ -8,6 +9,7 @@ import numpy as np
 from typing import Optional
 import asyncio
 import nest_asyncio
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # Apply nest_asyncio to solve event loop issues
 nest_asyncio.apply()
@@ -71,16 +73,35 @@ async def get_embedding_dim():
 
 
 # Initialize RAG instance
-rag = LightRAG(
+async def init():
+    embedding_dimension = await get_embedding_dim()
+    
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         embedding_func=EmbeddingFunc(
-        embedding_dim=asyncio.run(get_embedding_dim()),
+            embedding_dim=embedding_dimension,
             max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
             func=embedding_func,
         ),
-)
+    )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+    
+    return rag
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global rag
+    rag = await init()
+    print("done!")
+    yield
+
+
+app = FastAPI(
+    title="LightRAG API", description="API for RAG operations", lifespan=lifespan
+)
 
 # Data models
 

examples/lightrag_api_openai_compatible_demo_simplified.py

@@ -1,101 +0,0 @@
-import os
-from lightrag import LightRAG, QueryParam
-from lightrag.llm.openai import openai_complete_if_cache, openai_embed
-from lightrag.utils import EmbeddingFunc
-import numpy as np
-import asyncio
-import nest_asyncio
-
-# Apply nest_asyncio to solve event loop issues
-nest_asyncio.apply()
-
-DEFAULT_RAG_DIR = "index_default"
-
-# Configure working directory
-WORKING_DIR = os.environ.get("RAG_DIR", f"{DEFAULT_RAG_DIR}")
-print(f"WORKING_DIR: {WORKING_DIR}")
-LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4o-mini")
-print(f"LLM_MODEL: {LLM_MODEL}")
-EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-small")
-print(f"EMBEDDING_MODEL: {EMBEDDING_MODEL}")
-EMBEDDING_MAX_TOKEN_SIZE = int(os.environ.get("EMBEDDING_MAX_TOKEN_SIZE", 8192))
-print(f"EMBEDDING_MAX_TOKEN_SIZE: {EMBEDDING_MAX_TOKEN_SIZE}")
-BASE_URL = os.environ.get("BASE_URL", "https://api.openai.com/v1")
-print(f"BASE_URL: {BASE_URL}")
-API_KEY = os.environ.get("API_KEY", "xxxxxxxx")
-print(f"API_KEY: {API_KEY}")
-
-if not os.path.exists(WORKING_DIR):
-    os.mkdir(WORKING_DIR)
-
-
-# LLM model function
-
-
-async def llm_model_func(
-    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
-) -> str:
-    return await openai_complete_if_cache(
-        model=LLM_MODEL,
-        prompt=prompt,
-        system_prompt=system_prompt,
-        history_messages=history_messages,
-        base_url=BASE_URL,
-        api_key=API_KEY,
-        **kwargs,
-    )
-
-
-# Embedding function
-
-
-async def embedding_func(texts: list[str]) -> np.ndarray:
-    return await openai_embed(
-        texts=texts,
-        model=EMBEDDING_MODEL,
-        base_url=BASE_URL,
-        api_key=API_KEY,
-    )
-
-
-async def get_embedding_dim():
-    test_text = ["This is a test sentence."]
-    embedding = await embedding_func(test_text)
-    embedding_dim = embedding.shape[1]
-    print(f"{embedding_dim=}")
-    return embedding_dim
-
-
-# Initialize RAG instance
-rag = LightRAG(
-    working_dir=WORKING_DIR,
-    llm_model_func=llm_model_func,
-    embedding_func=EmbeddingFunc(
-        embedding_dim=asyncio.run(get_embedding_dim()),
-        max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
-        func=embedding_func,
-    ),
-)
-
-with open("./book.txt", "r", encoding="utf-8") as f:
-    rag.insert(f.read())
-
-# Perform naive search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
-
-# Perform local search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
-
-# Perform global search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
-
-# Perform hybrid search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)

examples/lightrag_api_oracle_demo.py

@@ -16,6 +16,7 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import openai_complete_if_cache, openai_embed
 from lightrag.utils import EmbeddingFunc
 import numpy as np
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 
 print(os.getcwd())
@@ -113,6 +114,9 @@ async def init():
         vector_storage="OracleVectorDBStorage",
     )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
     return rag
 
 

examples/lightrag_azure_openai_demo.py

@@ -6,6 +6,7 @@ import numpy as np
 from dotenv import load_dotenv
 import logging
 from openai import AzureOpenAI
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 logging.basicConfig(level=logging.INFO)
 
@@ -90,6 +91,9 @@ rag = LightRAG(
     ),
 )
 
+rag.initialize_storages()
+initialize_pipeline_status()
+
 book1 = open("./book_1.txt", encoding="utf-8")
 book2 = open("./book_2.txt", encoding="utf-8")
 

examples/lightrag_bedrock_demo.py

@@ -8,6 +8,12 @@ import logging
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.bedrock import bedrock_complete, bedrock_embed
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
+
+import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
 
 logging.getLogger("aiobotocore").setLevel(logging.WARNING)
 
@@ -15,19 +21,28 @@ WORKING_DIR = "./dickens"
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=bedrock_complete,
         llm_model_name="Anthropic Claude 3 Haiku // Amazon Bedrock",
         embedding_func=EmbeddingFunc(
             embedding_dim=1024, max_token_size=8192, func=bedrock_embed
         ),
-)
+    )
 
-with open("./book.txt", "r", encoding="utf-8") as f:
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+    
+    return rag
+
+def main():
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-for mode in ["naive", "local", "global", "hybrid"]:
+    for mode in ["naive", "local", "global", "hybrid"]:
         print("\n+-" + "-" * len(mode) + "-+")
         print(f"| {mode.capitalize()} |")
         print("+-" + "-" * len(mode) + "-+\n")

examples/lightrag_gemini_demo.py

@@ -8,6 +8,12 @@ from dotenv import load_dotenv
 from lightrag.utils import EmbeddingFunc
 from lightrag import LightRAG, QueryParam
 from sentence_transformers import SentenceTransformer
+from lightrag.kg.shared_storage import initialize_pipeline_status
+
+import asyncio
+import nest_asyncio
+# Apply nest_asyncio to solve event loop issues
+nest_asyncio.apply()
 
 load_dotenv()
 gemini_api_key = os.getenv("GEMINI_API_KEY")
@@ -60,7 +66,8 @@ async def embedding_func(texts: list[str]) -> np.ndarray:
     return embeddings
 
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         embedding_func=EmbeddingFunc(
@@ -68,17 +75,28 @@ rag = LightRAG(
             max_token_size=8192,
             func=embedding_func,
         ),
-)
+    )
 
-file_path = "story.txt"
+    await rag.initialize_storages()
-with open(file_path, "r") as file:
+    await initialize_pipeline_status()
+    
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+    file_path = "story.txt"
+    with open(file_path, "r") as file:
         text = file.read()
 
-rag.insert(text)
+    rag.insert(text)
 
-response = rag.query(
+    response = rag.query(
         query="What is the main theme of the story?",
         param=QueryParam(mode="hybrid", top_k=5, response_type="single line"),
-)
+    )
 
-print(response)
+    print(response)
+
+if __name__ == "__main__":
+    main()

examples/lightrag_hf_demo.py

@@ -4,13 +4,20 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm.hf import hf_model_complete, hf_embed
 from lightrag.utils import EmbeddingFunc
 from transformers import AutoModel, AutoTokenizer
+from lightrag.kg.shared_storage import initialize_pipeline_status
+
+import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
 
 WORKING_DIR = "./dickens"
 
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=hf_model_complete,
         llm_model_name="meta-llama/Llama-3.1-8B-Instruct",
@@ -27,28 +34,38 @@ rag = LightRAG(
                 ),
             ),
         ),
-)
+    )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
 
-with open("./book.txt", "r", encoding="utf-8") as f:
+    return rag
+
+def main():
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Perform naive search
+    # Perform naive search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-# Perform local search
+    # Perform local search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-# Perform global search
+    # Perform global search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-# Perform hybrid search
+    # Perform hybrid search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/lightrag_jinaai_demo.py

@@ -1,115 +0,0 @@
-import numpy as np
-from lightrag import LightRAG, QueryParam
-from lightrag.utils import EmbeddingFunc
-from lightrag.llm.jina import jina_embed
-from lightrag.llm.openai import openai_complete_if_cache
-import os
-import asyncio
-
-
-async def embedding_func(texts: list[str]) -> np.ndarray:
-    return await jina_embed(texts, api_key="YourJinaAPIKey")
-
-
-WORKING_DIR = "./dickens"
-
-if not os.path.exists(WORKING_DIR):
-    os.mkdir(WORKING_DIR)
-
-
-async def llm_model_func(
-    prompt, system_prompt=None, history_messages=[], **kwargs
-) -> str:
-    return await openai_complete_if_cache(
-        "solar-mini",
-        prompt,
-        system_prompt=system_prompt,
-        history_messages=history_messages,
-        api_key=os.getenv("UPSTAGE_API_KEY"),
-        base_url="https://api.upstage.ai/v1/solar",
-        **kwargs,
-    )
-
-
-rag = LightRAG(
-    working_dir=WORKING_DIR,
-    llm_model_func=llm_model_func,
-    embedding_func=EmbeddingFunc(
-        embedding_dim=1024, max_token_size=8192, func=embedding_func
-    ),
-)
-
-
-async def lightraginsert(file_path, semaphore):
-    async with semaphore:
-        try:
-            with open(file_path, "r", encoding="utf-8") as f:
-                content = f.read()
-        except UnicodeDecodeError:
-            # If UTF-8 decoding fails, try other encodings
-            with open(file_path, "r", encoding="gbk") as f:
-                content = f.read()
-        await rag.ainsert(content)
-
-
-async def process_files(directory, concurrency_limit):
-    semaphore = asyncio.Semaphore(concurrency_limit)
-    tasks = []
-    for root, dirs, files in os.walk(directory):
-        for f in files:
-            file_path = os.path.join(root, f)
-            if f.startswith("."):
-                continue
-            tasks.append(lightraginsert(file_path, semaphore))
-    await asyncio.gather(*tasks)
-
-
-async def main():
-    try:
-        rag = LightRAG(
-            working_dir=WORKING_DIR,
-            llm_model_func=llm_model_func,
-            embedding_func=EmbeddingFunc(
-                embedding_dim=1024,
-                max_token_size=8192,
-                func=embedding_func,
-            ),
-        )
-
-        asyncio.run(process_files(WORKING_DIR, concurrency_limit=4))
-
-        # Perform naive search
-        print(
-            await rag.aquery(
-                "What are the top themes in this story?", param=QueryParam(mode="naive")
-            )
-        )
-
-        # Perform local search
-        print(
-            await rag.aquery(
-                "What are the top themes in this story?", param=QueryParam(mode="local")
-            )
-        )
-
-        # Perform global search
-        print(
-            await rag.aquery(
-                "What are the top themes in this story?",
-                param=QueryParam(mode="global"),
-            )
-        )
-
-        # Perform hybrid search
-        print(
-            await rag.aquery(
-                "What are the top themes in this story?",
-                param=QueryParam(mode="hybrid"),
-            )
-        )
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

examples/lightrag_llamaindex_direct_demo.py

@@ -8,6 +8,11 @@ from lightrag.utils import EmbeddingFunc
 from llama_index.llms.openai import OpenAI
 from llama_index.embeddings.openai import OpenAIEmbedding
 import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
+
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # Configure working directory
 WORKING_DIR = "./index_default"
@@ -76,38 +81,53 @@ async def get_embedding_dim():
     return embedding_dim
 
 
-# Initialize RAG instance
+async def initialize_rag():
-rag = LightRAG(
+    embedding_dimension = await get_embedding_dim()
+    
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         embedding_func=EmbeddingFunc(
-        embedding_dim=asyncio.run(get_embedding_dim()),
+            embedding_dim=embedding_dimension,
             max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
             func=embedding_func,
         ),
-)
+    )
 
-# Insert example text
+    await rag.initialize_storages()
-with open("./book.txt", "r", encoding="utf-8") as f:
+    await initialize_pipeline_status()
+    
+    return rag
+
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    # Insert example text
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Test different query modes
+    # Test different query modes
-print("\nNaive Search:")
+    print("\nNaive Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-print("\nLocal Search:")
+    print("\nLocal Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-print("\nGlobal Search:")
+    print("\nGlobal Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-print("\nHybrid Search:")
+    print("\nHybrid Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/lightrag_llamaindex_litellm_demo.py

@@ -8,6 +8,11 @@ from lightrag.utils import EmbeddingFunc
 from llama_index.llms.litellm import LiteLLM
 from llama_index.embeddings.litellm import LiteLLMEmbedding
 import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
+
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # Configure working directory
 WORKING_DIR = "./index_default"
@@ -79,38 +84,53 @@ async def get_embedding_dim():
     return embedding_dim
 
 
-# Initialize RAG instance
+async def initialize_rag():
-rag = LightRAG(
+    embedding_dimension = await get_embedding_dim()
+    
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         embedding_func=EmbeddingFunc(
-        embedding_dim=asyncio.run(get_embedding_dim()),
+            embedding_dim=embedding_dimension,
             max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
             func=embedding_func,
         ),
-)
+    )
 
-# Insert example text
+    await rag.initialize_storages()
-with open("./book.txt", "r", encoding="utf-8") as f:
+    await initialize_pipeline_status()
+    
+    return rag
+
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    # Insert example text
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Test different query modes
+    # Test different query modes
-print("\nNaive Search:")
+    print("\nNaive Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-print("\nLocal Search:")
+    print("\nLocal Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-print("\nGlobal Search:")
+    print("\nGlobal Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-print("\nHybrid Search:")
+    print("\nHybrid Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/lightrag_lmdeploy_demo.py

@@ -5,6 +5,12 @@ from lightrag.llm.lmdeploy import lmdeploy_model_if_cache
 from lightrag.llm.hf import hf_embed
 from lightrag.utils import EmbeddingFunc
 from transformers import AutoModel, AutoTokenizer
+from lightrag.kg.shared_storage import initialize_pipeline_status
+
+import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
 
 WORKING_DIR = "./dickens"
 
@@ -35,8 +41,8 @@ async def lmdeploy_model_complete(
         **kwargs,
     )
 
-
+async def initialize_rag():
-rag = LightRAG(
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=lmdeploy_model_complete,
         llm_model_name="meta-llama/Llama-3.1-8B-Instruct",  # please use definite path for local model
@@ -53,28 +59,41 @@ rag = LightRAG(
                 ),
             ),
         ),
-)
+    )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
     
-with open("./book.txt", "r", encoding="utf-8") as f:
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    # Insert example text
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Perform naive search
+    # Test different query modes
-print(
+    print("\nNaive Search:")
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-# Perform local search
+    print("\nLocal Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-# Perform global search
+    print("\nGlobal Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-# Perform hybrid search
+    print("\nHybrid Search:")
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/lightrag_nvidia_demo.py

@@ -1,5 +1,9 @@
 import os
 import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
+
 from lightrag import LightRAG, QueryParam
 from lightrag.llm import (
     openai_complete_if_cache,
@@ -7,10 +11,12 @@ from lightrag.llm import (
 )
 from lightrag.utils import EmbeddingFunc
 import numpy as np
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # for custom llm_model_func
 from lightrag.utils import locate_json_string_body_from_string
 
+
 WORKING_DIR = "./dickens"
 
 if not os.path.exists(WORKING_DIR):
@@ -91,9 +97,7 @@ async def test_funcs():
 
 # asyncio.run(test_funcs())
 
-
+async def initialize_rag():
-async def main():
-    try:
     embedding_dimension = await get_embedding_dim()
     print(f"Detected embedding dimension: {embedding_dimension}")
 
@@ -111,22 +115,19 @@ async def main():
         ),
     )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+    
+    return rag
+async def main():
+    try:
+        # Initialize RAG instance
+        rag = asyncio.run(initialize_rag())
+
         # reading file
         with open("./book.txt", "r", encoding="utf-8") as f:
             await rag.ainsert(f.read())
 
-        # redefine rag to change embedding into query type
-        rag = LightRAG(
-            working_dir=WORKING_DIR,
-            llm_model_func=llm_model_func,
-            # llm_model_name="meta/llama3-70b-instruct", #un comment if
-            embedding_func=EmbeddingFunc(
-                embedding_dim=embedding_dimension,
-                max_token_size=512,
-                func=query_embedding_func,
-            ),
-        )
-
         # Perform naive search
         print("==============Naive===============")
         print(

examples/lightrag_ollama_age_demo.py

@@ -1,4 +1,8 @@
 import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
+
 import inspect
 import logging
 import os
@@ -6,6 +10,7 @@ import os
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.ollama import ollama_embed, ollama_model_complete
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens_age"
 
@@ -22,7 +27,8 @@ os.environ["AGE_POSTGRES_HOST"] = "localhost"
 os.environ["AGE_POSTGRES_PORT"] = "5455"
 os.environ["AGE_GRAPH_NAME"] = "dickens"
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=ollama_model_complete,
         llm_model_name="llama3.1:8b",
@@ -37,44 +43,56 @@ rag = LightRAG(
             ),
         ),
         graph_storage="AGEStorage",
-)
+    )
 
-with open("./book.txt", "r", encoding="utf-8") as f:
+    await rag.initialize_storages()
-    rag.insert(f.read())
+    await initialize_pipeline_status()
-
-# Perform naive search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
-
-# Perform local search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
-
-# Perform global search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
-
-# Perform hybrid search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
-
-# stream response
-resp = rag.query(
-    "What are the top themes in this story?",
-    param=QueryParam(mode="hybrid", stream=True),
-)
     
+    return rag
 
 async def print_stream(stream):
     async for chunk in stream:
         print(chunk, end="", flush=True)
 
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
 
-if inspect.isasyncgen(resp):
+    # Insert example text
+    with open("./book.txt", "r", encoding="utf-8") as f:
+        rag.insert(f.read())
+
+    # Test different query modes
+    print("\nNaive Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+    )
+
+    print("\nLocal Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+    )
+
+    print("\nGlobal Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+    )
+
+    print("\nHybrid Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+    )
+
+    # stream response
+    resp = rag.query(
+        "What are the top themes in this story?",
+        param=QueryParam(mode="hybrid", stream=True),
+    )
+
+    if inspect.isasyncgen(resp):
         asyncio.run(print_stream(resp))
-else:
+    else:
         print(resp)
+
+if __name__ == "__main__":
+    main()

examples/lightrag_ollama_demo.py

@@ -1,10 +1,14 @@
 import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
 import os
 import inspect
 import logging
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.ollama import ollama_model_complete, ollama_embed
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens"
 
@@ -13,7 +17,8 @@ logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=ollama_model_complete,
         llm_model_name="gemma2:2b",
@@ -27,44 +32,56 @@ rag = LightRAG(
                 texts, embed_model="nomic-embed-text", host="http://localhost:11434"
             ),
         ),
-)
+    )
 
-with open("./book.txt", "r", encoding="utf-8") as f:
+    await rag.initialize_storages()
-    rag.insert(f.read())
+    await initialize_pipeline_status()
-
-# Perform naive search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
-
-# Perform local search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
-
-# Perform global search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
-
-# Perform hybrid search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
-
-# stream response
-resp = rag.query(
-    "What are the top themes in this story?",
-    param=QueryParam(mode="hybrid", stream=True),
-)
     
+    return rag
 
 async def print_stream(stream):
     async for chunk in stream:
         print(chunk, end="", flush=True)
 
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
 
-if inspect.isasyncgen(resp):
+    # Insert example text
+    with open("./book.txt", "r", encoding="utf-8") as f:
+        rag.insert(f.read())
+
+    # Test different query modes
+    print("\nNaive Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+    )
+
+    print("\nLocal Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+    )
+
+    print("\nGlobal Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+    )
+
+    print("\nHybrid Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+    )
+
+    # stream response
+    resp = rag.query(
+        "What are the top themes in this story?",
+        param=QueryParam(mode="hybrid", stream=True),
+    )
+
+    if inspect.isasyncgen(resp):
         asyncio.run(print_stream(resp))
-else:
+    else:
         print(resp)
+
+if __name__ == "__main__":
+    main()

examples/lightrag_ollama_gremlin_demo.py

@@ -12,6 +12,7 @@ import os
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.ollama import ollama_embed, ollama_model_complete
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens_gremlin"
 
@@ -31,7 +32,8 @@ os.environ["GREMLIN_TRAVERSE_SOURCE"] = "g"
 os.environ["GREMLIN_USER"] = ""
 os.environ["GREMLIN_PASSWORD"] = ""
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=ollama_model_complete,
         llm_model_name="llama3.1:8b",
@@ -46,44 +48,56 @@ rag = LightRAG(
             ),
         ),
         graph_storage="GremlinStorage",
-)
+    )
 
-with open("./book.txt", "r", encoding="utf-8") as f:
+    await rag.initialize_storages()
-    rag.insert(f.read())
+    await initialize_pipeline_status()
-
-# Perform naive search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
-
-# Perform local search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
-
-# Perform global search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
-
-# Perform hybrid search
-print(
-    rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
-
-# stream response
-resp = rag.query(
-    "What are the top themes in this story?",
-    param=QueryParam(mode="hybrid", stream=True),
-)
     
+    return rag
 
 async def print_stream(stream):
     async for chunk in stream:
         print(chunk, end="", flush=True)
 
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
 
-if inspect.isasyncgen(resp):
+    # Insert example text
+    with open("./book.txt", "r", encoding="utf-8") as f:
+        rag.insert(f.read())
+
+    # Test different query modes
+    print("\nNaive Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+    )
+
+    print("\nLocal Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+    )
+
+    print("\nGlobal Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+    )
+
+    print("\nHybrid Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+    )
+
+    # stream response
+    resp = rag.query(
+        "What are the top themes in this story?",
+        param=QueryParam(mode="hybrid", stream=True),
+    )
+
+    if inspect.isasyncgen(resp):
         asyncio.run(print_stream(resp))
-else:
+    else:
         print(resp)
+
+if __name__ == "__main__":
+    main()

examples/lightrag_ollama_neo4j_milvus_mongo_demo.py

@@ -2,6 +2,11 @@ import os
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.ollama import ollama_model_complete, ollama_embed
 from lightrag.utils import EmbeddingFunc
+import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # WorkingDir
 ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -27,8 +32,8 @@ os.environ["MILVUS_USER"] = "root"
 os.environ["MILVUS_PASSWORD"] = "root"
 os.environ["MILVUS_DB_NAME"] = "lightrag"
 
-
+async def initialize_rag():
-rag = LightRAG(
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=ollama_model_complete,
         llm_model_name="qwen2.5:14b",
@@ -45,12 +50,41 @@ rag = LightRAG(
         kv_storage="MongoKVStorage",
         graph_storage="Neo4JStorage",
         vector_storage="MilvusVectorDBStorage",
-)
+    )
 
-file = "./book.txt"
+    await rag.initialize_storages()
-with open(file, "r") as f:
+    await initialize_pipeline_status()
+    
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    # Insert example text
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-print(
+    # Test different query modes
+    print("\nNaive Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+    )
+
+    print("\nLocal Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+    )
+
+    print("\nGlobal Search:")
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+    )
+
+    print("\nHybrid Search:")
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/lightrag_openai_compatible_demo.py

@@ -4,6 +4,7 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import openai_complete_if_cache, openai_embed
 from lightrag.utils import EmbeddingFunc
 import numpy as np
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens"
 
@@ -52,9 +53,7 @@ async def test_funcs():
 
 # asyncio.run(test_funcs())
 
-
+async def initialize_rag():
-async def main():
-    try:
     embedding_dimension = await get_embedding_dim()
     print(f"Detected embedding dimension: {embedding_dimension}")
 
@@ -68,6 +67,15 @@ async def main():
         ),
     )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+async def main():
+    try:
+        # Initialize RAG instance
+        rag = asyncio.run(initialize_rag())
+
         with open("./book.txt", "r", encoding="utf-8") as f:
             await rag.ainsert(f.read())
 

examples/lightrag_openai_compatible_demo_embedding_cache.py

@@ -4,6 +4,7 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import openai_complete_if_cache, openai_embed
 from lightrag.utils import EmbeddingFunc
 import numpy as np
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens"
 
@@ -52,9 +53,7 @@ async def test_funcs():
 
 # asyncio.run(test_funcs())
 
-
+async def initialize_rag():
-async def main():
-    try:
     embedding_dimension = await get_embedding_dim()
     print(f"Detected embedding dimension: {embedding_dimension}")
 
@@ -72,6 +71,16 @@ async def main():
         ),
     )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+async def main():
+    try:
+        # Initialize RAG instance
+        rag = asyncio.run(initialize_rag())
+
         with open("./book.txt", "r", encoding="utf-8") as f:
             await rag.ainsert(f.read())
 

examples/lightrag_openai_compatible_stream_demo.py

@@ -1,9 +1,11 @@
 import inspect
 import os
+import asyncio
 from lightrag import LightRAG
 from lightrag.llm import openai_complete, openai_embed
 from lightrag.utils import EmbeddingFunc, always_get_an_event_loop
 from lightrag import QueryParam
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # WorkingDir
 ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -13,7 +15,8 @@ if not os.path.exists(WORKING_DIR):
 print(f"WorkingDir: {WORKING_DIR}")
 
 api_key = "empty"
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=openai_complete,
         llm_model_name="qwen2.5-14b-instruct@4bit",
@@ -30,25 +33,36 @@ rag = LightRAG(
                 api_key=api_key,
             ),
         ),
-)
+    )
 
-with open("./book.txt", "r", encoding="utf-8") as f:
+    await rag.initialize_storages()
-    rag.insert(f.read())
+    await initialize_pipeline_status()
-
-resp = rag.query(
-    "What are the top themes in this story?",
-    param=QueryParam(mode="hybrid", stream=True),
-)
 
+    return rag
 
 async def print_stream(stream):
     async for chunk in stream:
         if chunk:
             print(chunk, end="", flush=True)
 
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
 
-loop = always_get_an_event_loop()
+    with open("./book.txt", "r", encoding="utf-8") as f:
-if inspect.isasyncgen(resp):
+        rag.insert(f.read())
+
+    resp = rag.query(
+        "What are the top themes in this story?",
+        param=QueryParam(mode="hybrid", stream=True),
+    )
+
+    loop = always_get_an_event_loop()
+    if inspect.isasyncgen(resp):
         loop.run_until_complete(print_stream(resp))
-else:
+    else:
         print(resp)
+
+if __name__ == "__main__":
+    main()
+

examples/lightrag_openai_demo.py

@@ -1,40 +1,54 @@
 import os
-
+import asyncio
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens"
 
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         embedding_func=openai_embed,
         llm_model_func=gpt_4o_mini_complete,
         # llm_model_func=gpt_4o_complete
-)
+    )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
 
-with open("./book.txt", "r", encoding="utf-8") as f:
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Perform naive search
+    # Perform naive search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-# Perform local search
+    # Perform local search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-# Perform global search
+    # Perform global search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-# Perform hybrid search
+    # Perform hybrid search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()
+

examples/lightrag_openai_mongodb_graph_demo.py

@@ -4,6 +4,7 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
 from lightrag.utils import EmbeddingFunc
 import numpy as np
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 #########
 # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
@@ -52,7 +53,7 @@ async def create_embedding_function_instance():
 async def initialize_rag():
     embedding_func_instance = await create_embedding_function_instance()
 
-    return LightRAG(
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=gpt_4o_mini_complete,
         embedding_func=embedding_func_instance,
@@ -60,14 +61,38 @@ async def initialize_rag():
         log_level="DEBUG",
     )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
 
-# Run the initialization
+    return rag
-rag = asyncio.run(initialize_rag())
 
-with open("book.txt", "r", encoding="utf-8") as f:
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Perform naive search
+    # Perform naive search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
+
+    # Perform local search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+    )
+
+    # Perform global search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+    )
+
+    # Perform hybrid search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+    )
+
+if __name__ == "__main__":
+    main()

examples/lightrag_openai_neo4j_milvus_redis_demo.py

@@ -1,7 +1,9 @@
 import os
+import asyncio
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.ollama import ollama_embed, openai_complete_if_cache
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # WorkingDir
 ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -48,8 +50,8 @@ embedding_func = EmbeddingFunc(
         texts, embed_model="shaw/dmeta-embedding-zh", host="http://117.50.173.35:11434"
     ),
 )
-
+async def initialize_rag():
-rag = LightRAG(
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         llm_model_max_token_size=32768,
@@ -60,11 +62,40 @@ rag = LightRAG(
         graph_storage="Neo4JStorage",
         vector_storage="MilvusVectorDBStorage",
         doc_status_storage="RedisKVStorage",
-)
+    )
 
-file = "../book.txt"
+    await rag.initialize_storages()
-with open(file, "r", encoding="utf-8") as f:
+    await initialize_pipeline_status()
+
+    return rag
+
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
+    # Perform naive search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+    )
 
-print(rag.query("谁会3D建模 ？", param=QueryParam(mode="mix")))
+    # Perform local search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+    )
+
+    # Perform global search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+    )
+
+    # Perform hybrid search
+    print(
+        rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+    )
+
+if __name__ == "__main__":
+    main()

examples/lightrag_oracle_demo.py

@@ -6,6 +6,7 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import openai_complete_if_cache, openai_embed
 from lightrag.utils import EmbeddingFunc
 import numpy as np
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 print(os.getcwd())
 script_directory = Path(__file__).resolve().parent.parent
@@ -63,9 +64,7 @@ async def get_embedding_dim():
     embedding_dim = embedding.shape[1]
     return embedding_dim
 
-
+async def initialize_rag():
-async def main():
-    try:
     # Detect embedding dimension
     embedding_dimension = await get_embedding_dim()
     print(f"Detected embedding dimension: {embedding_dimension}")
@@ -98,6 +97,15 @@ async def main():
             "insert_batch_size": 2,
         },
     )
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+async def main():
+    try:
+        # Initialize RAG instance
+        rag = asyncio.run(initialize_rag())
 
         # Extract and Insert into LightRAG storage
         with open(WORKING_DIR + "/docs.txt", "r", encoding="utf-8") as f:

examples/lightrag_siliconcloud_demo.py

@@ -5,6 +5,7 @@ from lightrag.llm.openai import openai_complete_if_cache
 from lightrag.llm.siliconcloud import siliconcloud_embedding
 from lightrag.utils import EmbeddingFunc
 import numpy as np
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens"
 
@@ -46,35 +47,48 @@ async def test_funcs():
 
 asyncio.run(test_funcs())
 
-
+async def initialize_rag():
-rag = LightRAG(
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         embedding_func=EmbeddingFunc(
             embedding_dim=768, max_token_size=512, func=embedding_func
         ),
-)
+    )
+
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
 
 
-with open("./book.txt") as f:
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Perform naive search
+    # Perform naive search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-# Perform local search
+    # Perform local search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-# Perform global search
+    # Perform global search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-# Perform hybrid search
+    # Perform hybrid search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()
+

examples/lightrag_tidb_demo.py

@@ -6,6 +6,7 @@ import numpy as np
 from lightrag import LightRAG, QueryParam
 from lightrag.llm import siliconcloud_embedding, openai_complete_if_cache
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens"
 
@@ -54,9 +55,7 @@ async def get_embedding_dim():
     embedding_dim = embedding.shape[1]
     return embedding_dim
 
-
+async def initialize_rag():
-async def main():
-    try:
     # Detect embedding dimension
     embedding_dimension = await get_embedding_dim()
     print(f"Detected embedding dimension: {embedding_dimension}")
@@ -78,9 +77,18 @@ async def main():
         graph_storage="TiDBGraphStorage",
     )
 
-        # Extract and Insert into LightRAG storage
+    await rag.initialize_storages()
-        with open("./dickens/demo.txt", "r", encoding="utf-8") as f:
+    await initialize_pipeline_status()
-            await rag.ainsert(f.read())
+
+    return rag
+
+async def main():
+    try:
+        # Initialize RAG instance
+        rag = asyncio.run(initialize_rag())
+
+        with open("./book.txt", "r", encoding="utf-8") as f:
+            rag.insert(f.read())
 
         # Perform search in different modes
         modes = ["naive", "local", "global", "hybrid"]

examples/lightrag_zhipu_demo.py

@@ -1,10 +1,12 @@
 import os
 import logging
+import asyncio
 
 
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.zhipu import zhipu_complete, zhipu_embedding
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 WORKING_DIR = "./dickens"
 
@@ -17,8 +19,8 @@ api_key = os.environ.get("ZHIPUAI_API_KEY")
 if api_key is None:
     raise Exception("Please set ZHIPU_API_KEY in your environment")
 
-
+async def initialize_rag():
-rag = LightRAG(
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=zhipu_complete,
         llm_model_name="glm-4-flashx",  # Using the most cost/performance balance model, but you can change it here.
@@ -29,27 +31,39 @@ rag = LightRAG(
             max_token_size=8192,
             func=lambda texts: zhipu_embedding(texts),
         ),
-)
+    )
 
-with open("./book.txt", "r", encoding="utf-8") as f:
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Perform naive search
+    # Perform naive search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-# Perform local search
+    # Perform local search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-# Perform global search
+    # Perform global search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-# Perform hybrid search
+    # Perform hybrid search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/lightrag_zhipu_postgres_demo.py

@@ -8,6 +8,7 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm.zhipu import zhipu_complete
 from lightrag.llm.ollama import ollama_embedding
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 load_dotenv()
 ROOT_DIR = os.environ.get("ROOT_DIR")
@@ -27,8 +28,7 @@ os.environ["POSTGRES_USER"] = "rag"
 os.environ["POSTGRES_PASSWORD"] = "rag"
 os.environ["POSTGRES_DATABASE"] = "rag"
 
-
+async def initialize_rag():
-async def main():
     rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=zhipu_complete,
@@ -50,9 +50,17 @@ async def main():
         auto_manage_storages_states=False,
     )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+async def main():
+        # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
     # add embedding_func for graph database, it's deleted in commit 5661d76860436f7bf5aef2e50d9ee4a59660146c
     rag.chunk_entity_relation_graph.embedding_func = rag.embedding_func
-    await rag.initialize_storages()
 
     with open(f"{ROOT_DIR}/book.txt", "r", encoding="utf-8") as f:
         await rag.ainsert(f.read())

examples/query_keyword_separation_example.py

@@ -6,6 +6,7 @@ import numpy as np
 from dotenv import load_dotenv
 import logging
 from openai import AzureOpenAI
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 logging.basicConfig(level=logging.INFO)
 
@@ -79,8 +80,8 @@ async def test_funcs():
 asyncio.run(test_funcs())
 
 embedding_dimension = 3072
-
+async def initialize_rag():
-rag = LightRAG(
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         embedding_func=EmbeddingFunc(
@@ -88,16 +89,23 @@ rag = LightRAG(
             max_token_size=8192,
             func=embedding_func,
         ),
-)
+    )
 
-book1 = open("./book_1.txt", encoding="utf-8")
+    await rag.initialize_storages()
-book2 = open("./book_2.txt", encoding="utf-8")
+    await initialize_pipeline_status()
 
-rag.insert([book1.read(), book2.read()])
+    return rag
 
 
 # Example function demonstrating the new query_with_separate_keyword_extraction usage
 async def run_example():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+    
+    book1 = open("./book_1.txt", encoding="utf-8")
+    book2 = open("./book_2.txt", encoding="utf-8")
+
+    rag.insert([book1.read(), book2.read()])
     query = "What are the top themes in this story?"
     prompt = "Please simplify the response for a young audience."
 

examples/test.py

@@ -1,6 +1,7 @@
 import os
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import gpt_4o_mini_complete
+from lightrag.kg.shared_storage import initialize_pipeline_status
 #########
 # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
 # import nest_asyncio
@@ -12,31 +13,45 @@ WORKING_DIR = "./dickens"
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=gpt_4o_mini_complete,  # Use gpt_4o_mini_complete LLM model
         # llm_model_func=gpt_4o_complete  # Optionally, use a stronger model
-)
+    )
 
-with open("./dickens/book.txt", "r", encoding="utf-8") as f:
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Perform naive search
+    # Perform naive search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-# Perform local search
+    # Perform local search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-# Perform global search
+    # Perform global search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-# Perform hybrid search
+    # Perform hybrid search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/test_chromadb.py

@@ -4,6 +4,7 @@ from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
 from lightrag.utils import EmbeddingFunc
 import numpy as np
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 #########
 # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
@@ -67,7 +68,7 @@ async def create_embedding_function_instance():
 async def initialize_rag():
     embedding_func_instance = await create_embedding_function_instance()
     if CHROMADB_USE_LOCAL_PERSISTENT:
-        return LightRAG(
+        rag = LightRAG(
             working_dir=WORKING_DIR,
             llm_model_func=gpt_4o_mini_complete,
             embedding_func=embedding_func_instance,
@@ -87,7 +88,7 @@ async def initialize_rag():
             },
         )
     else:
-        return LightRAG(
+        rag = LightRAG(
             working_dir=WORKING_DIR,
             llm_model_func=gpt_4o_mini_complete,
             embedding_func=embedding_func_instance,
@@ -112,28 +113,36 @@ async def initialize_rag():
         )
 
 
-# Run the initialization
+    await rag.initialize_storages()
-rag = asyncio.run(initialize_rag())
+    await initialize_pipeline_status()
 
-# with open("./dickens/book.txt", "r", encoding="utf-8") as f:
+    return rag
-#     rag.insert(f.read())
 
-# Perform naive search
+    # Initialize RAG instance
-print(
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
+        rag.insert(f.read())
+
+    # Perform naive search
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-# Perform local search
+    # Perform local search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-# Perform global search
+    # Perform global search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-# Perform hybrid search
+    # Perform hybrid search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/test_faiss.py

@@ -1,5 +1,6 @@
 import os
 import logging
+import asyncio
 import numpy as np
 
 from dotenv import load_dotenv
@@ -8,7 +9,9 @@ from sentence_transformers import SentenceTransformer
 from openai import AzureOpenAI
 from lightrag import LightRAG, QueryParam
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
+WORKING_DIR = "./dickens"
 # Configure Logging
 logging.basicConfig(level=logging.INFO)
 
@@ -55,11 +58,7 @@ async def embedding_func(texts: list[str]) -> np.ndarray:
     embeddings = model.encode(texts, convert_to_numpy=True)
     return embeddings
 
-
+async def initialize_rag():
-def main():
-    WORKING_DIR = "./dickens"
-
-    # Initialize LightRAG with the LLM model function and embedding function
     rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
@@ -74,6 +73,15 @@ def main():
         },
     )
 
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+def main():
+    
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
     # Insert the custom chunks into LightRAG
     book1 = open("./book_1.txt", encoding="utf-8")
     book2 = open("./book_2.txt", encoding="utf-8")

examples/test_neo4j.py

@@ -1,7 +1,8 @@
 import os
+import asyncio
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.openai import gpt_4o_mini_complete
-
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 #########
 # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
@@ -14,33 +15,46 @@ WORKING_DIR = "./local_neo4jWorkDir"
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=gpt_4o_mini_complete,  # Use gpt_4o_mini_complete LLM model
         graph_storage="Neo4JStorage",
         log_level="INFO",
         # llm_model_func=gpt_4o_complete  # Optionally, use a stronger model
-)
+    )
 
-with open("./book.txt") as f:
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+
+    with open("./book.txt", "r", encoding="utf-8") as f:
         rag.insert(f.read())
 
-# Perform naive search
+    # Perform naive search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
-)
+    )
 
-# Perform local search
+    # Perform local search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
-)
+    )
 
-# Perform global search
+    # Perform global search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
-)
+    )
 
-# Perform hybrid search
+    # Perform hybrid search
-print(
+    print(
         rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
-)
+    )
+
+if __name__ == "__main__":
+    main()

examples/test_split_by_character.ipynb

@@ -18,6 +18,7 @@
     "from lightrag import LightRAG, QueryParam\n",
     "from lightrag.llm.openai import openai_complete_if_cache, openai_embed\n",
     "from lightrag.utils import EmbeddingFunc\n",
+    "from lightrag.kg.shared_storage import initialize_pipeline_status\n",
     "import nest_asyncio"
    ]
   },
@@ -25,7 +26,9 @@
    "cell_type": "markdown",
    "id": "dd17956ec322b361",
    "metadata": {},
-   "source": "#### split by character"
+   "source": [
+    "#### split by character"
+   ]
   },
   {
    "cell_type": "code",
@@ -109,14 +112,26 @@
     }
    ],
    "source": [
-    "rag = LightRAG(\n",
+    "import asyncio\n",
+    "import nest_asyncio\n",
+    "\n",
+    "nest_asyncio.apply()\n",
+    "\n",
+    "async def initialize_rag():\n",
+    "    rag = LightRAG(\n",
     "        working_dir=WORKING_DIR,\n",
     "        llm_model_func=llm_model_func,\n",
     "        embedding_func=EmbeddingFunc(\n",
     "            embedding_dim=4096, max_token_size=8192, func=embedding_func\n",
     "        ),\n",
     "        chunk_token_size=512,\n",
-    ")"
+    "    )\n",
+    "    await rag.initialize_storages()\n",
+    "    await initialize_pipeline_status()\n",
+    "\n",
+    "    return rag\n",
+    "\n",
+    "rag = asyncio.run(initialize_rag())"
    ]
   },
   {
@@ -908,7 +923,9 @@
    "cell_type": "markdown",
    "id": "4e5bfad24cb721a8",
    "metadata": {},
-   "source": "#### split by character only"
+   "source": [
+    "#### split by character only"
+   ]
   },
   {
    "cell_type": "code",

examples/vram_management_demo.py

@@ -1,8 +1,10 @@
 import os
 import time
+import asyncio
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.ollama import ollama_model_complete, ollama_embed
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # Working directory and the directory path for text files
 WORKING_DIR = "./dickens"
@@ -12,8 +14,9 @@ TEXT_FILES_DIR = "/llm/mt"
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-# Initialize LightRAG
+async def initialize_rag():
-rag = LightRAG(
+    # Initialize LightRAG
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=ollama_model_complete,
         llm_model_name="qwen2.5:3b-instruct-max-context",
@@ -22,7 +25,11 @@ rag = LightRAG(
             max_token_size=8192,
             func=lambda texts: ollama_embed(texts, embed_model="nomic-embed-text"),
         ),
-)
+    )
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
 
 # Read all .txt files from the TEXT_FILES_DIR directory
 texts = []
@@ -47,58 +54,65 @@ def insert_texts_with_retry(rag, texts, retries=3, delay=5):
     raise RuntimeError("Failed to insert texts after multiple retries.")
 
 
-insert_texts_with_retry(rag, texts)
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
 
-# Perform different types of queries and handle potential errors
+    insert_texts_with_retry(rag, texts)
-try:
+
+    # Perform different types of queries and handle potential errors
+    try:
         print(
             rag.query(
                 "What are the top themes in this story?", param=QueryParam(mode="naive")
             )
         )
-except Exception as e:
+    except Exception as e:
         print(f"Error performing naive search: {e}")
 
-try:
+    try:
         print(
             rag.query(
                 "What are the top themes in this story?", param=QueryParam(mode="local")
             )
         )
-except Exception as e:
+    except Exception as e:
         print(f"Error performing local search: {e}")
 
-try:
+    try:
         print(
             rag.query(
                 "What are the top themes in this story?", param=QueryParam(mode="global")
             )
         )
-except Exception as e:
+    except Exception as e:
         print(f"Error performing global search: {e}")
 
-try:
+    try:
         print(
             rag.query(
                 "What are the top themes in this story?", param=QueryParam(mode="hybrid")
             )
         )
-except Exception as e:
+    except Exception as e:
         print(f"Error performing hybrid search: {e}")
 
 
-# Function to clear VRAM resources
+    # Function to clear VRAM resources
-def clear_vram():
+    def clear_vram():
         os.system("sudo nvidia-smi --gpu-reset")
 
 
-# Regularly clear VRAM to prevent overflow
+    # Regularly clear VRAM to prevent overflow
-clear_vram_interval = 3600  # Clear once every hour
+    clear_vram_interval = 3600  # Clear once every hour
-start_time = time.time()
+    start_time = time.time()
 
-while True:
+    while True:
         current_time = time.time()
         if current_time - start_time > clear_vram_interval:
             clear_vram()
             start_time = current_time
         time.sleep(60)  # Check the time every minute
+
+if __name__ == "__main__":
+    main()

reproduce/Step_1.py

@@ -1,8 +1,10 @@
 import os
 import json
 import time
+import asyncio
 
 from lightrag import LightRAG
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 
 def insert_text(rag, file_path):
@@ -29,6 +31,19 @@ WORKING_DIR = f"../{cls}"
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-rag = LightRAG(working_dir=WORKING_DIR)
+async def initialize_rag():
+    rag = LightRAG(working_dir=WORKING_DIR)
 
-insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json")
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+    insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json")
+
+
+if __name__ == "__main__":
+    main()

reproduce/Step_1_openai_compatible.py

@@ -1,11 +1,13 @@
 import os
 import json
 import time
+import asyncio
 import numpy as np
 
 from lightrag import LightRAG
 from lightrag.utils import EmbeddingFunc
 from lightrag.llm.openai import openai_complete_if_cache, openai_embed
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 
 ## For Upstage API
@@ -60,12 +62,25 @@ WORKING_DIR = f"../{cls}"
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-rag = LightRAG(
+async def initialize_rag():
+    rag = LightRAG(
         working_dir=WORKING_DIR,
         llm_model_func=llm_model_func,
         embedding_func=EmbeddingFunc(
             embedding_dim=4096, max_token_size=8192, func=embedding_func
         ),
-)
+    )
 
-insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json")
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
+
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
+    insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json")
+
+
+if __name__ == "__main__":
+    main()

run_with_gunicorn.py

@@ -1,203 +0,0 @@
-#!/usr/bin/env python
-"""
-Start LightRAG server with Gunicorn
-"""
-
-import os
-import sys
-import signal
-import pipmaster as pm
-from lightrag.api.utils_api import parse_args, display_splash_screen
-from lightrag.kg.shared_storage import initialize_share_data, finalize_share_data
-
-
-def check_and_install_dependencies():
-    """Check and install required dependencies"""
-    required_packages = [
-        "gunicorn",
-        "tiktoken",
-        "psutil",
-        # Add other required packages here
-    ]
-
-    for package in required_packages:
-        if not pm.is_installed(package):
-            print(f"Installing {package}...")
-            pm.install(package)
-            print(f"{package} installed successfully")
-
-
-# Signal handler for graceful shutdown
-def signal_handler(sig, frame):
-    print("\n\n" + "=" * 80)
-    print("RECEIVED TERMINATION SIGNAL")
-    print(f"Process ID: {os.getpid()}")
-    print("=" * 80 + "\n")
-
-    # Release shared resources
-    finalize_share_data()
-
-    # Exit with success status
-    sys.exit(0)
-
-
-def main():
-    # Check and install dependencies
-    check_and_install_dependencies()
-
-    # Register signal handlers for graceful shutdown
-    signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
-    signal.signal(signal.SIGTERM, signal_handler)  # kill command
-
-    # Parse all arguments using parse_args
-    args = parse_args(is_uvicorn_mode=False)
-
-    # Display startup information
-    display_splash_screen(args)
-
-    print("🚀 Starting LightRAG with Gunicorn")
-    print(f"🔄 Worker management: Gunicorn (workers={args.workers})")
-    print("🔍 Preloading app: Enabled")
-    print("📝 Note: Using Gunicorn's preload feature for shared data initialization")
-    print("\n\n" + "=" * 80)
-    print("MAIN PROCESS INITIALIZATION")
-    print(f"Process ID: {os.getpid()}")
-    print(f"Workers setting: {args.workers}")
-    print("=" * 80 + "\n")
-
-    # Import Gunicorn's StandaloneApplication
-    from gunicorn.app.base import BaseApplication
-
-    # Define a custom application class that loads our config
-    class GunicornApp(BaseApplication):
-        def __init__(self, app, options=None):
-            self.options = options or {}
-            self.application = app
-            super().__init__()
-
-        def load_config(self):
-            # Define valid Gunicorn configuration options
-            valid_options = {
-                "bind",
-                "workers",
-                "worker_class",
-                "timeout",
-                "keepalive",
-                "preload_app",
-                "errorlog",
-                "accesslog",
-                "loglevel",
-                "certfile",
-                "keyfile",
-                "limit_request_line",
-                "limit_request_fields",
-                "limit_request_field_size",
-                "graceful_timeout",
-                "max_requests",
-                "max_requests_jitter",
-            }
-
-            # Special hooks that need to be set separately
-            special_hooks = {
-                "on_starting",
-                "on_reload",
-                "on_exit",
-                "pre_fork",
-                "post_fork",
-                "pre_exec",
-                "pre_request",
-                "post_request",
-                "worker_init",
-                "worker_exit",
-                "nworkers_changed",
-                "child_exit",
-            }
-
-            # Import and configure the gunicorn_config module
-            import gunicorn_config
-
-            # Set configuration variables in gunicorn_config, prioritizing command line arguments
-            gunicorn_config.workers = (
-                args.workers if args.workers else int(os.getenv("WORKERS", 1))
-            )
-
-            # Bind configuration prioritizes command line arguments
-            host = args.host if args.host != "0.0.0.0" else os.getenv("HOST", "0.0.0.0")
-            port = args.port if args.port != 9621 else int(os.getenv("PORT", 9621))
-            gunicorn_config.bind = f"{host}:{port}"
-
-            # Log level configuration prioritizes command line arguments
-            gunicorn_config.loglevel = (
-                args.log_level.lower()
-                if args.log_level
-                else os.getenv("LOG_LEVEL", "info")
-            )
-
-            # Timeout configuration prioritizes command line arguments
-            gunicorn_config.timeout = (
-                args.timeout if args.timeout else int(os.getenv("TIMEOUT", 150))
-            )
-
-            # Keepalive configuration
-            gunicorn_config.keepalive = int(os.getenv("KEEPALIVE", 5))
-
-            # SSL configuration prioritizes command line arguments
-            if args.ssl or os.getenv("SSL", "").lower() in (
-                "true",
-                "1",
-                "yes",
-                "t",
-                "on",
-            ):
-                gunicorn_config.certfile = (
-                    args.ssl_certfile
-                    if args.ssl_certfile
-                    else os.getenv("SSL_CERTFILE")
-                )
-                gunicorn_config.keyfile = (
-                    args.ssl_keyfile if args.ssl_keyfile else os.getenv("SSL_KEYFILE")
-                )
-
-            # Set configuration options from the module
-            for key in dir(gunicorn_config):
-                if key in valid_options:
-                    value = getattr(gunicorn_config, key)
-                    # Skip functions like on_starting and None values
-                    if not callable(value) and value is not None:
-                        self.cfg.set(key, value)
-                # Set special hooks
-                elif key in special_hooks:
-                    value = getattr(gunicorn_config, key)
-                    if callable(value):
-                        self.cfg.set(key, value)
-
-            if hasattr(gunicorn_config, "logconfig_dict"):
-                self.cfg.set(
-                    "logconfig_dict", getattr(gunicorn_config, "logconfig_dict")
-                )
-
-        def load(self):
-            # Import the application
-            from lightrag.api.lightrag_server import get_application
-
-            return get_application(args)
-
-    # Create the application
-    app = GunicornApp("")
-
-    # Force workers to be an integer and greater than 1 for multi-process mode
-    workers_count = int(args.workers)
-    if workers_count > 1:
-        # Set a flag to indicate we're in the main process
-        os.environ["LIGHTRAG_MAIN_PROCESS"] = "1"
-        initialize_share_data(workers_count)
-    else:
-        initialize_share_data(1)
-
-    # Run the application
-    print("\nStarting Gunicorn with direct Python API...")
-    app.run()
-
-
-if __name__ == "__main__":
-    main()