Merge branch 'patch-6'

Author: danielaskdd

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!")