diff --git a/examples/copy_postgres_llm_cache_to_json.py b/examples/copy_llm_cache_to_another_storage.py similarity index 57% rename from examples/copy_postgres_llm_cache_to_json.py rename to examples/copy_llm_cache_to_another_storage.py index f5fa0d51..b9378c7c 100644 --- a/examples/copy_postgres_llm_cache_to_json.py +++ b/examples/copy_llm_cache_to_another_storage.py @@ -1,3 +1,9 @@ +""" +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 @@ -8,7 +14,7 @@ from lightrag.storage import JsonKVStorage load_dotenv() ROOT_DIR = os.environ.get("ROOT_DIR") -WORKING_DIR = f"{ROOT_DIR}/dickens-pg" +WORKING_DIR = f"{ROOT_DIR}/dickens" logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO) @@ -24,12 +30,12 @@ postgres_db = PostgreSQLDB( "port": 15432, "user": "rag", "password": "rag", - "database": "r1", + "database": "r2", } ) -async def main(): +async def copy_from_postgres_to_json(): await postgres_db.initdb() from_llm_response_cache = PGKVStorage( @@ -62,5 +68,30 @@ async def main(): print("Mission accomplished!") +async def copy_from_json_to_postgres(): + await postgres_db.initdb() + + from_llm_response_cache = JsonKVStorage( + namespace="llm_response_cache", + global_config={"working_dir": WORKING_DIR}, + embedding_func=None, + ) + + to_llm_response_cache = PGKVStorage( + namespace="llm_response_cache", + global_config={"embedding_batch_num": 6}, + embedding_func=None, + db=postgres_db, + ) + + for mode in await from_llm_response_cache.all_keys(): + print(f"Copying {mode}") + caches = await from_llm_response_cache.get_by_id(mode) + for k, v in caches.items(): + item = {mode: {k: v}} + print(f"\tCopying {item}") + await to_llm_response_cache.upsert(item) + + if __name__ == "__main__": - asyncio.run(main()) + asyncio.run(copy_from_json_to_postgres())