chore: import upstream snapshot with attribution

examples/generate_query.py

@@ -0,0 +1,57 @@
+from openai import OpenAI
+
+# os.environ["OPENAI_API_KEY"] = ""
+
+
+def openai_complete_if_cache(
+    model="gpt-4o-mini", prompt=None, system_prompt=None, history_messages=[], **kwargs
+) -> str:
+    openai_client = OpenAI()
+
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+
+    response = openai_client.chat.completions.create(
+        model=model, messages=messages, **kwargs
+    )
+    if not response.choices or response.choices[0].message is None:
+        return ""
+    return response.choices[0].message.content
+
+
+if __name__ == "__main__":
+    description = ""
+    prompt = f"""
+    Given the following description of a dataset:
+
+    {description}
+
+    Please identify 5 potential users who would engage with this dataset. For each user, list 5 tasks they would perform with this dataset. Then, for each (user, task) combination, generate 5 questions that require a high-level understanding of the entire dataset.
+
+    Output the results in the following structure:
+    - User 1: [user description]
+        - Task 1: [task description]
+            - Question 1:
+            - Question 2:
+            - Question 3:
+            - Question 4:
+            - Question 5:
+        - Task 2: [task description]
+            ...
+        - Task 5: [task description]
+    - User 2: [user description]
+        ...
+    - User 5: [user description]
+        ...
+    """
+
+    result = openai_complete_if_cache(model="gpt-4o-mini", prompt=prompt)
+
+    file_path = "./queries.txt"
+    with open(file_path, "w") as file:
+        file.write(result)
+
+    print(f"Queries written to {file_path}")

examples/graph_visual_with_html.py

@@ -0,0 +1,34 @@
+import pipmaster as pm
+
+if not pm.is_installed("pyvis"):
+    pm.install("pyvis")
+if not pm.is_installed("networkx"):
+    pm.install("networkx")
+
+import networkx as nx
+from pyvis.network import Network
+import random
+
+# Load the GraphML file
+G = nx.read_graphml("./dickens/graph_chunk_entity_relation.graphml")
+
+# Create a Pyvis network
+net = Network(height="100vh", notebook=True)
+
+# Convert NetworkX graph to Pyvis network
+net.from_nx(G)
+
+
+# Add colors and title to nodes
+for node in net.nodes:
+    node["color"] = "#{:06x}".format(random.randint(0, 0xFFFFFF))
+    if "description" in node:
+        node["title"] = node["description"]
+
+# Add title to edges
+for edge in net.edges:
+    if "description" in edge:
+        edge["title"] = edge["description"]
+
+# Save and display the network
+net.show("knowledge_graph.html")

examples/graph_visual_with_neo4j.py

@@ -0,0 +1,186 @@
+import os
+import json
+import xml.etree.ElementTree as ET
+from neo4j import GraphDatabase
+
+# Constants
+WORKING_DIR = "./dickens"
+BATCH_SIZE_NODES = 500
+BATCH_SIZE_EDGES = 100
+
+# Neo4j connection credentials
+NEO4J_URI = "bolt://localhost:7687"
+NEO4J_USERNAME = "neo4j"
+NEO4J_PASSWORD = "your_password"
+
+
+def xml_to_json(xml_file):
+    try:
+        tree = ET.parse(xml_file)
+        root = tree.getroot()
+
+        # Print the root element's tag and attributes to confirm the file has been correctly loaded
+        print(f"Root element: {root.tag}")
+        print(f"Root attributes: {root.attrib}")
+
+        data = {"nodes": [], "edges": []}
+
+        # Use namespace
+        namespace = {"": "http://graphml.graphdrawing.org/xmlns"}
+
+        for node in root.findall(".//node", namespace):
+            node_data = {
+                "id": node.get("id").strip('"'),
+                "entity_type": node.find("./data[@key='d1']", namespace).text.strip('"')
+                if node.find("./data[@key='d1']", namespace) is not None
+                else "",
+                "description": node.find("./data[@key='d2']", namespace).text
+                if node.find("./data[@key='d2']", namespace) is not None
+                else "",
+                "source_id": node.find("./data[@key='d3']", namespace).text
+                if node.find("./data[@key='d3']", namespace) is not None
+                else "",
+            }
+            data["nodes"].append(node_data)
+
+        for edge in root.findall(".//edge", namespace):
+            edge_data = {
+                "source": edge.get("source").strip('"'),
+                "target": edge.get("target").strip('"'),
+                "weight": float(edge.find("./data[@key='d5']", namespace).text)
+                if edge.find("./data[@key='d5']", namespace) is not None
+                else 0.0,
+                "description": edge.find("./data[@key='d6']", namespace).text
+                if edge.find("./data[@key='d6']", namespace) is not None
+                else "",
+                "keywords": edge.find("./data[@key='d9']", namespace).text
+                if edge.find("./data[@key='d9']", namespace) is not None
+                else "",
+                "source_id": edge.find("./data[@key='d8']", namespace).text
+                if edge.find("./data[@key='d8']", namespace) is not None
+                else "",
+            }
+            data["edges"].append(edge_data)
+
+        # Print the number of nodes and edges found
+        print(f"Found {len(data['nodes'])} nodes and {len(data['edges'])} edges")
+
+        return data
+    except ET.ParseError as e:
+        print(f"Error parsing XML file: {e}")
+        return None
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return None
+
+
+def convert_xml_to_json(xml_path, output_path):
+    """Converts XML file to JSON and saves the output."""
+    if not os.path.exists(xml_path):
+        print(f"Error: File not found - {xml_path}")
+        return None
+
+    json_data = xml_to_json(xml_path)
+    if json_data:
+        with open(output_path, "w", encoding="utf-8") as f:
+            json.dump(json_data, f, ensure_ascii=False, indent=2)
+        print(f"JSON file created: {output_path}")
+        return json_data
+    else:
+        print("Failed to create JSON data")
+        return None
+
+
+def process_in_batches(tx, query, data, batch_size):
+    """Process data in batches and execute the given query."""
+    for i in range(0, len(data), batch_size):
+        batch = data[i : i + batch_size]
+        tx.run(query, {"nodes": batch} if "nodes" in query else {"edges": batch})
+
+
+def main():
+    # Paths
+    xml_file = os.path.join(WORKING_DIR, "graph_chunk_entity_relation.graphml")
+    json_file = os.path.join(WORKING_DIR, "graph_data.json")
+
+    # Convert XML to JSON
+    json_data = convert_xml_to_json(xml_file, json_file)
+    if json_data is None:
+        return
+
+    # Load nodes and edges
+    nodes = json_data.get("nodes", [])
+    edges = json_data.get("edges", [])
+
+    # Neo4j queries
+    create_nodes_query = """
+    UNWIND $nodes AS node
+    MERGE (e:Entity {id: node.id})
+    SET e.entity_type = node.entity_type,
+        e.description = node.description,
+        e.source_id = node.source_id,
+        e.displayName = node.id
+    REMOVE e:Entity
+    WITH e, node
+    CALL apoc.create.addLabels(e, [node.id]) YIELD node AS labeledNode
+    RETURN count(*)
+    """
+
+    create_edges_query = """
+    UNWIND $edges AS edge
+    MATCH (source {id: edge.source})
+    MATCH (target {id: edge.target})
+    WITH source, target, edge,
+         CASE
+            WHEN edge.keywords CONTAINS 'lead' THEN 'lead'
+            WHEN edge.keywords CONTAINS 'participate' THEN 'participate'
+            WHEN edge.keywords CONTAINS 'uses' THEN 'uses'
+            WHEN edge.keywords CONTAINS 'located' THEN 'located'
+            WHEN edge.keywords CONTAINS 'occurs' THEN 'occurs'
+           ELSE REPLACE(SPLIT(edge.keywords, ',')[0], '\"', '')
+         END AS relType
+    CALL apoc.create.relationship(source, relType, {
+      weight: edge.weight,
+      description: edge.description,
+      keywords: edge.keywords,
+      source_id: edge.source_id
+    }, target) YIELD rel
+    RETURN count(*)
+    """
+
+    set_displayname_and_labels_query = """
+    MATCH (n)
+    SET n.displayName = n.id
+    WITH n
+    CALL apoc.create.setLabels(n, [n.entity_type]) YIELD node
+    RETURN count(*)
+    """
+
+    # Create a Neo4j driver
+    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
+
+    try:
+        # Execute queries in batches
+        with driver.session() as session:
+            # Insert nodes in batches
+            session.execute_write(
+                process_in_batches, create_nodes_query, nodes, BATCH_SIZE_NODES
+            )
+
+            # Insert edges in batches
+            session.execute_write(
+                process_in_batches, create_edges_query, edges, BATCH_SIZE_EDGES
+            )
+
+            # Set displayName and labels
+            session.run(set_displayname_and_labels_query)
+
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
+    finally:
+        driver.close()
+
+
+if __name__ == "__main__":
+    main()

examples/graph_visual_with_opensearch.py

@@ -0,0 +1,166 @@
+"""
+Knowledge Graph Visualization with OpenSearch + LightRAG WebUI
+
+This script demonstrates two ways to visualize the knowledge graph
+stored in OpenSearch:
+
+1. **WebUI (recommended)**: Opens the LightRAG WebUI in your browser
+   for interactive graph exploration with search, filtering, and
+   force-directed layout.
+
+2. **Standalone HTML**: Fetches graph data from the LightRAG Server API
+   and generates an interactive HTML file using Pyvis, similar to
+   graph_visual_with_html.py but reading from OpenSearch instead of
+   a local .graphml file.
+
+Prerequisites:
+    1. LightRAG Server running with OpenSearch storage:
+       lightrag-server --host 0.0.0.0 --port 9621
+
+    2. Documents already indexed (e.g., via the WebUI or API)
+
+Usage:
+    # Open WebUI for interactive exploration
+    python examples/graph_visual_with_opensearch.py
+
+    # Generate standalone HTML file
+    python examples/graph_visual_with_opensearch.py --html
+
+    # Custom server URL and output file
+    python examples/graph_visual_with_opensearch.py --html --server http://localhost:9621 --output my_graph.html
+"""
+
+import argparse
+import os
+import sys
+import webbrowser
+
+import pipmaster as pm
+
+if not pm.is_installed("requests"):
+    pm.install("requests")
+if not pm.is_installed("pyvis"):
+    pm.install("pyvis")
+
+import requests
+from pyvis.network import Network
+
+
+def fetch_graph(server_url: str, label: str = "*", max_nodes: int = 300) -> dict:
+    """Fetch knowledge graph data from LightRAG Server API."""
+    url = f"{server_url}/graphs"
+    params = {"label": label, "max_nodes": max_nodes}
+    resp = requests.get(url, params=params, timeout=30)
+    resp.raise_for_status()
+    return resp.json()
+
+
+def generate_html(graph_data: dict, output_file: str) -> str:
+    """Generate an interactive HTML visualization from graph data."""
+    nodes = graph_data.get("nodes", [])
+    edges = graph_data.get("edges", [])
+
+    if not nodes:
+        print("No nodes found in the graph. Index some documents first.")
+        sys.exit(1)
+
+    print(f"Building visualization: {len(nodes)} nodes, {len(edges)} edges")
+
+    net = Network(height="100vh", notebook=False, cdn_resources="in_line")
+
+    # Add nodes with colors based on entity type
+    import hashlib
+
+    for node in nodes:
+        node_id = node.get("id", "")
+        props = node.get("properties", {})
+        entity_type = props.get("entity_type", "unknown")
+        description = props.get("description", "")
+
+        # Deterministic color from entity type
+        color_hash = int(hashlib.md5(entity_type.encode()).hexdigest()[:6], 16)
+        color = f"#{color_hash:06x}"
+
+        net.add_node(
+            node_id,
+            label=node_id,
+            title=f"[{entity_type}] {description[:200]}"
+            if description
+            else entity_type,
+            color=color,
+        )
+
+    # Add edges
+    for edge in edges:
+        source = edge.get("source", "")
+        target = edge.get("target", "")
+        props = edge.get("properties", {})
+        rel_type = edge.get("type", "")
+        description = props.get("description", "")
+
+        net.add_edge(
+            source,
+            target,
+            title=f"[{rel_type}] {description[:200]}" if description else rel_type,
+            label=rel_type,
+        )
+
+    net.save_graph(output_file)
+    print(f"Graph saved to {output_file}")
+    return output_file
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Visualize LightRAG knowledge graph from OpenSearch"
+    )
+    parser.add_argument(
+        "--html",
+        action="store_true",
+        help="Generate standalone HTML file instead of opening WebUI",
+    )
+    parser.add_argument(
+        "--server",
+        default="http://localhost:9621",
+        help="LightRAG Server URL (default: http://localhost:9621)",
+    )
+    parser.add_argument(
+        "--output",
+        default="knowledge_graph_opensearch.html",
+        help="Output HTML file (default: knowledge_graph_opensearch.html)",
+    )
+    parser.add_argument(
+        "--label",
+        default="*",
+        help="Starting node label, or '*' for all nodes (default: *)",
+    )
+    parser.add_argument(
+        "--max-nodes",
+        type=int,
+        default=300,
+        help="Maximum nodes to fetch (default: 300)",
+    )
+    args = parser.parse_args()
+
+    # Verify server is running
+    try:
+        requests.get(f"{args.server}/health", timeout=5)
+    except requests.ConnectionError:
+        print(f"Error: Cannot connect to LightRAG Server at {args.server}")
+        print("Start the server first: lightrag-server --host 0.0.0.0 --port 9621")
+        sys.exit(1)
+
+    if args.html:
+        # Generate standalone HTML
+        graph_data = fetch_graph(args.server, args.label, args.max_nodes)
+        output = generate_html(graph_data, args.output)
+        webbrowser.open(f"file://{os.path.abspath(output)}")
+    else:
+        # Open WebUI graph explorer
+        url = f"{args.server}/#/graph"
+        print(f"Opening LightRAG WebUI graph explorer: {url}")
+        webbrowser.open(url)
+
+
+if __name__ == "__main__":
+    main()

examples/insert_custom_kg.py

@@ -0,0 +1,115 @@
+import os
+from lightrag import LightRAG
+from lightrag.llm.openai import gpt_4o_mini_complete
+#########
+# Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
+# import nest_asyncio
+# nest_asyncio.apply()
+#########
+
+WORKING_DIR = "./custom_kg"
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+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
+)
+
+custom_kg = {
+    "entities": [
+        {
+            "entity_name": "CompanyA",
+            "entity_type": "Organization",
+            "description": "A major technology company",
+            "source_id": "Source1",
+        },
+        {
+            "entity_name": "ProductX",
+            "entity_type": "Product",
+            "description": "A popular product developed by CompanyA",
+            "source_id": "Source1",
+        },
+        {
+            "entity_name": "PersonA",
+            "entity_type": "Person",
+            "description": "A renowned researcher in AI",
+            "source_id": "Source2",
+        },
+        {
+            "entity_name": "UniversityB",
+            "entity_type": "Organization",
+            "description": "A leading university specializing in technology and sciences",
+            "source_id": "Source2",
+        },
+        {
+            "entity_name": "CityC",
+            "entity_type": "Location",
+            "description": "A large metropolitan city known for its culture and economy",
+            "source_id": "Source3",
+        },
+        {
+            "entity_name": "EventY",
+            "entity_type": "Event",
+            "description": "An annual technology conference held in CityC",
+            "source_id": "Source3",
+        },
+    ],
+    "relationships": [
+        {
+            "src_id": "CompanyA",
+            "tgt_id": "ProductX",
+            "description": "CompanyA develops ProductX",
+            "keywords": "develop, produce",
+            "weight": 1.0,
+            "source_id": "Source1",
+        },
+        {
+            "src_id": "PersonA",
+            "tgt_id": "UniversityB",
+            "description": "PersonA works at UniversityB",
+            "keywords": "employment, affiliation",
+            "weight": 0.9,
+            "source_id": "Source2",
+        },
+        {
+            "src_id": "CityC",
+            "tgt_id": "EventY",
+            "description": "EventY is hosted in CityC",
+            "keywords": "host, location",
+            "weight": 0.8,
+            "source_id": "Source3",
+        },
+    ],
+    "chunks": [
+        {
+            "content": "ProductX, developed by CompanyA, has revolutionized the market with its cutting-edge features.",
+            "source_id": "Source1",
+            "source_chunk_index": 0,
+        },
+        {
+            "content": "One outstanding feature of ProductX is its advanced AI capabilities.",
+            "source_id": "Source1",
+            "chunk_order_index": 1,
+        },
+        {
+            "content": "PersonA is a prominent researcher at UniversityB, focusing on artificial intelligence and machine learning.",
+            "source_id": "Source2",
+            "source_chunk_index": 0,
+        },
+        {
+            "content": "EventY, held in CityC, attracts technology enthusiasts and companies from around the globe.",
+            "source_id": "Source3",
+            "source_chunk_index": 0,
+        },
+        {
+            "content": "None",
+            "source_id": "UNKNOWN",
+            "source_chunk_index": 0,
+        },
+    ],
+}
+
+rag.insert_custom_kg(custom_kg)

