gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added update flag to avoid persistence if no data is changed for KV storage

Browse Source
This commit is contained in:
yangdx
2025-03-10 01:17:25 +08:00
parent 4977c718f1
commit d2708b966d
3 changed files with 58 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
lightrag/kg/json_kv_impl.py
View File

@@ -14,6 +14,9 @@ from .shared_storage import (
get_namespace_data,
get_storage_lock,
get_data_init_lock,
get_update_flag,
set_all_update_flags,
clear_all_update_flags,
try_initialize_namespace,
)
@@ -25,10 +28,13 @@ class JsonKVStorage(BaseKVStorage):
working_dir = self.global_config["working_dir"]
self._file_name = os.path.join(working_dir, f"kv_store_{self.namespace}.json")
self._data = None
self._storage_lock = None
self.storage_updated = None
async def initialize(self):
"""Initialize storage data"""
self._storage_lock = get_storage_lock()
self.storage_updated = await get_update_flag(self.namespace)
async with get_data_init_lock():
# check need_init must before get_namespace_data
need_init = await try_initialize_namespace(self.namespace)
@@ -51,21 +57,24 @@ class JsonKVStorage(BaseKVStorage):
async def index_done_callback(self) -> None:
async with self._storage_lock:
data_dict = (
dict(self._data) if hasattr(self._data, "_getvalue") else self._data
)
# Calculate data count based on namespace
if self.namespace.endswith("cache"):
# # For cache namespaces, sum the cache entries across all cache types
data_count = sum(len(first_level_dict) for first_level_dict in data_dict.values()
if isinstance(first_level_dict, dict))
else:
# For non-cache namespaces, use the original count method
data_count = len(data_dict)
logger.info(f"Process {os.getpid()} KV writting {data_count} records to {self.namespace}")
write_json(data_dict, self._file_name)
if self.storage_updated:
data_dict = (
dict(self._data) if hasattr(self._data, "_getvalue") else self._data
)
# Calculate data count based on namespace
if self.namespace.endswith("cache"):
# # For cache namespaces, sum the cache entries across all cache types
data_count = sum(len(first_level_dict) for first_level_dict in data_dict.values()
if isinstance(first_level_dict, dict))
else:
# For non-cache namespaces, use the original count method
data_count = len(data_dict)
logger.info(f"Process {os.getpid()} KV writting {data_count} records to {self.namespace}")
write_json(data_dict, self._file_name)
await clear_all_update_flags(self.namespace)
async def get_all(self) -> dict[str, Any]:
"""Get all data from storage
@@ -101,9 +110,11 @@ class JsonKVStorage(BaseKVStorage):
logger.info(f"Inserting {len(data)} to {self.namespace}")
async with self._storage_lock:
self._data.update(data)
await set_all_update_flags(self.namespace)
async def delete(self, ids: list[str]) -> None:
async with self._storage_lock:
for doc_id in ids:
self._data.pop(doc_id, None)
await set_all_update_flags(self.namespace)
await self.index_done_callback()