diff --git a/lightrag/kg/json_kv_impl.py b/lightrag/kg/json_kv_impl.py index 14565c86..ff184dbd 100644 --- a/lightrag/kg/json_kv_impl.py +++ b/lightrag/kg/json_kv_impl.py @@ -21,7 +21,7 @@ class JsonKVStorage(BaseKVStorage): self._data: dict[str, Any] = load_json(self._file_name) or {} self._lock = asyncio.Lock() logger.info(f"Load KV {self.namespace} with {len(self._data)} data") - + async def index_done_callback(self): write_json(self._data, self._file_name) @@ -46,4 +46,4 @@ class JsonKVStorage(BaseKVStorage): self._data.update(left_data) async def drop(self) -> None: - self._data = {} \ No newline at end of file + self._data = {} diff --git a/lightrag/kg/jsondocstatus_impl.py b/lightrag/kg/jsondocstatus_impl.py index 603487bc..31aa836a 100644 --- a/lightrag/kg/jsondocstatus_impl.py +++ b/lightrag/kg/jsondocstatus_impl.py @@ -50,7 +50,7 @@ Usage: import os from dataclasses import dataclass -from typing import Any, Union, Dict +from typing import Any, Union from lightrag.utils import ( logger, @@ -85,18 +85,18 @@ class JsonDocStatusStorage(DocStatusStorage): ] ) - async def get_status_counts(self) -> Dict[str, int]: + async def get_status_counts(self) -> dict[str, int]: """Get counts of documents in each status""" counts = {status: 0 for status in DocStatus} for doc in self._data.values(): counts[doc["status"]] += 1 return counts - async def get_failed_docs(self) -> Dict[str, DocProcessingStatus]: + async def get_failed_docs(self) -> dict[str, DocProcessingStatus]: """Get all failed documents""" return {k: v for k, v in self._data.items() if v["status"] == DocStatus.FAILED} - async def get_pending_docs(self) -> Dict[str, DocProcessingStatus]: + async def get_pending_docs(self) -> dict[str, DocProcessingStatus]: """Get all pending documents""" return {k: v for k, v in self._data.items() if v["status"] == DocStatus.PENDING} diff --git a/lightrag/kg/mongo_impl.py b/lightrag/kg/mongo_impl.py index 45d4bb07..35902d37 100644 --- a/lightrag/kg/mongo_impl.py +++ b/lightrag/kg/mongo_impl.py @@ -74,6 +74,7 @@ class MongoKVStorage(BaseKVStorage): """Drop the collection""" await self._data.drop() + @dataclass class MongoGraphStorage(BaseGraphStorage): """ diff --git a/lightrag/kg/redis_impl.py b/lightrag/kg/redis_impl.py index 025f293f..05da41b7 100644 --- a/lightrag/kg/redis_impl.py +++ b/lightrag/kg/redis_impl.py @@ -20,7 +20,7 @@ class RedisKVStorage(BaseKVStorage): redis_url = os.environ.get("REDIS_URI", "redis://localhost:6379") self._redis = Redis.from_url(redis_url, decode_responses=True) logger.info(f"Use Redis as KV {self.namespace}") - + async def get_by_id(self, id): data = await self._redis.get(f"{self.namespace}:{id}") return json.loads(data) if data else None @@ -53,4 +53,4 @@ class RedisKVStorage(BaseKVStorage): async def drop(self) -> None: keys = await self._redis.keys(f"{self.namespace}:*") if keys: - await self._redis.delete(*keys) \ No newline at end of file + await self._redis.delete(*keys)