gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix linting

Browse Source
This commit is contained in:
zrguo
2025-03-03 18:40:03 +08:00
parent 1611400854
commit ef2a5ad191
33 changed files with 320 additions and 1411 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
examples/lightrag_api_ollama_demo.py
View File

@@ -36,7 +36,10 @@ async def init():
llm_model_name="gemma2:9b", llm_model_name="gemma2:9b",
llm_model_max_async=4, llm_model_max_async=4,
llm_model_max_token_size=8192, llm_model_max_token_size=8192,
llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 8192}}, llm_model_kwargs={
"host": "http://localhost:11434",
"options": {"num_ctx": 8192},
},
embedding_func=EmbeddingFunc( embedding_func=EmbeddingFunc(
embedding_dim=768, embedding_dim=768,
max_token_size=8192, max_token_size=8192,
@@ -64,6 +67,8 @@ async def lifespan(app: FastAPI):
app = FastAPI( app = FastAPI(
title="LightRAG API", description="API for RAG operations", lifespan=lifespan title="LightRAG API", description="API for RAG operations", lifespan=lifespan
) )
# Data models # Data models
class QueryRequest(BaseModel): class QueryRequest(BaseModel):
query: str query: str

5
examples/lightrag_api_openai_compatible_demo.py
View File

@@ -75,7 +75,7 @@ async def get_embedding_dim():
# Initialize RAG instance # Initialize RAG instance
async def init(): async def init():
embedding_dimension = await get_embedding_dim() embedding_dimension = await get_embedding_dim()
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
llm_model_func=llm_model_func, llm_model_func=llm_model_func,
@@ -88,9 +88,10 @@ async def init():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
global rag global rag

8
examples/lightrag_bedrock_demo.py
View File

@@ -21,6 +21,7 @@ WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -33,9 +34,10 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
def main(): def main():
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -47,5 +49,7 @@ def main():
print(f"| {mode.capitalize()} |") print(f"| {mode.capitalize()} |")
print("+-" + "-" * len(mode) + "-+\n") print("+-" + "-" * len(mode) + "-+\n")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode=mode)) rag.query(
"What are the top themes in this story?", param=QueryParam(mode=mode)
)
) )

5
examples/lightrag_gemini_demo.py
View File

@@ -12,6 +12,7 @@ from lightrag.kg.shared_storage import initialize_pipeline_status
import asyncio import asyncio
import nest_asyncio import nest_asyncio
# Apply nest_asyncio to solve event loop issues # Apply nest_asyncio to solve event loop issues
nest_asyncio.apply() nest_asyncio.apply()
@@ -79,9 +80,10 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -98,5 +100,6 @@ def main():
print(response) print(response)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

19
examples/lightrag_hf_demo.py
View File

@@ -16,6 +16,7 @@ WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -41,6 +42,7 @@ async def initialize_rag():
return rag return rag
def main(): def main():
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -49,23 +51,32 @@ def main():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

21
examples/lightrag_llamaindex_direct_demo.py
View File

