This commit is contained in:
zrguo
2025-03-04 12:25:07 +08:00
parent acb074d0f8
commit 6c8fa95214
10 changed files with 58 additions and 39 deletions

View File

@@ -655,16 +655,19 @@ setup_logger("lightrag", level="INFO")
# Note: Default settings use NetworkX # Note: Default settings use NetworkX
# Initialize LightRAG with Neo4J implementation. # Initialize LightRAG with Neo4J implementation.
rag = LightRAG( async def initialize_rag():
working_dir=WORKING_DIR, rag = LightRAG(
llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model working_dir=WORKING_DIR,
graph_storage="Neo4JStorage", #<-----------override KG default llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model
) graph_storage="Neo4JStorage", #<-----------override KG default
)
# Initialize database connections # Initialize database connections
await rag.initialize_storages() await rag.initialize_storages()
# Initialize pipeline status for document processing # Initialize pipeline status for document processing
await initialize_pipeline_status() await initialize_pipeline_status()
return rag
``` ```
see test_neo4j.py for a working example. see test_neo4j.py for a working example.

View File

@@ -81,34 +81,46 @@ asyncio.run(test_funcs())
embedding_dimension = 3072 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() async def initialize_rag():
initialize_pipeline_status() 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") await rag.initialize_storages()
book2 = open("./book_2.txt", encoding="utf-8") await initialize_pipeline_status()
rag.insert([book1.read(), book2.read()]) return rag
query_text = "What are the main themes?"
print("Result (Naive):") def main():
print(rag.query(query_text, param=QueryParam(mode="naive"))) rag = asyncio.run(initialize_rag())
print("\nResult (Local):") book1 = open("./book_1.txt", encoding="utf-8")
print(rag.query(query_text, param=QueryParam(mode="local"))) book2 = open("./book_2.txt", encoding="utf-8")
print("\nResult (Global):") rag.insert([book1.read(), book2.read()])
print(rag.query(query_text, param=QueryParam(mode="global")))
print("\nResult (Hybrid):") query_text = "What are the main themes?"
print(rag.query(query_text, param=QueryParam(mode="hybrid")))
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()

View File

@@ -53,3 +53,7 @@ def main():
"What are the top themes in this story?", param=QueryParam(mode=mode) "What are the top themes in this story?", param=QueryParam(mode=mode)
) )
) )
if __name__ == "__main__":
main()

View File

@@ -125,7 +125,7 @@ async def initialize_rag():
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = await initialize_rag()
# reading file # reading file
with open("./book.txt", "r", encoding="utf-8") as f: with open("./book.txt", "r", encoding="utf-8") as f:

View File

@@ -77,7 +77,7 @@ async def initialize_rag():
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = await initialize_rag()
with open("./book.txt", "r", encoding="utf-8") as f: with open("./book.txt", "r", encoding="utf-8") as f:
await rag.ainsert(f.read()) await rag.ainsert(f.read())

View File

@@ -81,7 +81,7 @@ async def initialize_rag():
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = await initialize_rag()
with open("./book.txt", "r", encoding="utf-8") as f: with open("./book.txt", "r", encoding="utf-8") as f:
await rag.ainsert(f.read()) await rag.ainsert(f.read())

View File

@@ -107,7 +107,7 @@ async def initialize_rag():
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = await initialize_rag()
# Extract and Insert into LightRAG storage # Extract and Insert into LightRAG storage
with open(WORKING_DIR + "/docs.txt", "r", encoding="utf-8") as f: with open(WORKING_DIR + "/docs.txt", "r", encoding="utf-8") as f:

View File

@@ -87,7 +87,7 @@ async def initialize_rag():
async def main(): async def main():
try: try:
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = await initialize_rag()
with open("./book.txt", "r", encoding="utf-8") as f: with open("./book.txt", "r", encoding="utf-8") as f:
rag.insert(f.read()) rag.insert(f.read())

View File

@@ -59,7 +59,7 @@ async def initialize_rag():
async def main(): async def main():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = await 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
rag.chunk_entity_relation_graph.embedding_func = rag.embedding_func rag.chunk_entity_relation_graph.embedding_func = rag.embedding_func

View File

@@ -102,7 +102,7 @@ async def initialize_rag():
# Example function demonstrating the new query_with_separate_keyword_extraction usage # Example function demonstrating the new query_with_separate_keyword_extraction usage
async def run_example(): async def run_example():
# Initialize RAG instance # Initialize RAG instance
rag = asyncio.run(initialize_rag()) rag = await 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")