examples/lightrag_ag2_multiagent_demo.py

@@ -0,0 +1,296 @@
+"""LightRAG + AG2 Multi-Agent Demo.
+
+Demonstrates how AG2 agents can use LightRAG's knowledge graph retrieval
+as a tool. Multiple specialized agents collaborate to answer complex
+questions over indexed documents.
+
+Architecture:
+    User -> AG2 GroupChat (Researcher + Analyst + Writer) -> LightRAG queries
+    - Researcher: uses LightRAG hybrid search to gather facts
+    - Analyst: uses LightRAG naive (vector) search for complementary results
+    - Writer: synthesizes findings into a final answer
+
+Requires:
+    pip install lightrag-hku "ag2[openai]>=0.11.4,<1.0"
+    export OPENAI_API_KEY="..."
+
+Usage:
+    python examples/lightrag_ag2_multiagent_demo.py
+"""
+
+import asyncio
+import json
+import os
+import shutil
+import threading
+
+from autogen import (
+    AssistantAgent,
+    GroupChat,
+    GroupChatManager,
+    LLMConfig,
+    UserProxyAgent,
+)
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
+
+# --- Configuration ---
+
+WORKING_DIR = "./ag2_demo_workdir"
+
+SAMPLE_TEXT = """
+Artificial intelligence has transformed multiple industries. Machine learning,
+a subset of AI, enables systems to learn from data without explicit programming.
+Deep learning, using neural networks with many layers, has achieved breakthroughs
+in computer vision, natural language processing, and speech recognition.
+
+Transformer architectures, introduced in the 2017 paper "Attention Is All You Need"
+by Vaswani et al., revolutionized NLP. Models like GPT and BERT are built on
+transformers. GPT (Generative Pre-trained Transformer) uses decoder-only architecture
+for text generation, while BERT (Bidirectional Encoder Representations) uses
+encoder-only architecture for understanding tasks.
+
+Retrieval-Augmented Generation (RAG) combines the strengths of retrieval systems
+and generative models. Instead of relying solely on parametric knowledge, RAG
+systems retrieve relevant documents from a knowledge base and use them as context
+for generation. This approach reduces hallucination and enables models to access
+up-to-date information.
+
+Knowledge graphs represent information as entities and relationships. When combined
+with RAG, knowledge graphs enable structured reasoning over document collections.
+LightRAG implements this approach with dual-level retrieval: local search focuses
+on specific entities, while global search captures broader themes and relationships.
+"""
+
+
+# --- LightRAG Setup ---
+
+
+async def setup_lightrag() -> LightRAG:
+    """Initialize LightRAG and index sample documents."""
+    if os.path.exists(WORKING_DIR):
+        shutil.rmtree(WORKING_DIR)
+    os.makedirs(WORKING_DIR, exist_ok=True)
+
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        embedding_func=openai_embed,
+        llm_model_func=gpt_4o_mini_complete,
+    )
+    await rag.initialize_storages()
+    await rag.ainsert(SAMPLE_TEXT)
+    print("LightRAG initialized and documents indexed.\n")
+    return rag
+
+
+# --- Async Bridge ---
+# AG2 runs tools in a background thread without an event loop.
+# We maintain a dedicated event loop in a separate thread for LightRAG async calls.
+
+_bg_loop: asyncio.AbstractEventLoop = None
+
+
+def _start_background_loop(loop: asyncio.AbstractEventLoop):
+    asyncio.set_event_loop(loop)
+    loop.run_forever()
+
+
+def _run_async(coro):
+    """Submit a coroutine to the background event loop and wait for the result."""
+    future = asyncio.run_coroutine_threadsafe(coro, _bg_loop)
+    return future.result(timeout=120)
+
+
+# --- AG2 Agent Tools ---
+
+# Global reference to LightRAG instance (set in main)
+_rag_instance: LightRAG = None
+
+
+def create_agents():
+    """Create AG2 agents with LightRAG tools."""
+    llm_config = LLMConfig(
+        {
+            "model": os.environ.get("OPENAI_MODEL", "gpt-4o-mini"),
+            "api_key": os.environ["OPENAI_API_KEY"],
+            "api_type": "openai",
+        }
+    )
+
+    researcher = AssistantAgent(
+        name="Researcher",
+        system_message=(
+            "You are a research specialist. Use the lightrag_query tool to search "
+            "the knowledge base. Start with 'hybrid' mode for comprehensive results. "
+            "If you need specific entity details, use 'local' mode. "
+            "Present your findings as structured bullet points. "
+            "Always call the tool -- do NOT answer from your own knowledge."
+        ),
+        llm_config=llm_config,
+    )
+
+    analyst = AssistantAgent(
+        name="Analyst",
+        system_message=(
+            "You are a knowledge graph analyst. Your FIRST action MUST be calling "
+            "the lightrag_query tool with mode='naive' to run a direct vector search. "
+            "This gives different results from the Researcher's hybrid search. "
+            "After receiving the naive search results, compare them with the "
+            "Researcher's findings and highlight any additional insights. "
+            "You MUST call the tool before writing any analysis."
+        ),
+        llm_config=llm_config,
+    )
+
+    writer = AssistantAgent(
+        name="Writer",
+        system_message=(
+            "You are a technical writer. Synthesize the findings from the "
+            "Researcher and Analyst into a clear, well-structured answer. "
+            "Do NOT use the search tool -- work only with what the other agents "
+            "have found. End your response with TERMINATE."
+        ),
+        llm_config=llm_config,
+    )
+
+    def is_termination(msg):
+        return "TERMINATE" in (msg.get("content") or "")
+
+    user_proxy = UserProxyAgent(
+        name="User",
+        human_input_mode="NEVER",
+        max_consecutive_auto_reply=10,
+        code_execution_config=False,
+        is_termination_msg=is_termination,
+    )
+
+    # --- Register LightRAG as a tool ---
+
+    @user_proxy.register_for_execution()
+    @researcher.register_for_llm(
+        description=(
+            "Query the LightRAG knowledge base. "
+            "mode: 'naive' (simple vector), 'local' (entity-focused), "
+            "'global' (theme/relationship-focused), 'hybrid' (combined). "
+            "Returns retrieved context from indexed documents."
+        )
+    )
+    @analyst.register_for_llm(
+        description=(
+            "Query the LightRAG knowledge base. "
+            "mode: 'naive' (simple vector), 'local' (entity-focused), "
+            "'global' (theme/relationship-focused), 'hybrid' (combined). "
+            "Returns retrieved context from indexed documents."
+        )
+    )
+    def lightrag_query(query: str, mode: str = "hybrid") -> str:
+        """Query LightRAG synchronously (wraps async call)."""
+        valid_modes = {"naive", "local", "global", "hybrid"}
+        if mode not in valid_modes:
+            return json.dumps(
+                {"error": f"Invalid mode '{mode}'. Use one of: {valid_modes}"}
+            )
+        try:
+            result = _run_async(
+                _rag_instance.aquery(query, param=QueryParam(mode=mode))
+            )
+            return json.dumps({"mode": mode, "query": query, "result": result})
+        except Exception as e:
+            return json.dumps({"error": str(e)})
+
+    return user_proxy, researcher, analyst, writer
+
+
+def run_multiagent_query(user_proxy, researcher, analyst, writer, question: str):
+    """Run a multi-agent GroupChat to answer a question using LightRAG."""
+    # Enforce pipeline: Researcher -> Analyst -> Writer.
+    # func_call_filter (default True) automatically routes tool calls
+    # to/from user_proxy, so transitions only govern non-tool handoffs.
+    # User can only start with Researcher; Researcher advances to Analyst;
+    # Analyst advances to Writer. Writer terminates the conversation.
+    allowed_transitions = {
+        user_proxy: [researcher],
+        researcher: [user_proxy, analyst],
+        analyst: [user_proxy, writer],
+        writer: [],
+    }
+
+    group_chat = GroupChat(
+        agents=[user_proxy, researcher, analyst, writer],
+        messages=[],
+        max_round=12,
+        allowed_or_disallowed_speaker_transitions=allowed_transitions,
+        speaker_transitions_type="allowed",
+    )
+
+    manager = GroupChatManager(
+        groupchat=group_chat,
+        llm_config=LLMConfig(
+            {
+                "model": os.environ.get("OPENAI_MODEL", "gpt-4o-mini"),
+                "api_key": os.environ["OPENAI_API_KEY"],
+                "api_type": "openai",
+            }
+        ),
+        is_termination_msg=lambda msg: "TERMINATE" in (msg.get("content") or ""),
+    )
+
+    print(f"Question: {question}\n{'=' * 60}\n")
+    user_proxy.run(manager, message=question).process()
+    print(f"\n{'=' * 60}")
+
+
+# --- Main ---
+
+
+def main():
+    global _rag_instance, _bg_loop
+
+    if not os.getenv("OPENAI_API_KEY"):
+        print(
+            "Error: OPENAI_API_KEY environment variable is not set.\n"
+            "Set it by running: export OPENAI_API_KEY='your-openai-api-key'"
+        )
+        return
+
+    # Start a background event loop for LightRAG async calls.
+    # AG2 tools run in threads without an event loop, so we need a
+    # persistent loop that can accept coroutines from any thread.
+    _bg_loop = asyncio.new_event_loop()
+    bg_thread = threading.Thread(
+        target=_start_background_loop, args=(_bg_loop,), daemon=True
+    )
+    bg_thread.start()
+
+    try:
+        # Step 1: Set up LightRAG (async, runs on the background loop)
+        _rag_instance = _run_async(setup_lightrag())
+
+        # Step 2: Create AG2 agents with LightRAG tools
+        user_proxy, researcher, analyst, writer = create_agents()
+
+        # Step 3: Ask a complex question
+        run_multiagent_query(
+            user_proxy,
+            researcher,
+            analyst,
+            writer,
+            question=(
+                "How do transformer architectures relate to RAG systems? "
+                "What role do knowledge graphs play in improving retrieval quality?"
+            ),
+        )
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    finally:
+        if _rag_instance:
+            _run_async(_rag_instance.finalize_storages())
+        _bg_loop.call_soon_threadsafe(_bg_loop.stop)
+        bg_thread.join(timeout=5)
+        shutil.rmtree(WORKING_DIR, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    main()
+    print("\nDone!")

examples/lightrag_azure_openai_demo.py

@@ -0,0 +1,125 @@
+import os
+import asyncio
+from lightrag import LightRAG, QueryParam
+from lightrag.utils import EmbeddingFunc
+import numpy as np
+from dotenv import load_dotenv
+import logging
+from openai import AzureOpenAI
+
+logging.basicConfig(level=logging.INFO)
+
+load_dotenv()
+
+AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION")
+AZURE_OPENAI_DEPLOYMENT = os.getenv("AZURE_OPENAI_DEPLOYMENT")
+AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
+AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
+
+AZURE_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_EMBEDDING_DEPLOYMENT")
+AZURE_EMBEDDING_API_VERSION = os.getenv("AZURE_EMBEDDING_API_VERSION")
+
+WORKING_DIR = "./dickens"
+
+if os.path.exists(WORKING_DIR):
+    import shutil
+
+    shutil.rmtree(WORKING_DIR)
+
+os.mkdir(WORKING_DIR)
+
+
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    client = AzureOpenAI(
+        api_key=AZURE_OPENAI_API_KEY,
+        api_version=AZURE_OPENAI_API_VERSION,
+        azure_endpoint=AZURE_OPENAI_ENDPOINT,
+    )
+
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    if history_messages:
+        messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+
+    chat_completion = client.chat.completions.create(
+        model=AZURE_OPENAI_DEPLOYMENT,  # model = "deployment_name".
+        messages=messages,
+        temperature=kwargs.get("temperature", 0),
+        top_p=kwargs.get("top_p", 1),
+        n=kwargs.get("n", 1),
+    )
+    if not chat_completion.choices or chat_completion.choices[0].message is None:
+        return ""
+    return chat_completion.choices[0].message.content
+
+
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    client = AzureOpenAI(
+        api_key=AZURE_OPENAI_API_KEY,
+        api_version=AZURE_EMBEDDING_API_VERSION,
+        azure_endpoint=AZURE_OPENAI_ENDPOINT,
+    )
+    embedding = client.embeddings.create(model=AZURE_EMBEDDING_DEPLOYMENT, input=texts)
+
+    embeddings = [item.embedding for item in embedding.data]
+    return np.array(embeddings)
+
+
+async def test_funcs():
+    result = await llm_model_func("How are you?")
+    print("Resposta do llm_model_func: ", result)
+
+    result = await embedding_func(["How are you?"])
+    print("Resultado do embedding_func: ", result.shape)
+    print("Dimensão da embedding: ", result.shape[1])
+
+
+asyncio.run(test_funcs())
+
+embedding_dimension = 3072
+
+
+async def initialize_rag():
+    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,
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+    return rag
+
+
+def main():
+    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_text = "What are the main themes?"
+
+    print("Result (Naive):")
+    print(rag.query(query_text, param=QueryParam(mode="naive")))
+
+    print("\nResult (Local):")
+    print(rag.query(query_text, param=QueryParam(mode="local")))
+
+    print("\nResult (Global):")
+    print(rag.query(query_text, param=QueryParam(mode="global")))
+
+    print("\nResult (Hybrid):")
+    print(rag.query(query_text, param=QueryParam(mode="hybrid")))
+
+
+if __name__ == "__main__":
+    main()

examples/lightrag_gemini_demo.py

@@ -0,0 +1,122 @@
+"""
+LightRAG Demo with Google Gemini Models
+
+This example demonstrates how to use LightRAG with Google's Gemini 2.0 Flash model
+for text generation and the text-embedding-004 model for embeddings.
+
+Prerequisites:
+    1. Set GEMINI_API_KEY environment variable:
+       export GEMINI_API_KEY='your-actual-api-key'
+
+    2. Prepare a text file named 'book.txt' in the current directory
+       (or modify BOOK_FILE constant to point to your text file)
+
+Usage:
+    python examples/lightrag_gemini_demo.py
+"""
+
+import os
+import asyncio
+import nest_asyncio
+import numpy as np
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.gemini import gemini_model_complete, gemini_embed
+from lightrag.utils import wrap_embedding_func_with_attrs
+
+nest_asyncio.apply()
+
+WORKING_DIR = "./rag_storage"
+BOOK_FILE = "./book.txt"
+
+# Validate API key
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+if not GEMINI_API_KEY:
+    raise ValueError(
+        "GEMINI_API_KEY environment variable is not set. "
+        "Please set it with: export GEMINI_API_KEY='your-api-key'"
+    )
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+# --------------------------------------------------
+# LLM function
+# --------------------------------------------------
+async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
+    return await gemini_model_complete(
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=GEMINI_API_KEY,
+        model_name="gemini-2.0-flash",
+        **kwargs,
+    )
+
+
+# --------------------------------------------------
+# Embedding function
+# --------------------------------------------------
+@wrap_embedding_func_with_attrs(
+    embedding_dim=768,
+    send_dimensions=True,
+    max_token_size=2048,
+    model_name="models/text-embedding-004",
+)
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    return await gemini_embed.func(
+        texts, api_key=GEMINI_API_KEY, model="models/text-embedding-004"
+    )
+
+
+# --------------------------------------------------
+# Initialize RAG
+# --------------------------------------------------
+async def initialize_rag():
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        embedding_func=embedding_func,
+        llm_model_name="gemini-2.0-flash",
+    )
+
+    # 🔑 REQUIRED
+    await rag.initialize_storages()
+    return rag
+
+
+# --------------------------------------------------
+# Main
+# --------------------------------------------------
+def main():
+    # Validate book file exists
+    if not os.path.exists(BOOK_FILE):
+        raise FileNotFoundError(
+            f"'{BOOK_FILE}' not found. "
+            "Please provide a text file to index in the current directory."
+        )
+
+    rag = asyncio.run(initialize_rag())
+
+    # Insert text
+    with open(BOOK_FILE, "r", encoding="utf-8") as f:
+        rag.insert(f.read())
+
+    query = "What are the top themes?"
+
+    print("\nNaive Search:")
+    print(rag.query(query, param=QueryParam(mode="naive")))
+
+    print("\nLocal Search:")
+    print(rag.query(query, param=QueryParam(mode="local")))
+
+    print("\nGlobal Search:")
+    print(rag.query(query, param=QueryParam(mode="global")))
+
+    print("\nHybrid Search:")
+    print(rag.query(query, param=QueryParam(mode="hybrid")))
+
+
+if __name__ == "__main__":
+    main()

