From affa992286a1183525b53deeb28726e0fc6e91ff Mon Sep 17 00:00:00 2001
From: HuangHai <10402852@qq.com>
Date: Fri, 11 Jul 2025 16:15:26 +0800
Subject: [PATCH] 'commit'

---
 dsLightRag/Test/T1_WriteToPg.py  | 10 +++++++---
 dsLightRag/Test/T2_ReadFromPg.py |  4 ++--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/dsLightRag/Test/T1_WriteToPg.py b/dsLightRag/Test/T1_WriteToPg.py
index 0a33cf2d..b8c32e29 100644
--- a/dsLightRag/Test/T1_WriteToPg.py
+++ b/dsLightRag/Test/T1_WriteToPg.py
@@ -63,9 +63,13 @@ async def initialize_pg_rag():
 
 
 async def main():
-    rag = await initialize_pg_rag()
-    with open(f"../Txt/sushi.txt", "r", encoding="utf-8") as f:
-        await rag.ainsert(f.read())
+    try:
+        rag = await initialize_pg_rag()
+        with open(f"../Txt/sushi.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+    finally:
+        if rag:
+            await rag.finalize_storages()
 
 if __name__ == "__main__":
     asyncio.run(main())
diff --git a/dsLightRag/Test/T2_ReadFromPg.py b/dsLightRag/Test/T2_ReadFromPg.py
index 377d0b14..0ce130d6 100644
--- a/dsLightRag/Test/T2_ReadFromPg.py
+++ b/dsLightRag/Test/T2_ReadFromPg.py
@@ -3,12 +3,12 @@ import inspect
 import logging
 import os
 
+from lightrag import QueryParam, LightRAG
 from lightrag.kg.shared_storage import initialize_pipeline_status
 from lightrag.utils import EmbeddingFunc
 
 from Config.Config import LLM_MODEL_NAME, EMBED_DIM, EMBED_MAX_TOKEN_SIZE
-from Util.LightRagUtil import configure_logging, initialize_rag, print_stream, llm_model_func, embedding_func
-from lightrag import QueryParam, LightRAG
+from Util.LightRagUtil import configure_logging, print_stream, llm_model_func, embedding_func
 
 # 在程序开始时添加以下配置
 logging.basicConfig(