@@ -83,7 +83,7 @@ async def get_embedding_dim():
async def initialize_rag(): async def initialize_rag():
embedding_dimension = await get_embedding_dim() embedding_dimension = await get_embedding_dim()
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
llm_model_func=llm_model_func, llm_model_func=llm_model_func,
@@ -96,7 +96,7 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
@@ -111,23 +111,32 @@ def main():
# Test different query modes # Test different query modes
print("\nNaive Search:") print("\nNaive Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
print("\nLocal Search:") print("\nLocal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
print("\nGlobal Search:") print("\nGlobal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
print("\nHybrid Search:") print("\nHybrid Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

21
examples/lightrag_llamaindex_litellm_demo.py
View File

@@ -86,7 +86,7 @@ async def get_embedding_dim():
async def initialize_rag(): async def initialize_rag():
embedding_dimension = await get_embedding_dim() embedding_dimension = await get_embedding_dim()
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
llm_model_func=llm_model_func, llm_model_func=llm_model_func,
@@ -99,7 +99,7 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
@@ -114,23 +114,32 @@ def main():
# Test different query modes # Test different query modes
print("\nNaive Search:") print("\nNaive Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
print("\nLocal Search:") print("\nLocal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
print("\nGlobal Search:") print("\nGlobal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
print("\nHybrid Search:") print("\nHybrid Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

23
examples/lightrag_lmdeploy_demo.py
View File

@@ -41,6 +41,7 @@ async def lmdeploy_model_complete(
**kwargs, **kwargs,
) )
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -63,9 +64,10 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -77,23 +79,32 @@ def main():
# Test different query modes # Test different query modes
print("\nNaive Search:") print("\nNaive Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
print("\nLocal Search:") print("\nLocal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
print("\nGlobal Search:") print("\nGlobal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
print("\nHybrid Search:") print("\nHybrid Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

5
examples/lightrag_nvidia_demo.py
View File

@@ -97,6 +97,7 @@ async def test_funcs():
# asyncio.run(test_funcs()) # asyncio.run(test_funcs())
async def initialize_rag(): async def initialize_rag():
embedding_dimension = await get_embedding_dim() embedding_dimension = await get_embedding_dim()
print(f"Detected embedding dimension: {embedding_dimension}") print(f"Detected embedding dimension: {embedding_dimension}")
@@ -117,8 +118,10 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance

27
examples/lightrag_ollama_age_demo.py
View File

@@ -27,6 +27,7 @@ os.environ["AGE_POSTGRES_HOST"] = "localhost"
os.environ["AGE_POSTGRES_PORT"] = "5455" os.environ["AGE_POSTGRES_PORT"] = "5455"
os.environ["AGE_GRAPH_NAME"] = "dickens" os.environ["AGE_GRAPH_NAME"] = "dickens"
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -34,7 +35,10 @@ async def initialize_rag():
llm_model_name="llama3.1:8b", llm_model_name="llama3.1:8b",
llm_model_max_async=4, llm_model_max_async=4,
llm_model_max_token_size=32768, llm_model_max_token_size=32768,
llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}}, llm_model_kwargs={
"host": "http://localhost:11434",
"options": {"num_ctx": 32768},
},
embedding_func=EmbeddingFunc( embedding_func=EmbeddingFunc(
embedding_dim=768, embedding_dim=768,
max_token_size=8192, max_token_size=8192,
@@ -47,13 +51,15 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
async def print_stream(stream): async def print_stream(stream):
async for chunk in stream: async for chunk in stream:
print(chunk, end="", flush=True) print(chunk, end="", flush=True)
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -65,22 +71,30 @@ def main():
# Test different query modes # Test different query modes
print("\nNaive Search:") print("\nNaive Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
print("\nLocal Search:") print("\nLocal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
print("\nGlobal Search:") print("\nGlobal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
print("\nHybrid Search:") print("\nHybrid Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
# stream response # stream response
@@ -94,5 +108,6 @@ def main():
else: else:
print(resp) print(resp)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

27
examples/lightrag_ollama_demo.py
View File

@@ -17,6 +17,7 @@ logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -24,7 +25,10 @@ async def initialize_rag():
llm_model_name="gemma2:2b", llm_model_name="gemma2:2b",
llm_model_max_async=4, llm_model_max_async=4,
llm_model_max_token_size=32768, llm_model_max_token_size=32768,
llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}}, llm_model_kwargs={
"host": "http://localhost:11434",
"options": {"num_ctx": 32768},
},
embedding_func=EmbeddingFunc( embedding_func=EmbeddingFunc(
embedding_dim=768, embedding_dim=768,
max_token_size=8192, max_token_size=8192,
@@ -36,13 +40,15 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
async def print_stream(stream): async def print_stream(stream):
async for chunk in stream: async for chunk in stream:
print(chunk, end="", flush=True) print(chunk, end="", flush=True)
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -54,22 +60,30 @@ def main():
# Test different query modes # Test different query modes
print("\nNaive Search:") print("\nNaive Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
print("\nLocal Search:") print("\nLocal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
print("\nGlobal Search:") print("\nGlobal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
print("\nHybrid Search:") print("\nHybrid Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
# stream response # stream response
@@ -83,5 +97,6 @@ def main():
else: else:
print(resp) print(resp)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

27
examples/lightrag_ollama_gremlin_demo.py
View File

@@ -32,6 +32,7 @@ os.environ["GREMLIN_TRAVERSE_SOURCE"] = "g"
os.environ["GREMLIN_USER"] = "" os.environ["GREMLIN_USER"] = ""
os.environ["GREMLIN_PASSWORD"] = "" os.environ["GREMLIN_PASSWORD"] = ""
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -39,7 +40,10 @@ async def initialize_rag():
llm_model_name="llama3.1:8b", llm_model_name="llama3.1:8b",
llm_model_max_async=4, llm_model_max_async=4,
llm_model_max_token_size=32768, llm_model_max_token_size=32768,
llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}}, llm_model_kwargs={
"host": "http://localhost:11434",
"options": {"num_ctx": 32768},
},
embedding_func=EmbeddingFunc( embedding_func=EmbeddingFunc(
embedding_dim=768, embedding_dim=768,
max_token_size=8192, max_token_size=8192,
@@ -52,13 +56,15 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
async def print_stream(stream): async def print_stream(stream):
async for chunk in stream: async for chunk in stream:
print(chunk, end="", flush=True) print(chunk, end="", flush=True)
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -70,22 +76,30 @@ def main():
# Test different query modes # Test different query modes
print("\nNaive Search:") print("\nNaive Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
print("\nLocal Search:") print("\nLocal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
print("\nGlobal Search:") print("\nGlobal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
print("\nHybrid Search:") print("\nHybrid Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
# stream response # stream response
@@ -99,5 +113,6 @@ def main():
else: else:
print(resp) print(resp)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

26
examples/lightrag_ollama_neo4j_milvus_mongo_demo.py
View File

@@ -32,6 +32,7 @@ os.environ["MILVUS_USER"] = "root"
os.environ["MILVUS_PASSWORD"] = "root" os.environ["MILVUS_PASSWORD"] = "root"
os.environ["MILVUS_DB_NAME"] = "lightrag" os.environ["MILVUS_DB_NAME"] = "lightrag"
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -39,7 +40,10 @@ async def initialize_rag():
llm_model_name="qwen2.5:14b", llm_model_name="qwen2.5:14b",
llm_model_max_async=4, llm_model_max_async=4,
llm_model_max_token_size=32768, llm_model_max_token_size=32768,
llm_model_kwargs={"host": "http://127.0.0.1:11434", "options": {"num_ctx": 32768}}, llm_model_kwargs={
"host": "http://127.0.0.1:11434",
"options": {"num_ctx": 32768},
},
embedding_func=EmbeddingFunc( embedding_func=EmbeddingFunc(
embedding_dim=1024, embedding_dim=1024,
max_token_size=8192, max_token_size=8192,
@@ -54,9 +58,10 @@ async def initialize_rag():
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -68,23 +73,32 @@ def main():
# Test different query modes # Test different query modes
print("\nNaive Search:") print("\nNaive Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
print("\nLocal Search:") print("\nLocal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
print("\nGlobal Search:") print("\nGlobal Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
print("\nHybrid Search:") print("\nHybrid Search:")
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

3
examples/lightrag_openai_compatible_demo.py
View File

@@ -53,6 +53,7 @@ async def test_funcs():
# asyncio.run(test_funcs()) # asyncio.run(test_funcs())
async def initialize_rag(): async def initialize_rag():
embedding_dimension = await get_embedding_dim() embedding_dimension = await get_embedding_dim()
print(f"Detected embedding dimension: {embedding_dimension}") print(f"Detected embedding dimension: {embedding_dimension}")
@@ -71,6 +72,8 @@ async def initialize_rag():
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance

2
examples/lightrag_openai_compatible_demo_embedding_cache.py
View File

@@ -53,6 +53,7 @@ async def test_funcs():
# asyncio.run(test_funcs()) # asyncio.run(test_funcs())
async def initialize_rag(): async def initialize_rag():
embedding_dimension = await get_embedding_dim() embedding_dimension = await get_embedding_dim()
print(f"Detected embedding dimension: {embedding_dimension}") print(f"Detected embedding dimension: {embedding_dimension}")
@@ -76,6 +77,7 @@ async def initialize_rag():
return rag return rag
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance

6
examples/lightrag_openai_compatible_stream_demo.py
View File

@@ -15,6 +15,8 @@ if not os.path.exists(WORKING_DIR):
print(f"WorkingDir: {WORKING_DIR}") print(f"WorkingDir: {WORKING_DIR}")
api_key = "empty" api_key = "empty"
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -40,11 +42,13 @@ async def initialize_rag():
return rag return rag
async def print_stream(stream): async def print_stream(stream):
async for chunk in stream: async for chunk in stream:
if chunk: if chunk:
print(chunk, end="", flush=True) print(chunk, end="", flush=True)
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -63,6 +67,6 @@ def main():
else: else:
print(resp) print(resp)
if __name__ == "__main__": if __name__ == "__main__":
main() main()

20
examples/lightrag_openai_demo.py
View File

@@ -9,6 +9,7 @@ WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -22,6 +23,7 @@ async def initialize_rag():
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -31,24 +33,32 @@ def main():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

19
examples/lightrag_openai_mongodb_graph_demo.py
View File

@@ -76,23 +76,32 @@ def main():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

19
examples/lightrag_openai_neo4j_milvus_redis_demo.py
View File

@@ -50,6 +50,8 @@ embedding_func = EmbeddingFunc(
texts, embed_model="shaw/dmeta-embedding-zh", host="http://117.50.173.35:11434" texts, embed_model="shaw/dmeta-embedding-zh", host="http://117.50.173.35:11434"
), ),
) )
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -79,23 +81,32 @@ def main():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

2
examples/lightrag_oracle_demo.py
View File

@@ -64,6 +64,7 @@ async def get_embedding_dim():
embedding_dim = embedding.shape[1] embedding_dim = embedding.shape[1]
return embedding_dim return embedding_dim
async def initialize_rag(): async def initialize_rag():
# Detect embedding dimension # Detect embedding dimension
embedding_dimension = await get_embedding_dim() embedding_dimension = await get_embedding_dim()
@@ -102,6 +103,7 @@ async def initialize_rag():
return rag return rag
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance

19
examples/lightrag_siliconcloud_demo.py
View File

@@ -47,6 +47,7 @@ async def test_funcs():
asyncio.run(test_funcs()) asyncio.run(test_funcs())
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -71,24 +72,32 @@ def main():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

2
examples/lightrag_tidb_demo.py
View File

@@ -55,6 +55,7 @@ async def get_embedding_dim():
embedding_dim = embedding.shape[1] embedding_dim = embedding.shape[1]
return embedding_dim return embedding_dim
async def initialize_rag(): async def initialize_rag():
# Detect embedding dimension # Detect embedding dimension
embedding_dimension = await get_embedding_dim() embedding_dimension = await get_embedding_dim()
@@ -82,6 +83,7 @@ async def initialize_rag():
return rag return rag
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance

19
examples/lightrag_zhipu_demo.py
View File

@@ -19,6 +19,7 @@ api_key = os.environ.get("ZHIPUAI_API_KEY")
if api_key is None: if api_key is None:
raise Exception("Please set ZHIPU_API_KEY in your environment") raise Exception("Please set ZHIPU_API_KEY in your environment")
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -38,6 +39,7 @@ async def initialize_rag():
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -47,23 +49,32 @@ def main():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

4
examples/lightrag_zhipu_postgres_demo.py
View File

@@ -28,6 +28,7 @@ os.environ["POSTGRES_USER"] = "rag"
os.environ["POSTGRES_PASSWORD"] = "rag" os.environ["POSTGRES_PASSWORD"] = "rag"
os.environ["POSTGRES_DATABASE"] = "rag" os.environ["POSTGRES_DATABASE"] = "rag"
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -55,8 +56,9 @@ async def initialize_rag():
return rag return rag
async def main(): async def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
# add embedding_func for graph database, it's deleted in commit 5661d76860436f7bf5aef2e50d9ee4a59660146c # add embedding_func for graph database, it's deleted in commit 5661d76860436f7bf5aef2e50d9ee4a59660146c

4
examples/query_keyword_separation_example.py
View File

@@ -80,6 +80,8 @@ async def test_funcs():
asyncio.run(test_funcs()) asyncio.run(test_funcs())
embedding_dimension = 3072 embedding_dimension = 3072
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -101,7 +103,7 @@ async def initialize_rag():
async def run_example(): async def run_example():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
book1 = open("./book_1.txt", encoding="utf-8") book1 = open("./book_1.txt", encoding="utf-8")
book2 = open("./book_2.txt", encoding="utf-8") book2 = open("./book_2.txt", encoding="utf-8")

21
examples/test.py
View File

@@ -1,4 +1,5 @@
import os import os
import asyncio
from lightrag import LightRAG, QueryParam from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import gpt_4o_mini_complete from lightrag.llm.openai import gpt_4o_mini_complete
from lightrag.kg.shared_storage import initialize_pipeline_status from lightrag.kg.shared_storage import initialize_pipeline_status
@@ -13,6 +14,7 @@ WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -35,23 +37,32 @@ def main():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

20
examples/test_chromadb.py
View File

@@ -112,12 +112,13 @@ async def initialize_rag():
}, },
) )
await rag.initialize_storages() await rag.initialize_storages()
await initialize_pipeline_status() await initialize_pipeline_status()
return rag return rag
def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -126,23 +127,32 @@ async def initialize_rag():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

3
examples/test_faiss.py
View File

@@ -58,6 +58,7 @@ async def embedding_func(texts: list[str]) -> np.ndarray:
embeddings = model.encode(texts, convert_to_numpy=True) embeddings = model.encode(texts, convert_to_numpy=True)
return embeddings return embeddings
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -78,8 +79,8 @@ async def initialize_rag():
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
# Insert the custom chunks into LightRAG # Insert the custom chunks into LightRAG

19
examples/test_neo4j.py
View File

@@ -15,6 +15,7 @@ WORKING_DIR = "./local_neo4jWorkDir"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -29,6 +30,7 @@ async def initialize_rag():
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())
@@ -38,23 +40,32 @@ def main():
# Perform naive search # Perform naive search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="naive")
)
) )
# Perform local search # Perform local search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="local")
)
) )
# Perform global search # Perform global search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global")
)
) )
# Perform hybrid search # Perform hybrid search
print( print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
)
) )
if __name__ == "__main__": if __name__ == "__main__":
main() main()

1313
examples/test_split_by_character.ipynb
View File

File diff suppressed because it is too large Load Diff

11
examples/vram_management_demo.py
View File

@@ -14,6 +14,7 @@ TEXT_FILES_DIR = "/llm/mt"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
# Initialize LightRAG # Initialize LightRAG
rag = LightRAG( rag = LightRAG(
@@ -31,6 +32,7 @@ async def initialize_rag():
return rag return rag
# Read all .txt files from the TEXT_FILES_DIR directory # Read all .txt files from the TEXT_FILES_DIR directory
texts = [] texts = []
for filename in os.listdir(TEXT_FILES_DIR): for filename in os.listdir(TEXT_FILES_DIR):
@@ -82,7 +84,8 @@ def main():
try: try:
print( print(
rag.query( rag.query(
"What are the top themes in this story?", param=QueryParam(mode="global") "What are the top themes in this story?",
param=QueryParam(mode="global"),
) )
) )
except Exception as e: except Exception as e:
@@ -91,18 +94,17 @@ def main():
try: try:
print( print(
rag.query( rag.query(
"What are the top themes in this story?", param=QueryParam(mode="hybrid") "What are the top themes in this story?",
param=QueryParam(mode="hybrid"),
) )
) )
except Exception as e: except Exception as e:
print(f"Error performing hybrid search: {e}") print(f"Error performing hybrid search: {e}")
# Function to clear VRAM resources # Function to clear VRAM resources
def clear_vram(): def clear_vram():
os.system("sudo nvidia-smi --gpu-reset") os.system("sudo nvidia-smi --gpu-reset")
# Regularly clear VRAM to prevent overflow # Regularly clear VRAM to prevent overflow
clear_vram_interval = 3600 # Clear once every hour clear_vram_interval = 3600 # Clear once every hour
start_time = time.time() start_time = time.time()
@@ -114,5 +116,6 @@ def main():
start_time = current_time start_time = current_time
time.sleep(60) # Check the time every minute time.sleep(60) # Check the time every minute
if __name__ == "__main__": if __name__ == "__main__":
main() main()

2
reproduce/Step_1.py
View File

@@ -31,6 +31,7 @@ WORKING_DIR = f"../{cls}"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
rag = LightRAG(working_dir=WORKING_DIR) rag = LightRAG(working_dir=WORKING_DIR)
@@ -39,6 +40,7 @@ async def initialize_rag():
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())

2
reproduce/Step_1_openai_compatible.py
View File

@@ -62,6 +62,7 @@ WORKING_DIR = f"../{cls}"
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
async def initialize_rag(): async def initialize_rag():
rag = LightRAG( rag = LightRAG(
working_dir=WORKING_DIR, working_dir=WORKING_DIR,
@@ -76,6 +77,7 @@ async def initialize_rag():
return rag return rag
def main(): def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = asyncio.run(initialize_rag())