examples/lightrag_gemini_postgres_demo.py

@@ -0,0 +1,178 @@
+"""
+LightRAG Demo with PostgreSQL + Google Gemini
+
+This example demonstrates how to use LightRAG with:
+- Google Gemini (LLM + Embeddings)
+- PostgreSQL-backed storages for:
+  - Vector storage
+  - Graph storage
+  - KV storage
+  - Document status storage
+
+Prerequisites:
+1. PostgreSQL database running and accessible
+2. Required tables will be auto-created by LightRAG
+3. Set environment variables (example .env):
+
+   POSTGRES_HOST=localhost
+   POSTGRES_PORT=5432
+   POSTGRES_USER=admin
+   POSTGRES_PASSWORD=admin
+   POSTGRES_DATABASE=ai
+
+   LIGHTRAG_KV_STORAGE=PGKVStorage
+   LIGHTRAG_DOC_STATUS_STORAGE=PGDocStatusStorage
+   LIGHTRAG_GRAPH_STORAGE=PGGraphStorage
+   LIGHTRAG_VECTOR_STORAGE=PGVectorStorage
+
+   GEMINI_API_KEY=your-api-key
+
+4. Prepare a text file to index (default: Data/book-small.txt)
+
+Usage:
+    python examples/lightrag_postgres_demo.py
+"""
+
+import os
+import asyncio
+import numpy as np
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.gemini import gemini_model_complete, gemini_embed
+from lightrag.utils import setup_logger, wrap_embedding_func_with_attrs
+
+
+# --------------------------------------------------
+# Logger
+# --------------------------------------------------
+setup_logger("lightrag", level="INFO")
+
+
+# --------------------------------------------------
+# Config
+# --------------------------------------------------
+WORKING_DIR = "./rag_storage"
+BOOK_FILE = "Data/book.txt"
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+if not GEMINI_API_KEY:
+    raise ValueError("GEMINI_API_KEY environment variable is not set")
+
+
+# --------------------------------------------------
+# LLM function (Gemini)
+# --------------------------------------------------
+async def llm_model_func(
+    prompt,
+    system_prompt=None,
+    history_messages=[],
+    keyword_extraction=False,
+    **kwargs,
+) -> str:
+    return await gemini_model_complete(
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=GEMINI_API_KEY,
+        model_name="gemini-2.0-flash",
+        **kwargs,
+    )
+
+
+# --------------------------------------------------
+# Embedding function (Gemini)
+# --------------------------------------------------
+@wrap_embedding_func_with_attrs(
+    embedding_dim=768,
+    max_token_size=2048,
+    model_name="models/text-embedding-004",
+)
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    return await gemini_embed.func(
+        texts,
+        api_key=GEMINI_API_KEY,
+        model="models/text-embedding-004",
+    )
+
+
+# --------------------------------------------------
+# Initialize RAG with PostgreSQL storages
+# --------------------------------------------------
+async def initialize_rag() -> LightRAG:
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_name="gemini-2.0-flash",
+        llm_model_func=llm_model_func,
+        embedding_func=embedding_func,
+        # Performance tuning
+        embedding_func_max_async=4,
+        embedding_batch_num=8,
+        llm_model_max_async=2,
+        # Chunking
+        chunk_token_size=1200,
+        chunk_overlap_token_size=100,
+        # PostgreSQL-backed storages
+        graph_storage="PGGraphStorage",
+        vector_storage="PGVectorStorage",
+        doc_status_storage="PGDocStatusStorage",
+        kv_storage="PGKVStorage",
+    )
+
+    # REQUIRED: initialize all storage backends
+    await rag.initialize_storages()
+    return rag
+
+
+# --------------------------------------------------
+# Main
+# --------------------------------------------------
+async def main():
+    rag = None
+    try:
+        print("Initializing LightRAG with PostgreSQL + Gemini...")
+        rag = await initialize_rag()
+
+        if not os.path.exists(BOOK_FILE):
+            raise FileNotFoundError(
+                f"'{BOOK_FILE}' not found. Please provide a text file to index."
+            )
+
+        print(f"\nReading document: {BOOK_FILE}")
+        with open(BOOK_FILE, "r", encoding="utf-8") as f:
+            content = f.read()
+
+        print(f"Loaded document ({len(content)} characters)")
+
+        print("\nInserting document into LightRAG (this may take some time)...")
+        await rag.ainsert(content)
+        print("Document indexed successfully!")
+
+        print("\n" + "=" * 60)
+        print("Running sample queries")
+        print("=" * 60)
+
+        query = "What are the top themes in this document?"
+
+        for mode in ["naive", "local", "global", "hybrid"]:
+            print(f"\n[{mode.upper()} MODE]")
+            result = await rag.aquery(query, param=QueryParam(mode=mode))
+            print(result[:400] + "..." if len(result) > 400 else result)
+
+        print("\nRAG system is ready for use!")
+
+    except Exception as e:
+        print("An error occurred:", e)
+        import traceback
+
+        traceback.print_exc()
+
+    finally:
+        if rag is not None:
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/lightrag_gemini_workspace_demo.py

@@ -0,0 +1,131 @@
+"""
+LightRAG Data Isolation Demo: Workspace Management
+
+This example demonstrates how to maintain multiple isolated knowledge bases
+within a single application using LightRAG's 'workspace' feature.
+
+Key Concepts:
+- Workspace Isolation: Each RAG instance is assigned a unique workspace name,
+  which ensures that Knowledge Graphs, Vector Databases, and Chunks are
+  stored in separate, non-conflicting directories.
+- Independent Configuration: Different workspaces can utilize different
+  entity type guidance and document sets simultaneously.
+
+Prerequisites:
+1. Set the following environment variables:
+   - GEMINI_API_KEY: Your Google Gemini API key.
+2. Ensure your data directory contains:
+   - Data/book-small.txt
+   - Data/HR_policies.txt
+
+Usage:
+    python lightrag_workspace_demo.py
+"""
+
+import os
+import asyncio
+import numpy as np
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.gemini import gemini_model_complete, gemini_embed
+from lightrag.utils import wrap_embedding_func_with_attrs
+
+
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    """Wrapper for Gemini LLM completion."""
+    return await gemini_model_complete(
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=os.getenv("GEMINI_API_KEY"),
+        model_name="gemini-2.0-flash-exp",
+        **kwargs,
+    )
+
+
+@wrap_embedding_func_with_attrs(
+    embedding_dim=768, max_token_size=2048, model_name="models/text-embedding-004"
+)
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    """Wrapper for Gemini embedding model."""
+    return await gemini_embed.func(
+        texts, api_key=os.getenv("GEMINI_API_KEY"), model="models/text-embedding-004"
+    )
+
+
+async def initialize_rag(
+    workspace: str = "default_workspace",
+) -> LightRAG:
+    """
+    Initializes a LightRAG instance with data isolation.
+
+    Entity type guidance can be customized by passing
+    addon_params={'entity_types_guidance': '...'} to LightRAG.
+    """
+
+    rag = LightRAG(
+        workspace=workspace,
+        llm_model_name="gemini-2.0-flash",
+        llm_model_func=llm_model_func,
+        embedding_func=embedding_func,
+        embedding_func_max_async=4,
+        embedding_batch_num=8,
+        llm_model_max_async=2,
+    )
+
+    await rag.initialize_storages()
+    return rag
+
+
+async def main():
+    rag_1 = None
+    rag_2 = None
+    try:
+        # 1. Initialize Isolated Workspaces
+        # Instance 1: Dedicated to literary analysis
+        # Instance 2: Dedicated to corporate HR documentation
+        print("Initializing isolated LightRAG workspaces...")
+        rag_1 = await initialize_rag("rag_workspace_book")
+        rag_2 = await initialize_rag("rag_workspace_hr")
+
+        # 2. Populate Workspace 1 (Literature)
+        book_path = "Data/book-small.txt"
+        if os.path.exists(book_path):
+            with open(book_path, "r", encoding="utf-8") as f:
+                print(f"Indexing {book_path} into Literature Workspace...")
+                await rag_1.ainsert(f.read())
+
+        # 3. Populate Workspace 2 (Corporate)
+        hr_path = "Data/HR_policies.txt"
+        if os.path.exists(hr_path):
+            with open(hr_path, "r", encoding="utf-8") as f:
+                print(f"Indexing {hr_path} into HR Workspace...")
+                await rag_2.ainsert(f.read())
+
+        # 4. Context-Specific Querying
+        print("\n--- Querying Literature Workspace ---")
+        res1 = await rag_1.aquery(
+            "What is the main theme?",
+            param=QueryParam(mode="hybrid", stream=False),
+        )
+        print(f"Book Analysis: {res1[:200]}...")
+
+        print("\n--- Querying HR Workspace ---")
+        res2 = await rag_2.aquery(
+            "What is the leave policy?", param=QueryParam(mode="hybrid")
+        )
+        print(f"HR Response: {res2[:200]}...")
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    finally:
+        # Finalize storage to safely close DB connections and write buffers
+        if rag_1:
+            await rag_1.finalize_storages()
+        if rag_2:
+            await rag_2.finalize_storages()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/lightrag_ollama_demo.py

