From 6fc54d3625a9272521e8565ea691cd8084b5343d Mon Sep 17 00:00:00 2001
From: yangdx <gzdaniel@me.com>
Date: Sat, 8 Nov 2025 18:33:13 +0800
Subject: [PATCH] Move LLM cache migration tool to lightrag.tools module

- Relocated tool to proper package structure
- Updated import paths and documentation
- Added shared storage initialization
- Fixed module path resolution
- Updated usage instructions
---
 .../tools}/README_MIGRATE_LLM_CACHE.md        |  8 +++++---
 .../tools}/migrate_llm_cache.py               | 19 ++++++++++++++++---
 2 files changed, 21 insertions(+), 6 deletions(-)
 rename {tools => lightrag/tools}/README_MIGRATE_LLM_CACHE.md (98%)
 rename {tools => lightrag/tools}/migrate_llm_cache.py (97%)

diff --git a/tools/README_MIGRATE_LLM_CACHE.md b/lightrag/tools/README_MIGRATE_LLM_CACHE.md
similarity index 98%
rename from tools/README_MIGRATE_LLM_CACHE.md
rename to lightrag/tools/README_MIGRATE_LLM_CACHE.md
index 364083d7..3fef0303 100644
--- a/tools/README_MIGRATE_LLM_CACHE.md
+++ b/lightrag/tools/README_MIGRATE_LLM_CACHE.md
@@ -78,7 +78,9 @@ pip install -r requirements.txt
 Run from the LightRAG project root directory:
 
 ```bash
-python tools/migrate_llm_cache.py
+python -m lightrag.tools.migrate_llm_cache
+# or
+python lightrag/tools/migrate_llm_cache.py
 ```
 
 ### Interactive Workflow
@@ -328,7 +330,7 @@ MONGO_URI=mongodb://user:pass@prod-server:27017/
 MONGO_DATABASE=LightRAG
 
 # 2. Run tool
-python tools/migrate_llm_cache.py
+python -m lightrag.tools.migrate_llm_cache
 
 # 3. Select: 1 (JsonKVStorage) -> 4 (MongoKVStorage)
 ```
@@ -356,7 +358,7 @@ POSTGRES_HOST=new-postgres-server
 # ... Other PostgreSQL configs
 
 # 2. Run tool
-python tools/migrate_llm_cache.py
+python -m lightrag.tools.migrate_llm_cache
 
 # 3. Select: 2 (RedisKVStorage) -> 3 (PGKVStorage)
 ```
diff --git a/tools/migrate_llm_cache.py b/lightrag/tools/migrate_llm_cache.py
similarity index 97%
rename from tools/migrate_llm_cache.py
rename to lightrag/tools/migrate_llm_cache.py
index 26b7c81c..68866264 100644
--- a/tools/migrate_llm_cache.py
+++ b/lightrag/tools/migrate_llm_cache.py
@@ -6,7 +6,9 @@ This tool migrates LLM response cache (default:extract:* and default:summary:*)
 between different KV storage implementations while preserving workspace isolation.
 
 Usage:
-    python tools/migrate_llm_cache.py
+    python -m lightrag.tools.migrate_llm_cache
+    # or
+    python lightrag/tools/migrate_llm_cache.py
 
 Supported KV Storage Types:
     - JsonKVStorage
@@ -23,8 +25,8 @@ from typing import Any, Dict, List
 from dataclasses import dataclass, field
 from dotenv import load_dotenv
 
-# Add parent directory to path for imports
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+# Add project root to path for imports
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 
 from lightrag.kg import STORAGE_ENV_REQUIREMENTS
 from lightrag.namespace import NameSpace
@@ -657,6 +659,10 @@ class MigrationTool:
     async def run(self):
         """Run the migration tool"""
         try:
+            # Initialize shared storage (REQUIRED for storage classes to work)
+            from lightrag.kg.shared_storage import initialize_share_data
+            initialize_share_data(workers=1)
+            
             # Print header
             self.print_header()
             self.print_storage_types()
@@ -751,6 +757,13 @@ class MigrationTool:
                     await self.target_storage.finalize()
                 except Exception:
                     pass
+            
+            # Finalize shared storage
+            try:
+                from lightrag.kg.shared_storage import finalize_share_data
+                finalize_share_data()
+            except Exception:
+                pass
 
 
 async def main():