Commit
·
cf71c1d
1
Parent(s):
e3f4921
Add some script in examples to copy llm cache from one solution to another
Browse files
examples/{copy_postgres_llm_cache_to_json.py → copy_llm_cache_to_another_storage.py}
RENAMED
|
@@ -1,3 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
import asyncio
|
| 2 |
import logging
|
| 3 |
import os
|
|
@@ -8,7 +14,7 @@ from lightrag.storage import JsonKVStorage
|
|
| 8 |
|
| 9 |
load_dotenv()
|
| 10 |
ROOT_DIR = os.environ.get("ROOT_DIR")
|
| 11 |
-
WORKING_DIR = f"{ROOT_DIR}/dickens
|
| 12 |
|
| 13 |
logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
|
| 14 |
|
|
@@ -24,12 +30,12 @@ postgres_db = PostgreSQLDB(
|
|
| 24 |
"port": 15432,
|
| 25 |
"user": "rag",
|
| 26 |
"password": "rag",
|
| 27 |
-
"database": "
|
| 28 |
}
|
| 29 |
)
|
| 30 |
|
| 31 |
|
| 32 |
-
async def
|
| 33 |
await postgres_db.initdb()
|
| 34 |
|
| 35 |
from_llm_response_cache = PGKVStorage(
|
|
@@ -62,5 +68,30 @@ async def main():
|
|
| 62 |
print("Mission accomplished!")
|
| 63 |
|
| 64 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 65 |
if __name__ == "__main__":
|
| 66 |
-
asyncio.run(
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Sometimes you need to switch a storage solution, but you want to save LLM token and time.
|
| 3 |
+
This handy script helps you to copy the LLM caches from one storage solution to another.
|
| 4 |
+
(Not all the storage impl are supported)
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
import asyncio
|
| 8 |
import logging
|
| 9 |
import os
|
|
|
|
| 14 |
|
| 15 |
load_dotenv()
|
| 16 |
ROOT_DIR = os.environ.get("ROOT_DIR")
|
| 17 |
+
WORKING_DIR = f"{ROOT_DIR}/dickens"
|
| 18 |
|
| 19 |
logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
|
| 20 |
|
|
|
|
| 30 |
"port": 15432,
|
| 31 |
"user": "rag",
|
| 32 |
"password": "rag",
|
| 33 |
+
"database": "r2",
|
| 34 |
}
|
| 35 |
)
|
| 36 |
|
| 37 |
|
| 38 |
+
async def copy_from_postgres_to_json():
|
| 39 |
await postgres_db.initdb()
|
| 40 |
|
| 41 |
from_llm_response_cache = PGKVStorage(
|
|
|
|
| 68 |
print("Mission accomplished!")
|
| 69 |
|
| 70 |
|
| 71 |
+
async def copy_from_json_to_postgres():
|
| 72 |
+
await postgres_db.initdb()
|
| 73 |
+
|
| 74 |
+
from_llm_response_cache = JsonKVStorage(
|
| 75 |
+
namespace="llm_response_cache",
|
| 76 |
+
global_config={"working_dir": WORKING_DIR},
|
| 77 |
+
embedding_func=None,
|
| 78 |
+
)
|
| 79 |
+
|
| 80 |
+
to_llm_response_cache = PGKVStorage(
|
| 81 |
+
namespace="llm_response_cache",
|
| 82 |
+
global_config={"embedding_batch_num": 6},
|
| 83 |
+
embedding_func=None,
|
| 84 |
+
db=postgres_db,
|
| 85 |
+
)
|
| 86 |
+
|
| 87 |
+
for mode in await from_llm_response_cache.all_keys():
|
| 88 |
+
print(f"Copying {mode}")
|
| 89 |
+
caches = await from_llm_response_cache.get_by_id(mode)
|
| 90 |
+
for k, v in caches.items():
|
| 91 |
+
item = {mode: {k: v}}
|
| 92 |
+
print(f"\tCopying {item}")
|
| 93 |
+
await to_llm_response_cache.upsert(item)
|
| 94 |
+
|
| 95 |
+
|
| 96 |
if __name__ == "__main__":
|
| 97 |
+
asyncio.run(copy_from_json_to_postgres())
|