@@ -0,0 +1,219 @@
+import asyncio
+import os
+import inspect
+import logging
+import logging.config
+from functools import partial
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.ollama import ollama_model_complete, ollama_embed
+from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
+
+from dotenv import load_dotenv
+
+load_dotenv(dotenv_path=".env", override=False)
+
+WORKING_DIR = "./dickens"
+
+
+def configure_logging():
+    """Configure logging for the application"""
+
+    # Reset any existing handlers to ensure clean configuration
+    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
+        logger_instance = logging.getLogger(logger_name)
+        logger_instance.handlers = []
+        logger_instance.filters = []
+
+    # Get log directory path from environment variable or use current directory
+    log_dir = os.getenv("LOG_DIR", os.getcwd())
+    log_file_path = os.path.abspath(os.path.join(log_dir, "lightrag_ollama_demo.log"))
+
+    print(f"\nLightRAG compatible demo log file: {log_file_path}\n")
+    os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
+
+    # Get log file max size and backup count from environment variables
+    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))  # Default 10MB
+    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))  # Default 5 backups
+
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "default": {
+                    "format": "%(levelname)s: %(message)s",
+                },
+                "detailed": {
+                    "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+                },
+            },
+            "handlers": {
+                "console": {
+                    "formatter": "default",
+                    "class": "logging.StreamHandler",
+                    "stream": "ext://sys.stderr",
+                },
+                "file": {
+                    "formatter": "detailed",
+                    "class": "logging.handlers.RotatingFileHandler",
+                    "filename": log_file_path,
+                    "maxBytes": log_max_bytes,
+                    "backupCount": log_backup_count,
+                    "encoding": "utf-8",
+                },
+            },
+            "loggers": {
+                "lightrag": {
+                    "handlers": ["console", "file"],
+                    "level": "INFO",
+                    "propagate": False,
+                },
+            },
+        }
+    )
+
+    # Set the logger level to INFO
+    logger.setLevel(logging.INFO)
+    # Enable verbose debug if needed
+    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
+
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+async def initialize_rag():
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=ollama_model_complete,
+        llm_model_name=os.getenv("LLM_MODEL", "qwen2.5-coder:7b"),
+        summary_max_tokens=8192,
+        llm_model_kwargs={
+            "host": os.getenv("LLM_BINDING_HOST", "http://localhost:11434"),
+            "options": {"num_ctx": 8192},
+            "timeout": int(os.getenv("TIMEOUT", "300")),
+        },
+        # Note: ollama_embed is decorated with @wrap_embedding_func_with_attrs,
+        # which wraps it in an EmbeddingFunc. Using .func accesses the original
+        # unwrapped function to avoid double wrapping when we create our own
+        # EmbeddingFunc with custom configuration (embedding_dim, max_token_size).
+        embedding_func=EmbeddingFunc(
+            embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
+            max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "8192")),
+            func=partial(
+                ollama_embed.func,  # Access the unwrapped function to avoid double EmbeddingFunc wrapping
+                embed_model=os.getenv("EMBEDDING_MODEL", "bge-m3:latest"),
+                host=os.getenv("EMBEDDING_BINDING_HOST", "http://localhost:11434"),
+            ),
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+    return rag
+
+
+async def print_stream(stream):
+    async for chunk in stream:
+        print(chunk, end="", flush=True)
+
+
+async def main():
+    try:
+        # Clear old data files
+        files_to_delete = [
+            "graph_chunk_entity_relation.graphml",
+            "kv_store_doc_status.json",
+            "kv_store_full_docs.json",
+            "kv_store_text_chunks.json",
+            "vdb_chunks.json",
+            "vdb_entities.json",
+            "vdb_relationships.json",
+        ]
+
+        for file in files_to_delete:
+            file_path = os.path.join(WORKING_DIR, file)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                print(f"Deleting old file:: {file_path}")
+
+        # Initialize RAG instance
+        rag = await initialize_rag()
+
+        # Test embedding function
+        test_text = ["This is a test string for embedding."]
+        embedding = await rag.embedding_func(test_text)
+        embedding_dim = embedding.shape[1]
+        print("\n=======================")
+        print("Test embedding function")
+        print("========================")
+        print(f"Test dict: {test_text}")
+        print(f"Detected embedding dimension: {embedding_dim}\n\n")
+
+        with open("./book.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+
+        # Perform naive search
+        print("\n=====================")
+        print("Query mode: naive")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="naive", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform local search
+        print("\n=====================")
+        print("Query mode: local")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="local", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform global search
+        print("\n=====================")
+        print("Query mode: global")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="global", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform hybrid search
+        print("\n=====================")
+        print("Query mode: hybrid")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="hybrid", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    finally:
+        if rag:
+            await rag.llm_response_cache.index_done_callback()
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    # Configure logging before running the main function
+    configure_logging()
+    asyncio.run(main())
+    print("\nDone!")

examples/lightrag_openai_compatible_demo.py

@@ -0,0 +1,229 @@
+import os
+import asyncio
+import inspect
+import logging
+import logging.config
+from functools import partial
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import openai_complete_if_cache
+from lightrag.llm.ollama import ollama_embed
+from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
+
+from dotenv import load_dotenv
+
+load_dotenv(dotenv_path=".env", override=False)
+
+WORKING_DIR = "./dickens"
+
+
+def configure_logging():
+    """Configure logging for the application"""
+
+    # Reset any existing handlers to ensure clean configuration
+    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
+        logger_instance = logging.getLogger(logger_name)
+        logger_instance.handlers = []
+        logger_instance.filters = []
+
+    # Get log directory path from environment variable or use current directory
+    log_dir = os.getenv("LOG_DIR", os.getcwd())
+    log_file_path = os.path.abspath(
+        os.path.join(log_dir, "lightrag_compatible_demo.log")
+    )
+
+    print(f"\nLightRAG compatible demo log file: {log_file_path}\n")
+    os.makedirs(os.path.dirname(log_dir), exist_ok=True)
+
+    # Get log file max size and backup count from environment variables
+    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))  # Default 10MB
+    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))  # Default 5 backups
+
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "default": {
+                    "format": "%(levelname)s: %(message)s",
+                },
+                "detailed": {
+                    "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+                },
+            },
+            "handlers": {
+                "console": {
+                    "formatter": "default",
+                    "class": "logging.StreamHandler",
+                    "stream": "ext://sys.stderr",
+                },
+                "file": {
+                    "formatter": "detailed",
+                    "class": "logging.handlers.RotatingFileHandler",
+                    "filename": log_file_path,
+                    "maxBytes": log_max_bytes,
+                    "backupCount": log_backup_count,
+                    "encoding": "utf-8",
+                },
+            },
+            "loggers": {
+                "lightrag": {
+                    "handlers": ["console", "file"],
+                    "level": "INFO",
+                    "propagate": False,
+                },
+            },
+        }
+    )
+
+    # Set the logger level to INFO
+    logger.setLevel(logging.INFO)
+    # Enable verbose debug if needed
+    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
+
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        os.getenv("LLM_MODEL", "deepseek-chat"),
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=os.getenv("LLM_BINDING_API_KEY") or os.getenv("OPENAI_API_KEY"),
+        base_url=os.getenv("LLM_BINDING_HOST", "https://api.deepseek.com"),
+        **kwargs,
+    )
+
+
+async def print_stream(stream):
+    async for chunk in stream:
+        if chunk:
+            print(chunk, end="", flush=True)
+
+
+async def initialize_rag():
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        # Note: ollama_embed is decorated with @wrap_embedding_func_with_attrs,
+        # which wraps it in an EmbeddingFunc. Using .func accesses the original
+        # unwrapped function to avoid double wrapping when we create our own
+        # EmbeddingFunc with custom configuration (embedding_dim, max_token_size).
+        embedding_func=EmbeddingFunc(
+            embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
+            max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "8192")),
+            func=partial(
+                ollama_embed.func,  # Access the unwrapped function to avoid double EmbeddingFunc wrapping
+                embed_model=os.getenv("EMBEDDING_MODEL", "bge-m3:latest"),
+                host=os.getenv("EMBEDDING_BINDING_HOST", "http://localhost:11434"),
+            ),
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+    return rag
+
+
+async def main():
+    try:
+        # Clear old data files
+        files_to_delete = [
+            "graph_chunk_entity_relation.graphml",
+            "kv_store_doc_status.json",
+            "kv_store_full_docs.json",
+            "kv_store_text_chunks.json",
+            "vdb_chunks.json",
+            "vdb_entities.json",
+            "vdb_relationships.json",
+        ]
+
+        for file in files_to_delete:
+            file_path = os.path.join(WORKING_DIR, file)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                print(f"Deleting old file:: {file_path}")
+
+        # Initialize RAG instance
+        rag = await initialize_rag()
+
+        # Test embedding function
+        test_text = ["This is a test string for embedding."]
+        embedding = await rag.embedding_func(test_text)
+        embedding_dim = embedding.shape[1]
+        print("\n=======================")
+        print("Test embedding function")
+        print("========================")
+        print(f"Test dict: {test_text}")
+        print(f"Detected embedding dimension: {embedding_dim}\n\n")
+
+        with open("./book.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+
+        # Perform naive search
+        print("\n=====================")
+        print("Query mode: naive")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="naive", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform local search
+        print("\n=====================")
+        print("Query mode: local")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="local", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform global search
+        print("\n=====================")
+        print("Query mode: global")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="global", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform hybrid search
+        print("\n=====================")
+        print("Query mode: hybrid")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="hybrid", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    finally:
+        if rag:
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    # Configure logging before running the main function
+    configure_logging()
+    asyncio.run(main())
+    print("\nDone!")

examples/lightrag_openai_demo.py

@@ -0,0 +1,187 @@
+import os
+import asyncio
+import logging
+import logging.config
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
+from lightrag.utils import logger, set_verbose_debug
+
+WORKING_DIR = "./dickens"
+
+
+def configure_logging():
+    """Configure logging for the application"""
+
+    # Reset any existing handlers to ensure clean configuration
+    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
+        logger_instance = logging.getLogger(logger_name)
+        logger_instance.handlers = []
+        logger_instance.filters = []
+
+    # Get log directory path from environment variable or use current directory
+    log_dir = os.getenv("LOG_DIR", os.getcwd())
+    log_file_path = os.path.abspath(os.path.join(log_dir, "lightrag_demo.log"))
+
+    print(f"\nLightRAG demo log file: {log_file_path}\n")
+    os.makedirs(os.path.dirname(log_dir), exist_ok=True)
+
+    # Get log file max size and backup count from environment variables
+    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))  # Default 10MB
+    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))  # Default 5 backups
+
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "default": {
+                    "format": "%(levelname)s: %(message)s",
+                },
+                "detailed": {
+                    "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+                },
+            },
+            "handlers": {
+                "console": {
+                    "formatter": "default",
+                    "class": "logging.StreamHandler",
+                    "stream": "ext://sys.stderr",
+                },
+                "file": {
+                    "formatter": "detailed",
+                    "class": "logging.handlers.RotatingFileHandler",
+                    "filename": log_file_path,
+                    "maxBytes": log_max_bytes,
+                    "backupCount": log_backup_count,
+                    "encoding": "utf-8",
+                },
+            },
+            "loggers": {
+                "lightrag": {
+                    "handlers": ["console", "file"],
+                    "level": "INFO",
+                    "propagate": False,
+                },
+            },
+        }
+    )
+
+    # Set the logger level to INFO
+    logger.setLevel(logging.INFO)
+    # Enable verbose debug if needed
+    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
+
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+async def initialize_rag():
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        embedding_func=openai_embed,
+        llm_model_func=gpt_4o_mini_complete,
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+
+    return rag
+
+
+async def main():
+    # Check if OPENAI_API_KEY environment variable exists
+    if not os.getenv("OPENAI_API_KEY"):
+        print(
+            "Error: OPENAI_API_KEY environment variable is not set. Please set this variable before running the program."
+        )
+        print("You can set the environment variable by running:")
+        print("  export OPENAI_API_KEY='your-openai-api-key'")
+        return  # Exit the async function
+
+    try:
+        # Clear old data files
+        files_to_delete = [
+            "graph_chunk_entity_relation.graphml",
+            "kv_store_doc_status.json",
+            "kv_store_full_docs.json",
+            "kv_store_text_chunks.json",
+            "vdb_chunks.json",
+            "vdb_entities.json",
+            "vdb_relationships.json",
+        ]
+
+        for file in files_to_delete:
+            file_path = os.path.join(WORKING_DIR, file)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                print(f"Deleting old file:: {file_path}")
+
+        # Initialize RAG instance
+        rag = await initialize_rag()
+
+        # Test embedding function
+        test_text = ["This is a test string for embedding."]
+        embedding = await rag.embedding_func(test_text)
+        embedding_dim = embedding.shape[1]
+        print("\n=======================")
+        print("Test embedding function")
+        print("========================")
+        print(f"Test dict: {test_text}")
+        print(f"Detected embedding dimension: {embedding_dim}\n\n")
+
+        with open("./book.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+
+        # Perform naive search
+        print("\n=====================")
+        print("Query mode: naive")
+        print("=====================")
+        print(
+            await rag.aquery(
+                "What are the top themes in this story?", param=QueryParam(mode="naive")
+            )
+        )
+
+        # Perform local search
+        print("\n=====================")
+        print("Query mode: local")
+        print("=====================")
+        print(
+            await rag.aquery(
+                "What are the top themes in this story?", param=QueryParam(mode="local")
+            )
+        )
+
+        # Perform global search
+        print("\n=====================")
+        print("Query mode: global")
+        print("=====================")
+        print(
+            await rag.aquery(
+                "What are the top themes in this story?",
+                param=QueryParam(mode="global"),
+            )
+        )
+
+        # Perform hybrid search
+        print("\n=====================")
+        print("Query mode: hybrid")
+        print("=====================")
+        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}")
+    finally:
+        if rag:
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    # Configure logging before running the main function
+    configure_logging()
+    asyncio.run(main())
+    print("\nDone!")

examples/lightrag_openai_mongodb_graph_demo.py

@@ -0,0 +1,108 @@
+import os
+import asyncio
+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
+
+#########
+# Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert()
+# import nest_asyncio
+# nest_asyncio.apply()
+#########
+WORKING_DIR = "./mongodb_test_dir"
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+os.environ["OPENAI_API_KEY"] = "sk-"
+os.environ["MONGO_URI"] = "mongodb://0.0.0.0:27017/?directConnection=true"
+os.environ["MONGO_DATABASE"] = "LightRAG"
+os.environ["MONGO_KG_COLLECTION"] = "MDB_KG"
+
+# Embedding Configuration and Functions
+EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-large")
+EMBEDDING_MAX_TOKEN_SIZE = int(os.environ.get("EMBEDDING_MAX_TOKEN_SIZE", 8192))
+
+
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    # Note: openai_embed is decorated with @wrap_embedding_func_with_attrs,
+    # which wraps it in an EmbeddingFunc. Using .func accesses the original
+    # unwrapped function to avoid double wrapping when we create our own
+    # EmbeddingFunc with custom configuration in create_embedding_function_instance().
+    return await openai_embed.func(
+        texts,
+        model=EMBEDDING_MODEL,
+    )
+
+
+async def get_embedding_dimension():
+    test_text = ["This is a test sentence."]
+    embedding = await embedding_func(test_text)
+    return embedding.shape[1]
+
+
+async def create_embedding_function_instance():
+    # Get embedding dimension
+    embedding_dimension = await get_embedding_dimension()
+    # Create embedding function instance
+    return EmbeddingFunc(
+        embedding_dim=embedding_dimension,
+        max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
+        func=embedding_func,
+    )
+
+
+async def initialize_rag():
+    embedding_func_instance = await create_embedding_function_instance()
+
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=gpt_4o_mini_complete,
+        embedding_func=embedding_func_instance,
+        graph_storage="MongoGraphStorage",
+        log_level="DEBUG",
+    )
+
+    await rag.initialize_storages()  # Auto-initializes 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()

examples/lightrag_openai_opensearch_graph_demo.py

@@ -0,0 +1,178 @@
+"""
+LightRAG Demo with OpenSearch + OpenAI
+
+This example demonstrates how to use LightRAG with:
+- OpenAI (LLM + Embeddings)
+- OpenSearch-backed storages for:
+  - KV storage
+  - Vector storage (k-NN)
+  - Graph storage (dual-index nodes + edges)
+  - Document status storage
+
+Prerequisites:
+1. OpenSearch cluster running and accessible (3.x or higher with k-NN plugin)
+2. Required indices will be auto-created by LightRAG
+3. Set environment variables (example .env):
+
+   OPENSEARCH_HOSTS=localhost:9200
+   OPENSEARCH_USER=admin
+   OPENSEARCH_PASSWORD=your-password
+   OPENSEARCH_USE_SSL=false
+   OPENSEARCH_VERIFY_CERTS=false
+
+   OPENAI_API_KEY=your-api-key
+
+4. Prepare a text file to index (default: ./book.txt)
+
+Usage:
+    python examples/lightrag_openai_opensearch_graph_demo.py
+"""
+
+import os
+import asyncio
+import numpy as np
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
+from lightrag.utils import setup_logger, EmbeddingFunc
+
+
+# --------------------------------------------------
+# Logger
+# --------------------------------------------------
+setup_logger("lightrag", level="INFO")
+
+
+# --------------------------------------------------
+# Config
+# --------------------------------------------------
+WORKING_DIR = "./opensearch_rag_storage"
+BOOK_FILE = "./book.txt"
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+# Replace with your API key, or set via environment variable
+if not os.getenv("OPENAI_API_KEY"):
+    os.environ["OPENAI_API_KEY"] = "sk-"
+
+EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-large")
+EMBEDDING_MAX_TOKEN_SIZE = int(os.environ.get("EMBEDDING_MAX_TOKEN_SIZE", 8192))
+
+
+# --------------------------------------------------
+# Embedding function (OpenAI)
+# --------------------------------------------------
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    return await openai_embed.func(
+        texts,
+        model=EMBEDDING_MODEL,
+    )
+
+
+async def get_embedding_dimension():
+    test_text = ["This is a test sentence."]
+    embedding = await embedding_func(test_text)
+    return embedding.shape[1]
+
+
+async def create_embedding_function_instance():
+    embedding_dimension = await get_embedding_dimension()
+    return EmbeddingFunc(
+        embedding_dim=embedding_dimension,
+        max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
+        func=embedding_func,
+    )
+
+
+# --------------------------------------------------
+# Initialize RAG with OpenSearch storages
+# --------------------------------------------------
+async def initialize_rag() -> LightRAG:
+    embedding_func_instance = await create_embedding_function_instance()
+
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=gpt_4o_mini_complete,
+        embedding_func=embedding_func_instance,
+        # OpenSearch-backed storages
+        kv_storage="OpenSearchKVStorage",
+        doc_status_storage="OpenSearchDocStatusStorage",
+        graph_storage="OpenSearchGraphStorage",
+        vector_storage="OpenSearchVectorDBStorage",
+    )
+
+    # REQUIRED: initialize all storage backends
+    await rag.initialize_storages()
+
+    # Clean previous data so the example is re-runnable
+    # (LLM response cache is preserved for faster reruns)
+    for storage in [
+        rag.full_docs,
+        rag.text_chunks,
+        rag.full_entities,
+        rag.full_relations,
+        rag.entity_chunks,
+        rag.relation_chunks,
+        rag.entities_vdb,
+        rag.relationships_vdb,
+        rag.chunks_vdb,
+        rag.chunk_entity_relation_graph,
+        rag.doc_status,
+    ]:
+        await storage.drop()
+    print("Cleared previous data.")
+
+    return rag
+
+
+# --------------------------------------------------
+# Main
+# --------------------------------------------------
+async def main():
+    rag = None
+    try:
+        print("Initializing LightRAG with OpenSearch + OpenAI...")
+        rag = await initialize_rag()
+
+        if not os.path.exists(BOOK_FILE):
+            raise FileNotFoundError(
+                f"'{BOOK_FILE}' not found. Please provide a text file to index."
+            )
+
+        print(f"\nReading document: {BOOK_FILE}")
+        with open(BOOK_FILE, "r", encoding="utf-8") as f:
+            content = f.read()
+
+        print(f"Loaded document ({len(content)} characters)")
+
+        print("\nInserting document into LightRAG (this may take some time)...")
+        await rag.ainsert(content)
+        print("Document indexed successfully!")
+
+        print("\n" + "=" * 60)
+        print("Running sample queries")
+        print("=" * 60)
+
+        query = "What are the top themes in this document?"
+
+        for mode in ["naive", "local", "global", "hybrid"]:
+            print(f"\n[{mode.upper()} MODE]")
+            result = await rag.aquery(query, param=QueryParam(mode=mode))
+            print(result)
+
+        print("\nRAG system is ready for use!")
+
+    except Exception as e:
+        print("An error occurred:", e)
+        import traceback
+
+        traceback.print_exc()
+
+    finally:
+        if rag is not None:
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/lightrag_vllm_demo.py

