samuel-z-chen commited on
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-pg"
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": "r1",
28
  }
29
  )
30
 
31
 
32
- async def main():
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(main())
 
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())