diff --git a/dsLightRag/Test/testPg2.py b/dsLightRag/Test/T1_WriteToPg.py similarity index 70% rename from dsLightRag/Test/testPg2.py rename to dsLightRag/Test/T1_WriteToPg.py index 298cd0f6..167db27e 100644 --- a/dsLightRag/Test/testPg2.py +++ b/dsLightRag/Test/T1_WriteToPg.py @@ -67,26 +67,10 @@ async def initialize_rag(): async def main(): - try: - rag = await initialize_rag() - - #with open(f"{ROOT_DIR}/book.txt", "r", encoding="utf-8") as f: - # await rag.ainsert(f.read()) - - print("==== Trying to test the rag queries ====") - print("**** Start Naive Query ****") - start_time = time.time() - # Perform naive search - await rag.aquery( - "What are the top themes in this story?", param=QueryParam(mode="naive") - ) - print(f"Naive Query Time: {time.time() - start_time} seconds") - except Exception as e: - #logger.error(f"Main execution error: {e}") - pass - finally: - # 确保所有资源正确释放 - await asyncio.sleep(0.1) # 给pending任务完成的时间 + rag = await initialize_rag() + + with open(f"{ROOT_DIR}/book.txt", "r", encoding="utf-8") as f: + await rag.ainsert(f.read()) if __name__ == "__main__": diff --git a/dsLightRag/Test/testPg3.py b/dsLightRag/Test/T2_ReadFromPg.py similarity index 100% rename from dsLightRag/Test/testPg3.py rename to dsLightRag/Test/T2_ReadFromPg.py diff --git a/dsLightRag/Test/testPg1.py b/dsLightRag/Test/testPg1.py deleted file mode 100644 index 244b6d92..00000000 --- a/dsLightRag/Test/testPg1.py +++ /dev/null @@ -1,51 +0,0 @@ -import os -import asyncio -from lightrag.kg.postgres_impl import PGGraphStorage -from lightrag.utils import EmbeddingFunc - -from Config.Config import EMBED_DIM, EMBED_MAX_TOKEN_SIZE -from Util.LightRagUtil import embedding_func - -######### -# 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 = "./local_neo4jWorkDir" - -if not os.path.exists(WORKING_DIR): - os.mkdir(WORKING_DIR) - -# AGE -os.environ["AGE_GRAPH_NAME"] = "dickens" - -os.environ["POSTGRES_HOST"] = "10.10.14.208" -os.environ["POSTGRES_PORT"] = "5432" -os.environ["POSTGRES_USER"] = "postgres" -os.environ["POSTGRES_PASSWORD"] = "postgres" -os.environ["POSTGRES_DATABASE"] = "rag" - - -async def main(): - graph_db = PGGraphStorage( - namespace="dickens", - embedding_func=EmbeddingFunc( - embedding_dim=EMBED_DIM, - max_token_size=EMBED_MAX_TOKEN_SIZE, - func=embedding_func - ), - global_config={}, - ) - await graph_db.initialize() - labels = await graph_db.get_all_labels() - print("all labels", labels) - - res = await graph_db.get_knowledge_graph("FEZZIWIG") - print("knowledge graphs", res) - - await graph_db.finalize() - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file