@@ -0,0 +1,180 @@
+"""
+LightRAG Demo with vLLM (LLM, Embeddings, and Reranker)
+
+This example demonstrates how to use LightRAG with:
+- vLLM-served LLM (OpenAI-compatible API)
+- vLLM-served embedding model
+- Jina-compatible reranker (also vLLM-served)
+
+Prerequisites:
+    1. Create a .env file or export environment variables:
+       - LLM_MODEL
+       - LLM_BINDING_HOST
+       - LLM_BINDING_API_KEY
+       - EMBEDDING_MODEL
+       - EMBEDDING_BINDING_HOST
+       - EMBEDDING_BINDING_API_KEY
+       - EMBEDDING_DIM
+       - EMBEDDING_TOKEN_LIMIT
+       - RERANK_MODEL
+       - RERANK_BINDING_HOST
+       - RERANK_BINDING_API_KEY
+
+    2. Prepare a text file to index (default: Data/book-small.txt)
+
+    3. Configure storage backends via environment variables or modify
+       the storage parameters in initialize_rag() below.
+
+Usage:
+    python examples/lightrag_vllm_demo.py
+"""
+
+import os
+import asyncio
+from functools import partial
+from dotenv import load_dotenv
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import openai_complete_if_cache, openai_embed
+from lightrag.utils import EmbeddingFunc
+from lightrag.rerank import jina_rerank
+
+load_dotenv()
+
+# --------------------------------------------------
+# Constants
+# --------------------------------------------------
+
+WORKING_DIR = "./LightRAG_Data"
+BOOK_FILE = "Data/book-small.txt"
+
+# --------------------------------------------------
+# LLM function (vLLM, OpenAI-compatible)
+# --------------------------------------------------
+
+
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        model=os.getenv("LLM_MODEL", "Qwen/Qwen3-14B-AWQ"),
+        prompt=prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        base_url=os.getenv("LLM_BINDING_HOST", "http://0.0.0.0:4646/v1"),
+        api_key=os.getenv("LLM_BINDING_API_KEY", "not_needed"),
+        timeout=600,
+        **kwargs,
+    )
+
+
+# --------------------------------------------------
+# Embedding function (vLLM)
+# --------------------------------------------------
+
+vLLM_emb_func = EmbeddingFunc(
+    model_name=os.getenv("EMBEDDING_MODEL", "Qwen/Qwen3-Embedding-0.6B"),
+    send_dimensions=False,
+    embedding_dim=int(os.getenv("EMBEDDING_DIM", 1024)),
+    max_token_size=int(os.getenv("EMBEDDING_TOKEN_LIMIT", 4096)),
+    func=partial(
+        openai_embed.func,
+        model=os.getenv("EMBEDDING_MODEL", "Qwen/Qwen3-Embedding-0.6B"),
+        base_url=os.getenv(
+            "EMBEDDING_BINDING_HOST",
+            "http://0.0.0.0:1234/v1",
+        ),
+        api_key=os.getenv("EMBEDDING_BINDING_API_KEY", "not_needed"),
+    ),
+)
+
+# --------------------------------------------------
+# Reranker (Jina-compatible, vLLM-served)
+# --------------------------------------------------
+
+jina_rerank_model_func = partial(
+    jina_rerank,
+    model=os.getenv("RERANK_MODEL", "Qwen/Qwen3-Reranker-0.6B"),
+    api_key=os.getenv("RERANK_BINDING_API_KEY"),
+    base_url=os.getenv(
+        "RERANK_BINDING_HOST",
+        "http://0.0.0.0:3535/v1/rerank",
+    ),
+)
+
+# --------------------------------------------------
+# Initialize RAG
+# --------------------------------------------------
+
+
+async def initialize_rag():
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        embedding_func=vLLM_emb_func,
+        rerank_model_func=jina_rerank_model_func,
+        # Storage backends (configurable via environment or modify here)
+        kv_storage=os.getenv("KV_STORAGE", "PGKVStorage"),
+        doc_status_storage=os.getenv("DOC_STATUS_STORAGE", "PGDocStatusStorage"),
+        vector_storage=os.getenv("VECTOR_STORAGE", "PGVectorStorage"),
+        graph_storage=os.getenv("GRAPH_STORAGE", "Neo4JStorage"),
+    )
+
+    await rag.initialize_storages()
+    return rag
+
+
+# --------------------------------------------------
+# Main
+# --------------------------------------------------
+
+
+async def main():
+    rag = None
+    try:
+        # Validate book file exists
+        if not os.path.exists(BOOK_FILE):
+            raise FileNotFoundError(
+                f"'{BOOK_FILE}' not found. Please provide a text file to index."
+            )
+
+        rag = await initialize_rag()
+
+        # --------------------------------------------------
+        # Data Ingestion
+        # --------------------------------------------------
+        print(f"Indexing {BOOK_FILE}...")
+        with open(BOOK_FILE, "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+        print("Indexing complete.")
+
+        # --------------------------------------------------
+        # Query
+        # --------------------------------------------------
+        query = (
+            "What are the main themes of the book, and how do the key characters "
+            "evolve throughout the story?"
+        )
+
+        print("\nHybrid Search with Reranking:")
+        result = await rag.aquery(
+            query,
+            param=QueryParam(
+                mode="hybrid",
+                stream=False,
+                enable_rerank=True,
+            ),
+        )
+
+        print("\nResult:\n", result)
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    finally:
+        if rag:
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+    print("\nDone!")

examples/milvus_kwargs_configuration_demo.py

@@ -0,0 +1,113 @@
+"""
+Example: Configuring Milvus Index Parameters via vector_db_storage_cls_kwargs
+
+This example demonstrates how to configure Milvus indexing parameters through
+vector_db_storage_cls_kwargs, which is the recommended approach when using
+frameworks that build on top of LightRAG (like RAGAnything).
+
+This approach allows configuration to be passed through framework layers without
+requiring environment variable changes or direct code modifications.
+"""
+
+import os
+import asyncio
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import openai_complete_if_cache, openai_embed
+
+
+async def main():
+    # Configure Milvus connection
+    os.environ["MILVUS_URI"] = "http://localhost:19530"
+    # os.environ["MILVUS_USER"] = "root"
+    # os.environ["MILVUS_PASSWORD"] = "your_password"
+    # os.environ["MILVUS_DB_NAME"] = "lightrag"
+
+    # Initialize LightRAG with Milvus index configuration via vector_db_storage_cls_kwargs
+    # This is the recommended approach for framework integration (e.g., RAGAnything)
+    rag = LightRAG(
+        working_dir="./demo_index",
+        llm_model_func=openai_complete_if_cache,
+        embedding_func=openai_embed,
+        # Specify Milvus as the vector storage backend
+        vector_storage="MilvusVectorDBStorage",
+        # Configure Milvus indexing parameters via vector_db_storage_cls_kwargs
+        # These parameters are extracted and passed to MilvusIndexConfig
+        vector_db_storage_cls_kwargs={
+            # Required parameter for all vector storage backends
+            "cosine_better_than_threshold": 0.2,
+            # Milvus index configuration parameters
+            # All of these can be configured via vector_db_storage_cls_kwargs
+            # Index type (AUTOINDEX, HNSW, HNSW_SQ, IVF_FLAT, etc.)
+            "index_type": "HNSW",
+            # Distance metric (COSINE, L2, IP)
+            "metric_type": "COSINE",
+            # HNSW parameters
+            "hnsw_m": 32,  # Number of connections per layer (2-2048)
+            "hnsw_ef_construction": 256,  # Size of dynamic candidate list during construction
+            "hnsw_ef": 150,  # Size of dynamic candidate list during search
+            # IVF parameters (used when index_type is IVF_FLAT, IVF_SQ8, IVF_PQ)
+            # "ivf_nlist": 2048,              # Number of cluster units
+            # "ivf_nprobe": 32,               # Number of units to query
+            # HNSW_SQ parameters (requires Milvus 2.6.8+)
+            # "sq_type": "SQ8",               # Quantization type (SQ4U, SQ6, SQ8, BF16, FP16)
+            # "sq_refine": True,              # Enable refinement
+            # "sq_refine_type": "FP32",       # Refinement type
+            # "sq_refine_k": 20,              # Number of candidates to refine
+        },
+    )
+
+    # Initialize storage backends
+    await rag.initialize_storages()
+
+    print(
+        "✅ LightRAG initialized with Milvus index configuration via vector_db_storage_cls_kwargs"
+    )
+    print(
+        f"   Index Type: {rag.vector_db_storages['entities'].index_config.index_type}"
+    )
+    print(
+        f"   Metric Type: {rag.vector_db_storages['entities'].index_config.metric_type}"
+    )
+    print(f"   HNSW M: {rag.vector_db_storages['entities'].index_config.hnsw_m}")
+    print(
+        f"   HNSW EF Construction: {rag.vector_db_storages['entities'].index_config.hnsw_ef_construction}"
+    )
+    print(f"   HNSW EF: {rag.vector_db_storages['entities'].index_config.hnsw_ef}")
+
+    # Example: Insert some text
+    sample_text = """
+    LightRAG is a Retrieval-Augmented Generation framework that uses graph-based
+    knowledge representation for enhanced information retrieval. It supports multiple
+    vector storage backends including Milvus, which offers advanced indexing options
+    for optimal performance.
+    """
+
+    await rag.ainsert(sample_text)
+    print("\n✅ Sample text inserted")
+
+    # Example: Query with different modes
+    result = await rag.aquery("What is LightRAG?", param=QueryParam(mode="hybrid"))
+    print(f"\n✅ Query result: {result[:200]}...")
+
+    # Cleanup
+    await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    print("=" * 80)
+    print("Milvus Configuration via vector_db_storage_cls_kwargs Example")
+    print("=" * 80)
+    print()
+    print("This example shows how to configure Milvus indexing parameters through")
+    print("vector_db_storage_cls_kwargs, which is ideal for framework integration.")
+    print()
+    print("Key Benefits:")
+    print("  • No environment variable changes required")
+    print("  • Configuration can be passed through framework layers")
+    print("  • Perfect for RAGAnything and similar frameworks")
+    print("  • All 11 index parameters are supported")
+    print()
+    print("=" * 80)
+    print()
+
+    asyncio.run(main())

examples/opensearch_storage_demo.py

@@ -0,0 +1,355 @@
+"""
+Integration test for OpenSearch Storage in LightRAG.
+
+Tests all 4 storage types against a live OpenSearch cluster:
+- KV Storage: CRUD, filter_keys
+- DocStatus Storage: CRUD, pagination (PIT + search_after), status counts
+- Graph Storage: nodes, edges, BFS traversal, search_labels
+- Vector Storage: k-NN upsert, query, get/delete
+
+Prerequisites:
+    OpenSearch cluster running with k-NN plugin enabled.
+    Set env vars: OPENSEARCH_HOSTS, OPENSEARCH_USER, OPENSEARCH_PASSWORD,
+                  OPENSEARCH_USE_SSL, OPENSEARCH_VERIFY_CERTS
+
+Usage:
+    OPENSEARCH_HOSTS=localhost:9200 OPENSEARCH_USER=admin \
+    OPENSEARCH_PASSWORD=<password> OPENSEARCH_USE_SSL=true \
+    OPENSEARCH_VERIFY_CERTS=false python examples/opensearch_storage_demo.py
+"""
+
+import asyncio
+import numpy as np
+from lightrag.kg.opensearch_impl import (
+    OpenSearchKVStorage,
+    OpenSearchDocStatusStorage,
+    OpenSearchGraphStorage,
+    OpenSearchVectorDBStorage,
+    ClientManager,
+)
+from lightrag.kg.shared_storage import initialize_share_data
+from lightrag.base import DocStatus
+
+
+class MockEmbeddingFunc:
+    """Mock embedding function for testing."""
+
+    def __init__(self, dim=128):
+        self.embedding_dim = dim
+        self.max_token_size = 512
+        self.model_name = "mock-embedding"
+
+    async def __call__(self, texts, **kwargs):
+        return np.random.rand(len(texts), self.embedding_dim).astype(np.float32)
+
+
+CONFIG = {
+    "embedding_batch_num": 10,
+    "max_graph_nodes": 1000,
+    "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+}
+EMBED = MockEmbeddingFunc()
+PASSED = 0
+FAILED = 0
+
+
+def check(condition, msg):
+    global PASSED, FAILED
+    if condition:
+        print(f"  ✓ {msg}")
+        PASSED += 1
+    else:
+        print(f"  ✗ {msg}")
+        FAILED += 1
+
+
+async def test_connection_manager():
+    print("\n=== Connection Manager ===")
+    client1 = await ClientManager.get_client()
+    client2 = await ClientManager.get_client()
+    check(client1 is client2, "Singleton pattern (same instance)")
+    await ClientManager.release_client(client1)
+    await ClientManager.release_client(client2)
+    check(True, "Released clients")
+
+
+async def test_kv_storage():
+    print("\n=== KV Storage ===")
+    s = OpenSearchKVStorage(
+        namespace="integ_kv",
+        global_config=CONFIG,
+        embedding_func=EMBED,
+        workspace="integ",
+    )
+    await s.initialize()
+    try:
+        await s.upsert({"k1": {"content": "hello"}, "k2": {"content": "world"}})
+        await s.index_done_callback()
+
+        doc = await s.get_by_id("k1")
+        check(doc is not None and doc.get("content") == "hello", "get_by_id")
+
+        docs = await s.get_by_ids(["k1", "k2", "missing"])
+        check(docs[0] is not None and docs[2] is None, "get_by_ids preserves order")
+
+        missing = await s.filter_keys({"k1", "k99"})
+        check(missing == {"k99"}, f"filter_keys: {missing}")
+
+        check(not await s.is_empty(), "is_empty=False")
+
+        await s.delete(["k2"])
+        await s.index_done_callback()
+        check(await s.get_by_id("k2") is None, "delete + verify")
+    finally:
+        await s.drop()
+        await s.finalize()
+
+
+async def test_doc_status_storage():
+    print("\n=== DocStatus Storage ===")
+    s = OpenSearchDocStatusStorage(
+        namespace="integ_ds",
+        global_config=CONFIG,
+        embedding_func=EMBED,
+        workspace="integ",
+    )
+    await s.initialize()
+    try:
+        # Insert docs
+        await s.upsert(
+            {
+                f"d{i}": {
+                    "status": "processed" if i % 2 == 0 else "pending",
+                    "file_path": f"/file{i}.txt",
+                    "content_summary": f"summary {i}",
+                    "content_length": i * 10,
+                    "chunks_count": i,
+                    "created_at": 1000 + i,
+                    "updated_at": 2000 + i,
+                }
+                for i in range(20)
+            }
+        )
+        await s.index_done_callback()
+
+        # Status counts
+        counts = await s.get_all_status_counts()
+        check(counts.get("all") == 20, f"all_status_counts: {counts}")
+        check(
+            counts.get("processed") == 10, f"processed count: {counts.get('processed')}"
+        )
+
+        # get_docs_by_status (uses PIT + search_after)
+        processed = await s.get_docs_by_status(DocStatus.PROCESSED)
+        check(len(processed) == 10, f"get_docs_by_status(processed): {len(processed)}")
+
+        # get_docs_by_track_id (uses PIT + search_after)
+        await s.upsert(
+            {
+                "tracked1": {
+                    "status": "processed",
+                    "file_path": "/t.txt",
+                    "content_summary": "s",
+                    "content_length": 1,
+                    "chunks_count": 1,
+                    "created_at": 100,
+                    "updated_at": 200,
+                    "track_id": "batch-42",
+                }
+            }
+        )
+        await s.index_done_callback()
+        tracked = await s.get_docs_by_track_id("batch-42")
+        check(len(tracked) == 1, f"get_docs_by_track_id: {len(tracked)}")
+
+        # Paginated (uses PIT + search_after)
+        page1, total = await s.get_docs_paginated(page=1, page_size=10)
+        check(total == 21, f"paginated total: {total}")
+        check(len(page1) == 10, f"page1 size: {len(page1)}")
+
+        page2, _ = await s.get_docs_paginated(page=2, page_size=10)
+        check(len(page2) == 10, f"page2 size: {len(page2)}")
+
+        page3, _ = await s.get_docs_paginated(page=3, page_size=10)
+        check(len(page3) == 1, f"page3 size: {len(page3)}")
+
+        # With status filter
+        filtered, ftotal = await s.get_docs_paginated(
+            status_filter=DocStatus.PENDING, page=1, page_size=50
+        )
+        check(ftotal == 10, f"filtered total: {ftotal}")
+
+        # get_doc_by_file_path
+        doc = await s.get_doc_by_file_path("/file0.txt")
+        check(doc is not None and doc["_id"] == "d0", "get_doc_by_file_path")
+    finally:
+        await s.drop()
+        await s.finalize()
+
+
+async def test_graph_storage():
+    print("\n=== Graph Storage ===")
+    s = OpenSearchGraphStorage(
+        namespace="integ_graph",
+        global_config=CONFIG,
+        embedding_func=EMBED,
+        workspace="integ",
+    )
+    await s.initialize()
+    try:
+        # Upsert nodes and edges
+        await s.upsert_node(
+            "Alice", {"entity_type": "person", "description": "A researcher"}
+        )
+        await s.upsert_node(
+            "Bob", {"entity_type": "person", "description": "A developer"}
+        )
+        await s.upsert_node(
+            "Quantum", {"entity_type": "topic", "description": "Quantum computing"}
+        )
+        await s.upsert_edge(
+            "Alice",
+            "Bob",
+            {"relationship": "knows", "weight": "1.0", "keywords": "collab"},
+        )
+        await s.upsert_edge(
+            "Alice",
+            "Quantum",
+            {"relationship": "researches", "weight": "2.0", "keywords": "research"},
+        )
+        await s.upsert_edge(
+            "Bob",
+            "Quantum",
+            {"relationship": "studies", "weight": "0.5", "keywords": "learning"},
+        )
+        await s.index_done_callback()
+
+        check(await s.has_node("Alice"), "has_node(Alice)")
+        check(not await s.has_node("Nobody"), "has_node(Nobody)=False")
+        check(await s.has_edge("Alice", "Bob"), "has_edge(Alice,Bob)")
+
+        node = await s.get_node("Alice")
+        check(node is not None and node.get("entity_type") == "person", "get_node")
+        check(node.get("entity_id") == "Alice", "entity_id field present")
+
+        check(
+            await s.node_degree("Alice") == 2,
+            f"node_degree(Alice)={await s.node_degree('Alice')}",
+        )
+
+        edges = await s.get_node_edges("Alice")
+        check(len(edges) == 2, f"get_node_edges: {len(edges)}")
+
+        # Batch ops
+        batch = await s.get_nodes_batch(["Alice", "Bob", "Missing"])
+        check("Alice" in batch and "Missing" not in batch, "get_nodes_batch")
+
+        degrees = await s.node_degrees_batch(["Alice", "Bob", "Quantum"])
+        check(degrees.get("Alice") == 2, f"node_degrees_batch: {degrees}")
+
+        # Knowledge graph (BFS)
+        kg = await s.get_knowledge_graph("Alice", max_depth=2)
+        check(len(kg.nodes) == 3, f"BFS nodes: {len(kg.nodes)}")
+        check(len(kg.edges) == 3, f"BFS edges: {len(kg.edges)}")
+
+        # get_all_labels (uses PIT)
+        labels = await s.get_all_labels()
+        check("Alice" in labels and "Bob" in labels, f"get_all_labels: {labels}")
+
+        # get_all_nodes (uses PIT)
+        all_nodes = await s.get_all_nodes()
+        check(len(all_nodes) == 3, f"get_all_nodes: {len(all_nodes)}")
+
+        # get_all_edges (uses PIT)
+        all_edges = await s.get_all_edges()
+        check(len(all_edges) == 3, f"get_all_edges: {len(all_edges)}")
+
+        # search_labels
+        found = await s.search_labels("ali", limit=10)
+        check("Alice" in found, f"search_labels('ali'): {found}")
+
+        # popular_labels
+        popular = await s.get_popular_labels(limit=10)
+        check(len(popular) > 0, f"get_popular_labels: {popular}")
+
+        # Delete node (cascading)
+        await s.delete_node("Bob")
+        await s.index_done_callback()
+        check(not await s.has_node("Bob"), "delete_node cascade")
+        check(not await s.has_edge("Alice", "Bob"), "edges removed after delete_node")
+
+        print(f"  (PPL graphlookup: {s._ppl_graphlookup_available})")
+    finally:
+        await s.drop()
+        await s.finalize()
+
+
+async def test_vector_storage():
+    print("\n=== Vector Storage ===")
+    s = OpenSearchVectorDBStorage(
+        namespace="integ_vec",
+        global_config=CONFIG,
+        embedding_func=EMBED,
+        workspace="integ",
+        meta_fields={"content", "entity_name"},
+    )
+    await s.initialize()
+    try:
+        await s.upsert(
+            {
+                "v1": {"content": "apple fruit"},
+                "v2": {"content": "banana fruit"},
+                "v3": {"content": "quantum physics"},
+            }
+        )
+        await s.index_done_callback()
+
+        results = await s.query("apple", top_k=3)
+        check(len(results) > 0, f"query returned {len(results)} results")
+        check(all("distance" in r for r in results), "results have distance")
+
+        doc = await s.get_by_id("v1")
+        check(doc is not None and doc["id"] == "v1", "get_by_id")
+
+        docs = await s.get_by_ids(["v1", "v2", "missing"])
+        check(docs[0] is not None and docs[2] is None, "get_by_ids")
+
+        vecs = await s.get_vectors_by_ids(["v1"])
+        check("v1" in vecs and len(vecs["v1"]) == 128, "get_vectors_by_ids")
+
+        await s.delete(["v3"])
+        await s.index_done_callback()
+        check(await s.get_by_id("v3") is None, "delete + verify")
+    finally:
+        await s.drop()
+        await s.finalize()
+
+
+async def main():
+    print("=" * 60)
+    print("OpenSearch Storage Integration Tests")
+    print("=" * 60)
+
+    initialize_share_data(workers=1)
+
+    try:
+        await test_connection_manager()
+        await test_kv_storage()
+        await test_doc_status_storage()
+        await test_graph_storage()
+        await test_vector_storage()
+    except Exception as e:
+        print(f"\n✗ Fatal error: {e}")
+        import traceback
+
+        traceback.print_exc()
+
+    print(f"\n{'=' * 60}")
+    print(f"Results: {PASSED} passed, {FAILED} failed")
+    print(f"{'=' * 60}")
+    if FAILED > 0:
+        exit(1)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/rerank_example.py

@@ -0,0 +1,234 @@
+"""
+LightRAG Rerank Integration Example
+
+This example demonstrates how to use rerank functionality with LightRAG
+to improve retrieval quality across different query modes.
+
+Configuration Required:
+1. Set your OpenAI LLM API key and base URL with env vars
+    LLM_MODEL
+    LLM_BINDING_HOST
+    LLM_BINDING_API_KEY
+2. Set your OpenAI embedding API key and base URL with env vars:
+    EMBEDDING_MODEL
+    EMBEDDING_DIM
+    EMBEDDING_BINDING_HOST
+    EMBEDDING_BINDING_API_KEY
+3. Set your vLLM deployed AI rerank model setting with env vars:
+    RERANK_BINDING=cohere
+    RERANK_MODEL (e.g., answerai-colbert-small-v1 or rerank-v3.5)
+    RERANK_BINDING_HOST (e.g., https://api.cohere.com/v2/rerank or LiteLLM proxy)
+    RERANK_BINDING_API_KEY
+    RERANK_ENABLE_CHUNKING=true (optional, for models with token limits)
+    RERANK_MAX_TOKENS_PER_DOC=480 (optional, default 4096)
+
+Note: Rerank is controlled per query via the 'enable_rerank' parameter (default: True)
+"""
+
+import asyncio
+import os
+import numpy as np
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import openai_complete_if_cache, openai_embed
+from lightrag.utils import EmbeddingFunc, setup_logger
+
+from functools import partial
+from lightrag.rerank import cohere_rerank
+
+# Set up your working directory
+WORKING_DIR = "./test_rerank"
+setup_logger("test_rerank")
+
+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(
+        os.getenv("LLM_MODEL"),
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=os.getenv("LLM_BINDING_API_KEY"),
+        base_url=os.getenv("LLM_BINDING_HOST"),
+        **kwargs,
+    )
+
+
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    return await openai_embed(
+        texts,
+        model=os.getenv("EMBEDDING_MODEL"),
+        api_key=os.getenv("EMBEDDING_BINDING_API_KEY"),
+        base_url=os.getenv("EMBEDDING_BINDING_HOST"),
+    )
+
+
+rerank_model_func = partial(
+    cohere_rerank,
+    model=os.getenv("RERANK_MODEL", "rerank-v3.5"),
+    api_key=os.getenv("RERANK_BINDING_API_KEY"),
+    base_url=os.getenv("RERANK_BINDING_HOST", "https://api.cohere.com/v2/rerank"),
+    enable_chunking=os.getenv("RERANK_ENABLE_CHUNKING", "false").lower() == "true",
+    max_tokens_per_doc=int(os.getenv("RERANK_MAX_TOKENS_PER_DOC", "4096")),
+)
+
+
+async def create_rag_with_rerank():
+    """Create LightRAG instance with rerank configuration"""
+
+    # Get embedding dimension
+    test_embedding = await embedding_func(["test"])
+    embedding_dim = test_embedding.shape[1]
+    print(f"Detected embedding dimension: {embedding_dim}")
+
+    # Method 1: Using custom rerank function
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        embedding_func=EmbeddingFunc(
+            embedding_dim=embedding_dim,
+            max_token_size=8192,
+            func=embedding_func,
+        ),
+        # Rerank Configuration - provide the rerank function
+        rerank_model_func=rerank_model_func,
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+    return rag
+
+
+async def test_rerank_with_different_settings():
+    """
+    Test rerank functionality with different enable_rerank settings
+    """
+    print("\n\n🚀 Setting up LightRAG with Rerank functionality...")
+
+    rag = await create_rag_with_rerank()
+
+    # Insert sample documents
+    sample_docs = [
+        "Reranking improves retrieval quality by re-ordering documents based on relevance.",
+        "LightRAG is a powerful retrieval-augmented generation system with multiple query modes.",
+        "Vector databases enable efficient similarity search in high-dimensional embedding spaces.",
+        "Natural language processing has evolved with large language models and transformers.",
+        "Machine learning algorithms can learn patterns from data without explicit programming.",
+    ]
+
+    print("📄 Inserting sample documents...")
+    await rag.ainsert(sample_docs)
+
+    query = "How does reranking improve retrieval quality?"
+    print(f"\n🔍 Testing query: '{query}'")
+    print("=" * 80)
+
+    # Test with rerank enabled (default)
+    print("\n📊 Testing with enable_rerank=True (default):")
+    result_with_rerank = await rag.aquery(
+        query,
+        param=QueryParam(
+            mode="naive",
+            top_k=10,
+            chunk_top_k=5,
+            enable_rerank=True,  # Explicitly enable rerank
+        ),
+    )
+    print(f"   Result length: {len(result_with_rerank)} characters")
+    print(f"   Preview: {result_with_rerank[:100]}...")
+
+    # Test with rerank disabled
+    print("\n📊 Testing with enable_rerank=False:")
+    result_without_rerank = await rag.aquery(
+        query,
+        param=QueryParam(
+            mode="naive",
+            top_k=10,
+            chunk_top_k=5,
+            enable_rerank=False,  # Disable rerank
+        ),
+    )
+    print(f"   Result length: {len(result_without_rerank)} characters")
+    print(f"   Preview: {result_without_rerank[:100]}...")
+
+    # Test with default settings (enable_rerank defaults to True)
+    print("\n📊 Testing with default settings (enable_rerank defaults to True):")
+    result_default = await rag.aquery(
+        query, param=QueryParam(mode="naive", top_k=10, chunk_top_k=5)
+    )
+    print(f"   Result length: {len(result_default)} characters")
+    print(f"   Preview: {result_default[:100]}...")
+
+
+async def test_direct_rerank():
+    """Test rerank function directly"""
+    print("\n🔧 Direct Rerank API Test")
+    print("=" * 40)
+
+    documents = [
+        "Vector search finds semantically similar documents",
+        "LightRAG supports advanced reranking capabilities",
+        "Reranking significantly improves retrieval quality",
+        "Natural language processing with modern transformers",
+        "The quick brown fox jumps over the lazy dog",
+    ]
+
+    query = "rerank improve quality"
+    print(f"Query: '{query}'")
+    print(f"Documents: {len(documents)}")
+
+    try:
+        reranked_results = await rerank_model_func(
+            query=query,
+            documents=documents,
+            top_n=4,
+        )
+
+        print("\n✅ Rerank Results:")
+        i = 0
+        for result in reranked_results:
+            index = result["index"]
+            score = result["relevance_score"]
+            content = documents[index]
+            print(f"  {index}. Score: {score:.4f} | {content}...")
+            i += 1
+
+    except Exception as e:
+        print(f"❌ Rerank failed: {e}")
+
+
+async def main():
+    """Main example function"""
+    print("🎯 LightRAG Rerank Integration Example")
+    print("=" * 60)
+
+    try:
+        # Test direct rerank
+        await test_direct_rerank()
+
+        # Test rerank with different enable_rerank settings
+        await test_rerank_with_different_settings()
+
+        print("\n✅ Example completed successfully!")
+        print("\n💡 Key Points:")
+        print("   ✓ Rerank is now controlled per query via 'enable_rerank' parameter")
+        print("   ✓ Default value for enable_rerank is True")
+        print("   ✓ Rerank function is configured at LightRAG initialization")
+        print("   ✓ Per-query enable_rerank setting overrides default behavior")
+        print(
+            "   ✓ If enable_rerank=True but no rerank model is configured, a warning is issued"
+        )
+        print("   ✓ Monitor API usage and costs when using rerank services")
+
+    except Exception as e:
+        print(f"\n❌ Example failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/unofficial-sample/copy_llm_cache_to_another_storage.py

@@ -0,0 +1,114 @@
+"""
+Sometimes you need to switch a storage solution, but you want to save LLM token and time.
+This handy script helps you to copy the LLM caches from one storage solution to another.
+(Not all the storage impl are supported)
+"""
+
+import asyncio
+import logging
+import os
+from dotenv import load_dotenv
+
+from lightrag.kg.postgres_impl import PostgreSQLDB, PGKVStorage
+from lightrag.kg.json_kv_impl import JsonKVStorage
+from lightrag.namespace import NameSpace
+
+load_dotenv()
+ROOT_DIR = os.environ.get("ROOT_DIR")
+WORKING_DIR = f"{ROOT_DIR}/dickens"
+
+logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+# AGE
+os.environ["AGE_GRAPH_NAME"] = "chinese"
+
+postgres_db = PostgreSQLDB(
+    config={
+        "host": "localhost",
+        "port": 15432,
+        "user": "rag",
+        "password": "rag",
+        "database": "r2",
+    }
+)
+
+
+async def copy_from_postgres_to_json():
+    await postgres_db.initdb()
+
+    from_llm_response_cache = PGKVStorage(
+        namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
+        global_config={"embedding_batch_num": 6},
+        embedding_func=None,
+        db=postgres_db,
+    )
+
+    to_llm_response_cache = JsonKVStorage(
+        namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
+        global_config={"working_dir": WORKING_DIR},
+        embedding_func=None,
+    )
+
+    # Get all cache data using the new flattened structure
+    all_data = await from_llm_response_cache.get_all()
+
+    # Convert flattened data to hierarchical structure for JsonKVStorage
+    kv = {}
+    for flattened_key, cache_entry in all_data.items():
+        # Parse flattened key: {mode}:{cache_type}:{hash}
+        parts = flattened_key.split(":", 2)
+        if len(parts) == 3:
+            mode, cache_type, hash_value = parts
+            if mode not in kv:
+                kv[mode] = {}
+            kv[mode][hash_value] = cache_entry
+            print(f"Copying {flattened_key} -> {mode}[{hash_value}]")
+        else:
+            print(f"Skipping invalid key format: {flattened_key}")
+
+    await to_llm_response_cache.upsert(kv)
+    await to_llm_response_cache.index_done_callback()
+    print("Mission accomplished!")
+
+
+async def copy_from_json_to_postgres():
+    await postgres_db.initdb()
+
+    from_llm_response_cache = JsonKVStorage(
+        namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
+        global_config={"working_dir": WORKING_DIR},
+        embedding_func=None,
+    )
+
+    to_llm_response_cache = PGKVStorage(
+        namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
+        global_config={"embedding_batch_num": 6},
+        embedding_func=None,
+        db=postgres_db,
+    )
+
+    # Get all cache data from JsonKVStorage (hierarchical structure)
+    all_data = await from_llm_response_cache.get_all()
+
+    # Convert hierarchical data to flattened structure for PGKVStorage
+    flattened_data = {}
+    for mode, mode_data in all_data.items():
+        print(f"Processing mode: {mode}")
+        for hash_value, cache_entry in mode_data.items():
+            # Determine cache_type from cache entry or use default
+            cache_type = cache_entry.get("cache_type", "extract")
+            # Create flattened key: {mode}:{cache_type}:{hash}
+            flattened_key = f"{mode}:{cache_type}:{hash_value}"
+            flattened_data[flattened_key] = cache_entry
+            print(f"\tConverting {mode}[{hash_value}] -> {flattened_key}")
+
+    # Upsert the flattened data
+    await to_llm_response_cache.upsert(flattened_data)
+    print("Mission accomplished!")
+
+
+if __name__ == "__main__":
+    asyncio.run(copy_from_json_to_postgres())

examples/unofficial-sample/lightrag_bedrock_demo.py

@@ -0,0 +1,56 @@
+"""
+LightRAG meets Amazon Bedrock ⛰️
+"""
+
+import os
+import logging
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.bedrock import bedrock_complete, bedrock_embed
+from lightrag.utils import EmbeddingFunc
+
+import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
+
+logging.getLogger("aiobotocore").setLevel(logging.WARNING)
+
+WORKING_DIR = "./dickens"
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+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
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes 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"]:
+        print("\n+-" + "-" * len(mode) + "-+")
+        print(f"| {mode.capitalize()} |")
+        print("+-" + "-" * len(mode) + "-+\n")
+        print(
+            rag.query(
+                "What are the top themes in this story?", param=QueryParam(mode=mode)
+            )
+        )
+
+
+if __name__ == "__main__":
+    main()

examples/unofficial-sample/lightrag_cloudflare_demo.py

@@ -0,0 +1,354 @@
+import asyncio
+import os
+import inspect
+import logging
+import logging.config
+from lightrag import LightRAG, QueryParam
+from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
+
+import requests
+import numpy as np
+from dotenv import load_dotenv
+
+"""This code is a modified version of lightrag_openai_demo.py"""
+
+# ideally, as always, env!
+load_dotenv(dotenv_path=".env", override=False)
+
+
+"""    ----========= IMPORTANT CHANGE THIS! =========----    """
+cloudflare_api_key = "YOUR_API_KEY"
+account_id = "YOUR_ACCOUNT ID"  # This is unique to your Cloudflare account
+
+# Authomatically changes
+api_base_url = f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/"
+
+
+# choose an embedding model
+EMBEDDING_MODEL = "@cf/baai/bge-m3"
+# choose a generative model
+LLM_MODEL = "@cf/meta/llama-3.2-3b-instruct"
+
+WORKING_DIR = "../dickens"  # you can change output as desired
+
+
+# Cloudflare init
+class CloudflareWorker:
+    def __init__(
+        self,
+        cloudflare_api_key: str,
+        api_base_url: str,
+        llm_model_name: str,
+        embedding_model_name: str,
+        max_tokens: int = 4080,
+        max_response_tokens: int = 4080,
+    ):
+        self.cloudflare_api_key = cloudflare_api_key
+        self.api_base_url = api_base_url
+        self.llm_model_name = llm_model_name
+        self.embedding_model_name = embedding_model_name
+        self.max_tokens = max_tokens
+        self.max_response_tokens = max_response_tokens
+
+    async def _send_request(self, model_name: str, input_: dict, debug_log: str):
+        headers = {"Authorization": f"Bearer {self.cloudflare_api_key}"}
+
+        print(f"""
+        data sent to Cloudflare
+        ~~~~~~~~~~~
+        {debug_log}
+        """)
+
+        try:
+            response_raw = requests.post(
+                f"{self.api_base_url}{model_name}", headers=headers, json=input_
+            ).json()
+            print(f"""
+        Cloudflare worker responded with:
+        ~~~~~~~~~~~
+        {str(response_raw)}
+            """)
+            result = response_raw.get("result", {})
+
+            if "data" in result:  # Embedding case
+                return np.array(result["data"])
+
+            if "response" in result:  # LLM response
+                return result["response"]
+
+            raise ValueError("Unexpected Cloudflare response format")
+
+        except Exception as e:
+            print(f"""
+            Cloudflare API returned:
+            ~~~~~~~~~
+            Error: {e}
+            """)
+            input("Press Enter to continue...")
+            return None
+
+    async def query(self, prompt, system_prompt: str = "", **kwargs) -> str:
+        # since no caching is used and we don't want to mess with everything lightrag, pop the kwarg it is
+        kwargs.pop("hashing_kv", None)
+
+        message = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": prompt},
+        ]
+
+        input_ = {
+            "messages": message,
+            "max_tokens": self.max_tokens,
+            "response_token_limit": self.max_response_tokens,
+        }
+
+        return await self._send_request(
+            self.llm_model_name,
+            input_,
+            debug_log=f"\n- model used {self.llm_model_name}\n- system prompt: {system_prompt}\n- query: {prompt}",
+        )
+
+    async def embedding_chunk(self, texts: list[str]) -> np.ndarray:
+        print(f"""
+        TEXT inputted
+        ~~~~~
+        {texts}
+        """)
+
+        input_ = {
+            "text": texts,
+            "max_tokens": self.max_tokens,
+            "response_token_limit": self.max_response_tokens,
+        }
+
+        return await self._send_request(
+            self.embedding_model_name,
+            input_,
+            debug_log=f"\n-llm model name {self.embedding_model_name}\n- texts: {texts}",
+        )
+
+
+def configure_logging():
+    """Configure logging for the application"""
+
+    # Reset any existing handlers to ensure clean configuration
+    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
+        logger_instance = logging.getLogger(logger_name)
+        logger_instance.handlers = []
+        logger_instance.filters = []
+
+    # Get log directory path from environment variable or use current directory
+    log_dir = os.getenv("LOG_DIR", os.getcwd())
+    log_file_path = os.path.abspath(
+        os.path.join(log_dir, "lightrag_cloudflare_worker_demo.log")
+    )
+
+    print(f"\nLightRAG compatible demo log file: {log_file_path}\n")
+    os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
+
+    # Get log file max size and backup count from environment variables
+    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))  # Default 10MB
+    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))  # Default 5 backups
+
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "default": {
+                    "format": "%(levelname)s: %(message)s",
+                },
+                "detailed": {
+                    "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+                },
+            },
+            "handlers": {
+                "console": {
+                    "formatter": "default",
+                    "class": "logging.StreamHandler",
+                    "stream": "ext://sys.stderr",
+                },
+                "file": {
+                    "formatter": "detailed",
+                    "class": "logging.handlers.RotatingFileHandler",
+                    "filename": log_file_path,
+                    "maxBytes": log_max_bytes,
+                    "backupCount": log_backup_count,
+                    "encoding": "utf-8",
+                },
+            },
+            "loggers": {
+                "lightrag": {
+                    "handlers": ["console", "file"],
+                    "level": "INFO",
+                    "propagate": False,
+                },
+            },
+        }
+    )
+
+    # Set the logger level to INFO
+    logger.setLevel(logging.INFO)
+    # Enable verbose debug if needed
+    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
+
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+async def initialize_rag():
+    cloudflare_worker = CloudflareWorker(
+        cloudflare_api_key=cloudflare_api_key,
+        api_base_url=api_base_url,
+        embedding_model_name=EMBEDDING_MODEL,
+        llm_model_name=LLM_MODEL,
+    )
+
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        max_parallel_insert=2,
+        llm_model_func=cloudflare_worker.query,
+        llm_model_name=os.getenv("LLM_MODEL", LLM_MODEL),
+        summary_max_tokens=4080,
+        embedding_func=EmbeddingFunc(
+            embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
+            max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "2048")),
+            func=lambda texts: cloudflare_worker.embedding_chunk(
+                texts,
+            ),
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+    return rag
+
+
+async def print_stream(stream):
+    async for chunk in stream:
+        print(chunk, end="", flush=True)
+
+
+async def main():
+    try:
+        # Clear old data files
+        files_to_delete = [
+            "graph_chunk_entity_relation.graphml",
+            "kv_store_doc_status.json",
+            "kv_store_full_docs.json",
+            "kv_store_text_chunks.json",
+            "vdb_chunks.json",
+            "vdb_entities.json",
+            "vdb_relationships.json",
+        ]
+
+        for file in files_to_delete:
+            file_path = os.path.join(WORKING_DIR, file)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                print(f"Deleting old file:: {file_path}")
+
+        # Initialize RAG instance
+        rag = await initialize_rag()
+
+        # Test embedding function
+        test_text = ["This is a test string for embedding."]
+        embedding = await rag.embedding_func(test_text)
+        embedding_dim = embedding.shape[1]
+        print("\n=======================")
+        print("Test embedding function")
+        print("========================")
+        print(f"Test dict: {test_text}")
+        print(f"Detected embedding dimension: {embedding_dim}\n\n")
+
+        # Locate the location of what is needed to be added to the knowledge
+        # Can add several simultaneously by modifying code
+        with open("./book.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+
+        # Perform naive search
+        print("\n=====================")
+        print("Query mode: naive")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="naive", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform local search
+        print("\n=====================")
+        print("Query mode: local")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="local", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform global search
+        print("\n=====================")
+        print("Query mode: global")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="global", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform hybrid search
+        print("\n=====================")
+        print("Query mode: hybrid")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="hybrid", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        """ FOR TESTING (if you want to test straight away, after building. Uncomment this part"""
+
+        """
+        print("\n" + "=" * 60)
+        print("AI ASSISTANT READY!")
+        print("Ask questions about (your uploaded) regulations")
+        print("Type 'quit' to exit")
+        print("=" * 60)
+
+        while True:
+            question = input("\n🔥 Your question: ")
+
+            if question.lower() in ['quit', 'exit', 'bye']:
+                break
+
+            print("\nThinking...")
+            response = await rag.aquery(question, param=QueryParam(mode="hybrid"))
+            print(f"\nAnswer: {response}")
+
+        """
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    finally:
+        if rag:
+            await rag.llm_response_cache.index_done_callback()
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    # Configure logging before running the main function
+    configure_logging()
+    asyncio.run(main())
+    print("\nDone!")

examples/unofficial-sample/lightrag_embedding_prefixes.py

@@ -0,0 +1,235 @@
+import os
+import asyncio
+import inspect
+import logging
+import logging.config
+from functools import partial
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import openai_complete_if_cache
+from lightrag.llm.ollama import ollama_embed
+from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
+
+from dotenv import load_dotenv
+
+load_dotenv(dotenv_path=".env", override=False)
+
+WORKING_DIR = "./dickens"
+
+
+def configure_logging():
+    """Configure logging for the application"""
+
+    # Reset any existing handlers to ensure clean configuration
+    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
+        logger_instance = logging.getLogger(logger_name)
+        logger_instance.handlers = []
+        logger_instance.filters = []
+
+    # Get log directory path from environment variable or use current directory
+    log_dir = os.getenv("LOG_DIR", os.getcwd())
+    log_file_path = os.path.abspath(
+        os.path.join(log_dir, "lightrag_compatible_demo.log")
+    )
+
+    print(f"\nLightRAG compatible demo log file: {log_file_path}\n")
+    os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
+
+    # Get log file max size and backup count from environment variables
+    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))  # Default 10MB
+    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))  # Default 5 backups
+
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "default": {
+                    "format": "%(levelname)s: %(message)s",
+                },
+                "detailed": {
+                    "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+                },
+            },
+            "handlers": {
+                "console": {
+                    "formatter": "default",
+                    "class": "logging.StreamHandler",
+                    "stream": "ext://sys.stderr",
+                },
+                "file": {
+                    "formatter": "detailed",
+                    "class": "logging.handlers.RotatingFileHandler",
+                    "filename": log_file_path,
+                    "maxBytes": log_max_bytes,
+                    "backupCount": log_backup_count,
+                    "encoding": "utf-8",
+                },
+            },
+            "loggers": {
+                "lightrag": {
+                    "handlers": ["console", "file"],
+                    "level": "INFO",
+                    "propagate": False,
+                },
+            },
+        }
+    )
+
+    # Set the logger level to INFO
+    logger.setLevel(logging.INFO)
+    # Enable verbose debug if needed
+    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
+
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        os.getenv("LLM_MODEL", "deepseek-chat"),
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=os.getenv("LLM_BINDING_API_KEY") or os.getenv("OPENAI_API_KEY"),
+        base_url=os.getenv("LLM_BINDING_HOST", "https://api.deepseek.com"),
+        **kwargs,
+    )
+
+
+async def print_stream(stream):
+    async for chunk in stream:
+        if chunk:
+            print(chunk, end="", flush=True)
+
+
+async def initialize_rag():
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        # Note: ollama_embed is decorated with @wrap_embedding_func_with_attrs,
+        # which wraps it in an EmbeddingFunc. Using .func accesses the original
+        # unwrapped function to avoid double wrapping when we create our own
+        # EmbeddingFunc with custom configuration (embedding_dim, max_token_size and prefixes).
+        embedding_func=EmbeddingFunc(
+            embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
+            max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "8192")),
+            supports_asymmetric=True,
+            func=partial(
+                ollama_embed.func,  # Access the unwrapped function to avoid double EmbeddingFunc wrapping
+                embed_model=os.getenv("EMBEDDING_MODEL", "FRIDA:latest"),
+                host=os.getenv("EMBEDDING_BINDING_HOST", "http://localhost:11434"),
+                query_prefix=os.getenv("EMBEDDING_QUERY_PREFIX", "search_query: "),
+                document_prefix=os.getenv(
+                    "EMBEDDING_DOCUMENT_PREFIX", "search_document: "
+                ),
+            ),
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+    return rag
+
+
+async def main():
+    rag = None
+    try:
+        # Clear old data files
+        files_to_delete = [
+            "graph_chunk_entity_relation.graphml",
+            "kv_store_doc_status.json",
+            "kv_store_full_docs.json",
+            "kv_store_text_chunks.json",
+            "vdb_chunks.json",
+            "vdb_entities.json",
+            "vdb_relationships.json",
+        ]
+
+        for file in files_to_delete:
+            file_path = os.path.join(WORKING_DIR, file)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                print(f"Deleting old file:: {file_path}")
+
+        # Initialize RAG instance
+        rag = await initialize_rag()
+
+        # Test embedding function
+        test_text = ["This is a test string for embedding."]
+        embedding = await rag.embedding_func(test_text)
+        embedding_dim = embedding.shape[1]
+        print("\n=======================")
+        print("Test embedding function")
+        print("========================")
+        print(f"Test dict: {test_text}")
+        print(f"Detected embedding dimension: {embedding_dim}\n\n")
+
+        with open("./book.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+
+        # Perform naive search
+        print("\n=====================")
+        print("Query mode: naive")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="naive", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform local search
+        print("\n=====================")
+        print("Query mode: local")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="local", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform global search
+        print("\n=====================")
+        print("Query mode: global")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="global", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform hybrid search
+        print("\n=====================")
+        print("Query mode: hybrid")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="hybrid", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    finally:
+        if rag:
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    # Configure logging before running the main function
+    configure_logging()
+    asyncio.run(main())
+    print("\nDone!")

examples/unofficial-sample/lightrag_hf_demo.py

@@ -0,0 +1,79 @@
+import os
+
+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
+
+import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
+
+WORKING_DIR = "./dickens"
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+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",
+        embedding_func=EmbeddingFunc(
+            embedding_dim=384,
+            max_token_size=5000,
+            func=lambda texts: hf_embed(
+                texts,
+                tokenizer=AutoTokenizer.from_pretrained(
+                    "sentence-transformers/all-MiniLM-L6-v2"
+                ),
+                embed_model=AutoModel.from_pretrained(
+                    "sentence-transformers/all-MiniLM-L6-v2"
+                ),
+            ),
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes 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())
+
+    # 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()

examples/unofficial-sample/lightrag_llamaindex_direct_demo.py

@@ -0,0 +1,139 @@
+import os
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.llama_index_impl import (
+    llama_index_complete_if_cache,
+    llama_index_embed,
+)
+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()
+
+
+# Configure working directory
+WORKING_DIR = "./index_default"
+print(f"WORKING_DIR: {WORKING_DIR}")
+
+# Model configuration
+LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4")
+print(f"LLM_MODEL: {LLM_MODEL}")
+EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-large")
+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}")
+
+# OpenAI configuration
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "your-api-key-here")
+
+if not os.path.exists(WORKING_DIR):
+    print(f"Creating working directory: {WORKING_DIR}")
+    os.mkdir(WORKING_DIR)
+
+
+# Initialize LLM function
+async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
+    try:
+        # Initialize OpenAI if not in kwargs
+        if "llm_instance" not in kwargs:
+            llm_instance = OpenAI(
+                model=LLM_MODEL,
+                api_key=OPENAI_API_KEY,
+                temperature=0.7,
+            )
+            kwargs["llm_instance"] = llm_instance
+
+        response = await llama_index_complete_if_cache(
+            kwargs["llm_instance"],
+            prompt,
+            system_prompt=system_prompt,
+            history_messages=history_messages,
+            **kwargs,
+        )
+        return response
+    except Exception as e:
+        print(f"LLM request failed: {str(e)}")
+        raise
+
+
+# Initialize embedding function
+async def embedding_func(texts):
+    try:
+        embed_model = OpenAIEmbedding(
+            model=EMBEDDING_MODEL,
+            api_key=OPENAI_API_KEY,
+        )
+        return await llama_index_embed(texts, embed_model=embed_model)
+    except Exception as e:
+        print(f"Embedding failed: {str(e)}")
+        raise
+
+
+# Get embedding dimension
+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={embedding_dim}")
+    return embedding_dim
+
+
+async def initialize_rag():
+    embedding_dimension = await get_embedding_dim()
+
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        embedding_func=EmbeddingFunc(
+            embedding_dim=embedding_dimension,
+            max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
+            func=embedding_func,
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes 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
+    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/unofficial-sample/lightrag_llamaindex_litellm_demo.py

@@ -0,0 +1,141 @@
+import os
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.llama_index_impl import (
+    llama_index_complete_if_cache,
+    llama_index_embed,
+)
+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()
+
+
+# Configure working directory
+WORKING_DIR = "./index_default"
+print(f"WORKING_DIR: {WORKING_DIR}")
+
+# Model configuration
+LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4")
+print(f"LLM_MODEL: {LLM_MODEL}")
+EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-large")
+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}")
+
+# LiteLLM configuration
+LITELLM_URL = os.environ.get("LITELLM_URL", "http://localhost:4000")
+print(f"LITELLM_URL: {LITELLM_URL}")
+LITELLM_KEY = os.environ.get("LITELLM_KEY", "sk-1234")
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+# Initialize LLM function
+async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
+    try:
+        # Initialize LiteLLM if not in kwargs
+        if "llm_instance" not in kwargs:
+            llm_instance = LiteLLM(
+                model=f"openai/{LLM_MODEL}",  # Format: "provider/model_name"
+                api_base=LITELLM_URL,
+                api_key=LITELLM_KEY,
+                temperature=0.7,
+            )
+            kwargs["llm_instance"] = llm_instance
+
+        response = await llama_index_complete_if_cache(
+            kwargs["llm_instance"],
+            prompt,
+            system_prompt=system_prompt,
+            history_messages=history_messages,
+        )
+        return response
+    except Exception as e:
+        print(f"LLM request failed: {str(e)}")
+        raise
+
+
+# Initialize embedding function
+async def embedding_func(texts):
+    try:
+        embed_model = LiteLLMEmbedding(
+            model_name=f"openai/{EMBEDDING_MODEL}",
+            api_base=LITELLM_URL,
+            api_key=LITELLM_KEY,
+        )
+        return await llama_index_embed(texts, embed_model=embed_model)
+    except Exception as e:
+        print(f"Embedding failed: {str(e)}")
+        raise
+
+
+# Get embedding dimension
+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={embedding_dim}")
+    return embedding_dim
+
+
+async def initialize_rag():
+    embedding_dimension = await get_embedding_dim()
+
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        embedding_func=EmbeddingFunc(
+            embedding_dim=embedding_dimension,
+            max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
+            func=embedding_func,
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes 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
+    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/unofficial-sample/lightrag_llamaindex_litellm_opik_demo.py

@@ -0,0 +1,152 @@
+import os
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.llama_index_impl import (
+    llama_index_complete_if_cache,
+    llama_index_embed,
+)
+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()
+
+
+# Configure working directory
+WORKING_DIR = "./index_default"
+print(f"WORKING_DIR: {WORKING_DIR}")
+
+# Model configuration
+LLM_MODEL = os.environ.get("LLM_MODEL", "gemma-3-4b")
+print(f"LLM_MODEL: {LLM_MODEL}")
+EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "arctic-embed")
+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}")
+
+# LiteLLM configuration
+LITELLM_URL = os.environ.get("LITELLM_URL", "http://localhost:4000")
+print(f"LITELLM_URL: {LITELLM_URL}")
+LITELLM_KEY = os.environ.get("LITELLM_KEY", "sk-4JdvGFKqSA3S0k_5p0xufw")
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+# Initialize LLM function
+async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
+    try:
+        # Initialize LiteLLM if not in kwargs
+        if "llm_instance" not in kwargs:
+            llm_instance = LiteLLM(
+                model=f"openai/{LLM_MODEL}",  # Format: "provider/model_name"
+                api_base=LITELLM_URL,
+                api_key=LITELLM_KEY,
+                temperature=0.7,
+            )
+            kwargs["llm_instance"] = llm_instance
+
+        chat_kwargs = {}
+        chat_kwargs["litellm_params"] = {
+            "metadata": {
+                "opik": {
+                    "project_name": "lightrag_llamaindex_litellm_opik_demo",
+                    "tags": ["lightrag", "litellm"],
+                }
+            }
+        }
+
+        response = await llama_index_complete_if_cache(
+            kwargs["llm_instance"],
+            prompt,
+            system_prompt=system_prompt,
+            history_messages=history_messages,
+            chat_kwargs=chat_kwargs,
+        )
+        return response
+    except Exception as e:
+        print(f"LLM request failed: {str(e)}")
+        raise
+
+
+# Initialize embedding function
+async def embedding_func(texts):
+    try:
+        embed_model = LiteLLMEmbedding(
+            model_name=f"openai/{EMBEDDING_MODEL}",
+            api_base=LITELLM_URL,
+            api_key=LITELLM_KEY,
+        )
+        return await llama_index_embed(texts, embed_model=embed_model)
+    except Exception as e:
+        print(f"Embedding failed: {str(e)}")
+        raise
+
+
+# Get embedding dimension
+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={embedding_dim}")
+    return embedding_dim
+
+
+async def initialize_rag():
+    embedding_dimension = await get_embedding_dim()
+
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        embedding_func=EmbeddingFunc(
+            embedding_dim=embedding_dimension,
+            max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
+            func=embedding_func,
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes 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
+    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/unofficial-sample/lightrag_lmdeploy_demo.py

@@ -0,0 +1,107 @@
+import os
+
+from lightrag import LightRAG, QueryParam
+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
+
+import asyncio
+import nest_asyncio
+
+nest_asyncio.apply()
+
+WORKING_DIR = "./dickens"
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+async def lmdeploy_model_complete(
+    prompt=None,
+    system_prompt=None,
+    history_messages=[],
+    keyword_extraction=False,
+    **kwargs,
+) -> str:
+    model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
+    return await lmdeploy_model_if_cache(
+        model_name,
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        ## please specify chat_template if your local path does not follow original HF file name,
+        ## or model_name is a pytorch model on huggingface.co,
+        ## you can refer to https://github.com/InternLM/lmdeploy/blob/main/lmdeploy/model.py
+        ## for a list of chat_template available in lmdeploy.
+        chat_template="llama3",
+        # model_format ='awq', # if you are using awq quantization model.
+        # quant_policy=8, # if you want to use online kv cache, 4=kv int4, 8=kv int8.
+        **kwargs,
+    )
+
+
+async def initialize_rag():
+    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
+        embedding_func=EmbeddingFunc(
+            embedding_dim=384,
+            max_token_size=5000,
+            func=lambda texts: hf_embed(
+                texts,
+                tokenizer=AutoTokenizer.from_pretrained(
+                    "sentence-transformers/all-MiniLM-L6-v2"
+                ),
+                embed_model=AutoModel.from_pretrained(
+                    "sentence-transformers/all-MiniLM-L6-v2"
+                ),
+            ),
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes 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
+    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/unofficial-sample/lightrag_nvidia_demo.py

@@ -0,0 +1,168 @@
+import os
+import asyncio
+import nest_asyncio
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm import (
+    openai_complete_if_cache,
+    nvidia_openai_embed,
+)
+from lightrag.utils import EmbeddingFunc
+import numpy as np
+
+# for custom llm_model_func
+from lightrag.utils import locate_json_string_body_from_string
+
+nest_asyncio.apply()
+
+WORKING_DIR = "./dickens"
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+# some method to use your API key (choose one)
+# NVIDIA_OPENAI_API_KEY = os.getenv("NVIDIA_OPENAI_API_KEY")
+NVIDIA_OPENAI_API_KEY = "nvapi-xxxx"  # your api key
+
+# using pre-defined function for nvidia LLM API. OpenAI compatible
+# llm_model_func = nvidia_openai_complete
+
+
+# If you trying to make custom llm_model_func to use llm model on NVIDIA API like other example:
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    result = await openai_complete_if_cache(
+        "nvidia/llama-3.1-nemotron-70b-instruct",
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=NVIDIA_OPENAI_API_KEY,
+        base_url="https://integrate.api.nvidia.com/v1",
+        **kwargs,
+    )
+    if keyword_extraction:
+        return locate_json_string_body_from_string(result)
+    return result
+
+
+# custom embedding
+nvidia_embed_model = "nvidia/nv-embedqa-e5-v5"
+
+
+async def indexing_embedding_func(texts: list[str]) -> np.ndarray:
+    return await nvidia_openai_embed(
+        texts,
+        model=nvidia_embed_model,  # maximum 512 token
+        # model="nvidia/llama-3.2-nv-embedqa-1b-v1",
+        api_key=NVIDIA_OPENAI_API_KEY,
+        base_url="https://integrate.api.nvidia.com/v1",
+        input_type="passage",
+        trunc="END",  # handling on server side if input token is longer than maximum token
+        encode="float",
+    )
+
+
+async def query_embedding_func(texts: list[str]) -> np.ndarray:
+    return await nvidia_openai_embed(
+        texts,
+        model=nvidia_embed_model,  # maximum 512 token
+        # model="nvidia/llama-3.2-nv-embedqa-1b-v1",
+        api_key=NVIDIA_OPENAI_API_KEY,
+        base_url="https://integrate.api.nvidia.com/v1",
+        input_type="query",
+        trunc="END",  # handling on server side if input token is longer than maximum token
+        encode="float",
+    )
+
+
+# dimension are same
+async def get_embedding_dim():
+    test_text = ["This is a test sentence."]
+    embedding = await indexing_embedding_func(test_text)
+    embedding_dim = embedding.shape[1]
+    return embedding_dim
+
+
+# function test
+async def test_funcs():
+    result = await llm_model_func("How are you?")
+    print("llm_model_func: ", result)
+
+    result = await indexing_embedding_func(["How are you?"])
+    print("embedding_func: ", result)
+
+
+# asyncio.run(test_funcs())
+
+
+async def initialize_rag():
+    embedding_dimension = await get_embedding_dim()
+    print(f"Detected embedding dimension: {embedding_dimension}")
+
+    # lightRAG class during indexing
+    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,  # maximum token size, somehow it's still exceed maximum number of token
+            # so truncate (trunc) parameter on embedding_func will handle it and try to examine the tokenizer used in LightRAG
+            # so you can adjust to be able to fit the NVIDIA model (future work)
+            func=indexing_embedding_func,
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+    return rag
+
+
+async def main():
+    try:
+        # Initialize RAG instance
+        rag = await initialize_rag()
+
+        # reading file
+        with open("./book.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+
+        # Perform naive search
+        print("==============Naive===============")
+        print(
+            await rag.aquery(
+                "What are the top themes in this story?", param=QueryParam(mode="naive")
+            )
+        )
+
+        # Perform local search
+        print("==============local===============")
+        print(
+            await rag.aquery(
+                "What are the top themes in this story?", param=QueryParam(mode="local")
+            )
+        )
+
+        # Perform global search
+        print("==============global===============")
+        print(
+            await rag.aquery(
+                "What are the top themes in this story?",
+                param=QueryParam(mode="global"),
+            )
+        )
+
+        # Perform hybrid search
+        print("==============hybrid===============")
+        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/unofficial-sample/lightrag_openai_neo4j_milvus_redis_demo.py

@@ -0,0 +1,109 @@
+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
+
+# WorkingDir
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+WORKING_DIR = os.path.join(ROOT_DIR, "myKG")
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+print(f"WorkingDir: {WORKING_DIR}")
+
+# redis
+os.environ["REDIS_URI"] = "redis://localhost:6379"
+
+# neo4j
+BATCH_SIZE_NODES = 500
+BATCH_SIZE_EDGES = 100
+os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
+os.environ["NEO4J_USERNAME"] = "neo4j"
+os.environ["NEO4J_PASSWORD"] = "12345678"
+
+# milvus
+os.environ["MILVUS_URI"] = "http://localhost:19530"
+os.environ["MILVUS_USER"] = "root"
+os.environ["MILVUS_PASSWORD"] = "Milvus"
+os.environ["MILVUS_DB_NAME"] = "lightrag"
+
+
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        "deepseek-chat",
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key="",
+        base_url="",
+        **kwargs,
+    )
+
+
+embedding_func = EmbeddingFunc(
+    embedding_dim=768,
+    max_token_size=512,
+    func=lambda texts: ollama_embed(
+        texts, embed_model="shaw/dmeta-embedding-zh", host="http://117.50.173.35:11434"
+    ),
+)
+
+
+async def initialize_rag():
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        summary_max_tokens=10000,
+        embedding_func=embedding_func,
+        chunk_token_size=512,
+        chunk_overlap_token_size=256,
+        kv_storage="RedisKVStorage",
+        graph_storage="Neo4JStorage",
+        vector_storage="MilvusVectorDBStorage",
+        doc_status_storage="RedisKVStorage",
+    )
+
+    await rag.initialize_storages()  # Auto-initializes 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()