mirror of
https://github.com/HKUDS/LightRAG.git
synced 2025-08-01 21:31:54 +00:00
271722405f
Refactored the LLM cache to a flat Key-Value (KV) structure, replacing the previous nested format. The old structure used the 'mode' as a key and stored specific cache content as JSON nested under it. This change significantly enhances cache recall efficiency.
115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
"""
|
|
Sometimes you need to switch a storage solution, but you want to save LLM token and time.
|
|
This handy script helps you to copy the LLM caches from one storage solution to another.
|
|
(Not all the storage impl are supported)
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
from dotenv import load_dotenv
|
|
|
|
from lightrag.kg.postgres_impl import PostgreSQLDB, PGKVStorage
|
|
from lightrag.kg.json_kv_impl import JsonKVStorage
|
|
from lightrag.namespace import NameSpace
|
|
|
|
load_dotenv()
|
|
ROOT_DIR = os.environ.get("ROOT_DIR")
|
|
WORKING_DIR = f"{ROOT_DIR}/dickens"
|
|
|
|
logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
|
|
|
|
if not os.path.exists(WORKING_DIR):
|
|
os.mkdir(WORKING_DIR)
|
|
|
|
# AGE
|
|
os.environ["AGE_GRAPH_NAME"] = "chinese"
|
|
|
|
postgres_db = PostgreSQLDB(
|
|
config={
|
|
"host": "localhost",
|
|
"port": 15432,
|
|
"user": "rag",
|
|
"password": "rag",
|
|
"database": "r2",
|
|
}
|
|
)
|
|
|
|
|
|
async def copy_from_postgres_to_json():
|
|
await postgres_db.initdb()
|
|
|
|
from_llm_response_cache = PGKVStorage(
|
|
namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
|
|
global_config={"embedding_batch_num": 6},
|
|
embedding_func=None,
|
|
db=postgres_db,
|
|
)
|
|
|
|
to_llm_response_cache = JsonKVStorage(
|
|
namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
|
|
global_config={"working_dir": WORKING_DIR},
|
|
embedding_func=None,
|
|
)
|
|
|
|
# Get all cache data using the new flattened structure
|
|
all_data = await from_llm_response_cache.get_all()
|
|
|
|
# Convert flattened data to hierarchical structure for JsonKVStorage
|
|
kv = {}
|
|
for flattened_key, cache_entry in all_data.items():
|
|
# Parse flattened key: {mode}:{cache_type}:{hash}
|
|
parts = flattened_key.split(":", 2)
|
|
if len(parts) == 3:
|
|
mode, cache_type, hash_value = parts
|
|
if mode not in kv:
|
|
kv[mode] = {}
|
|
kv[mode][hash_value] = cache_entry
|
|
print(f"Copying {flattened_key} -> {mode}[{hash_value}]")
|
|
else:
|
|
print(f"Skipping invalid key format: {flattened_key}")
|
|
|
|
await to_llm_response_cache.upsert(kv)
|
|
await to_llm_response_cache.index_done_callback()
|
|
print("Mission accomplished!")
|
|
|
|
|
|
async def copy_from_json_to_postgres():
|
|
await postgres_db.initdb()
|
|
|
|
from_llm_response_cache = JsonKVStorage(
|
|
namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
|
|
global_config={"working_dir": WORKING_DIR},
|
|
embedding_func=None,
|
|
)
|
|
|
|
to_llm_response_cache = PGKVStorage(
|
|
namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
|
|
global_config={"embedding_batch_num": 6},
|
|
embedding_func=None,
|
|
db=postgres_db,
|
|
)
|
|
|
|
# Get all cache data from JsonKVStorage (hierarchical structure)
|
|
all_data = await from_llm_response_cache.get_all()
|
|
|
|
# Convert hierarchical data to flattened structure for PGKVStorage
|
|
flattened_data = {}
|
|
for mode, mode_data in all_data.items():
|
|
print(f"Processing mode: {mode}")
|
|
for hash_value, cache_entry in mode_data.items():
|
|
# Determine cache_type from cache entry or use default
|
|
cache_type = cache_entry.get("cache_type", "extract")
|
|
# Create flattened key: {mode}:{cache_type}:{hash}
|
|
flattened_key = f"{mode}:{cache_type}:{hash_value}"
|
|
flattened_data[flattened_key] = cache_entry
|
|
print(f"\tConverting {mode}[{hash_value}] -> {flattened_key}")
|
|
|
|
# Upsert the flattened data
|
|
await to_llm_response_cache.upsert(flattened_data)
|
|
print("Mission accomplished!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(copy_from_json_to_postgres())
|