Add some script in examples to copy llm cache from one solution to another
This commit is contained in:
@@ -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 asyncio
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
@@ -8,7 +14,7 @@ from lightrag.storage import JsonKVStorage
|
|||||||
|
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
ROOT_DIR = os.environ.get("ROOT_DIR")
|
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)
|
logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
|
||||||
|
|
||||||
@@ -24,12 +30,12 @@ postgres_db = PostgreSQLDB(
|
|||||||
"port": 15432,
|
"port": 15432,
|
||||||
"user": "rag",
|
"user": "rag",
|
||||||
"password": "rag",
|
"password": "rag",
|
||||||
"database": "r1",
|
"database": "r2",
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
async def main():
|
async def copy_from_postgres_to_json():
|
||||||
await postgres_db.initdb()
|
await postgres_db.initdb()
|
||||||
|
|
||||||
from_llm_response_cache = PGKVStorage(
|
from_llm_response_cache = PGKVStorage(
|
||||||
@@ -62,5 +68,30 @@ async def main():
|
|||||||
print("Mission accomplished!")
|
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__":
|
if __name__ == "__main__":
|
||||||
asyncio.run(main())
|
asyncio.run(copy_from_json_to_postgres())
|
Reference in New Issue
Block a user