From 6c8fa9521477b3a9440f640337b060fa46f2f5c8 Mon Sep 17 00:00:00 2001 From: zrguo Date: Tue, 4 Mar 2025 12:25:07 +0800 Subject: [PATCH] fix demo --- README.md | 21 ++++--- examples/lightrag_azure_openai_demo.py | 58 +++++++++++-------- examples/lightrag_bedrock_demo.py | 4 ++ examples/lightrag_nvidia_demo.py | 2 +- examples/lightrag_openai_compatible_demo.py | 2 +- ..._openai_compatible_demo_embedding_cache.py | 2 +- examples/lightrag_oracle_demo.py | 2 +- examples/lightrag_tidb_demo.py | 2 +- examples/lightrag_zhipu_postgres_demo.py | 2 +- examples/query_keyword_separation_example.py | 2 +- 10 files changed, 58 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 5e8c5a94..f863d9ed 100644 --- a/README.md +++ b/README.md @@ -655,16 +655,19 @@ setup_logger("lightrag", level="INFO") # Note: Default settings use NetworkX # Initialize LightRAG with Neo4J implementation. -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model - graph_storage="Neo4JStorage", #<-----------override KG default -) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model + graph_storage="Neo4JStorage", #<-----------override KG default + ) -# Initialize database connections -await rag.initialize_storages() -# Initialize pipeline status for document processing -await initialize_pipeline_status() + # Initialize database connections + await rag.initialize_storages() + # Initialize pipeline status for document processing + await initialize_pipeline_status() + + return rag ``` see test_neo4j.py for a working example. diff --git a/examples/lightrag_azure_openai_demo.py b/examples/lightrag_azure_openai_demo.py index e0840366..c101383d 100644 --- a/examples/lightrag_azure_openai_demo.py +++ b/examples/lightrag_azure_openai_demo.py @@ -81,34 +81,46 @@ asyncio.run(test_funcs()) embedding_dimension = 3072 -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, - ), -) -rag.initialize_storages() -initialize_pipeline_status() +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, + ), + ) -book1 = open("./book_1.txt", encoding="utf-8") -book2 = open("./book_2.txt", encoding="utf-8") + await rag.initialize_storages() + await initialize_pipeline_status() -rag.insert([book1.read(), book2.read()]) + return rag -query_text = "What are the main themes?" -print("Result (Naive):") -print(rag.query(query_text, param=QueryParam(mode="naive"))) +def main(): + rag = asyncio.run(initialize_rag()) -print("\nResult (Local):") -print(rag.query(query_text, param=QueryParam(mode="local"))) + book1 = open("./book_1.txt", encoding="utf-8") + book2 = open("./book_2.txt", encoding="utf-8") -print("\nResult (Global):") -print(rag.query(query_text, param=QueryParam(mode="global"))) + rag.insert([book1.read(), book2.read()]) -print("\nResult (Hybrid):") -print(rag.query(query_text, param=QueryParam(mode="hybrid"))) + 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() diff --git a/examples/lightrag_bedrock_demo.py b/examples/lightrag_bedrock_demo.py index 68e9f962..c7f41677 100644 --- a/examples/lightrag_bedrock_demo.py +++ b/examples/lightrag_bedrock_demo.py @@ -53,3 +53,7 @@ def main(): "What are the top themes in this story?", param=QueryParam(mode=mode) ) ) + + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_nvidia_demo.py b/examples/lightrag_nvidia_demo.py index 6de0814c..0e9259bc 100644 --- a/examples/lightrag_nvidia_demo.py +++ b/examples/lightrag_nvidia_demo.py @@ -125,7 +125,7 @@ async def initialize_rag(): async def main(): try: # Initialize RAG instance - rag = asyncio.run(initialize_rag()) + rag = await initialize_rag() # reading file with open("./book.txt", "r", encoding="utf-8") as f: diff --git a/examples/lightrag_openai_compatible_demo.py b/examples/lightrag_openai_compatible_demo.py index 1c4a7a92..d26a8de3 100644 --- a/examples/lightrag_openai_compatible_demo.py +++ b/examples/lightrag_openai_compatible_demo.py @@ -77,7 +77,7 @@ async def initialize_rag(): async def main(): try: # Initialize RAG instance - rag = asyncio.run(initialize_rag()) + rag = await initialize_rag() with open("./book.txt", "r", encoding="utf-8") as f: await rag.ainsert(f.read()) diff --git a/examples/lightrag_openai_compatible_demo_embedding_cache.py b/examples/lightrag_openai_compatible_demo_embedding_cache.py index 85408f3b..4638219f 100644 --- a/examples/lightrag_openai_compatible_demo_embedding_cache.py +++ b/examples/lightrag_openai_compatible_demo_embedding_cache.py @@ -81,7 +81,7 @@ async def initialize_rag(): async def main(): try: # Initialize RAG instance - rag = asyncio.run(initialize_rag()) + rag = await initialize_rag() with open("./book.txt", "r", encoding="utf-8") as f: await rag.ainsert(f.read()) diff --git a/examples/lightrag_oracle_demo.py b/examples/lightrag_oracle_demo.py index 420f1af0..6663f6a1 100644 --- a/examples/lightrag_oracle_demo.py +++ b/examples/lightrag_oracle_demo.py @@ -107,7 +107,7 @@ async def initialize_rag(): async def main(): try: # Initialize RAG instance - rag = asyncio.run(initialize_rag()) + rag = await initialize_rag() # Extract and Insert into LightRAG storage with open(WORKING_DIR + "/docs.txt", "r", encoding="utf-8") as f: diff --git a/examples/lightrag_tidb_demo.py b/examples/lightrag_tidb_demo.py index f167e9cc..52695560 100644 --- a/examples/lightrag_tidb_demo.py +++ b/examples/lightrag_tidb_demo.py @@ -87,7 +87,7 @@ async def initialize_rag(): async def main(): try: # Initialize RAG instance - rag = asyncio.run(initialize_rag()) + rag = await initialize_rag() with open("./book.txt", "r", encoding="utf-8") as f: rag.insert(f.read()) diff --git a/examples/lightrag_zhipu_postgres_demo.py b/examples/lightrag_zhipu_postgres_demo.py index 304c5f2c..e4a20f26 100644 --- a/examples/lightrag_zhipu_postgres_demo.py +++ b/examples/lightrag_zhipu_postgres_demo.py @@ -59,7 +59,7 @@ async def initialize_rag(): async def main(): # Initialize RAG instance - rag = asyncio.run(initialize_rag()) + rag = await initialize_rag() # add embedding_func for graph database, it's deleted in commit 5661d76860436f7bf5aef2e50d9ee4a59660146c rag.chunk_entity_relation_graph.embedding_func = rag.embedding_func diff --git a/examples/query_keyword_separation_example.py b/examples/query_keyword_separation_example.py index cbfdd930..092330f4 100644 --- a/examples/query_keyword_separation_example.py +++ b/examples/query_keyword_separation_example.py @@ -102,7 +102,7 @@ async def initialize_rag(): # Example function demonstrating the new query_with_separate_keyword_extraction usage async def run_example(): # Initialize RAG instance - rag = asyncio.run(initialize_rag()) + rag = await initialize_rag() book1 = open("./book_1.txt", encoding="utf-8") book2 = open("./book_2.txt", encoding="utf-8")