diff --git a/lightrag/kg/json_kv_impl.py b/lightrag/kg/json_kv_impl.py index 7e13dea7..0566fb26 100644 --- a/lightrag/kg/json_kv_impl.py +++ b/lightrag/kg/json_kv_impl.py @@ -50,3 +50,6 @@ class JsonKVStorage(BaseKVStorage): for doc_id in ids: self._data.pop(doc_id, None) await self.index_done_callback() + + async def drop(self) -> None: + self._data = {} diff --git a/lightrag/kg/mongo_impl.py b/lightrag/kg/mongo_impl.py index 8cfc84b9..87c9f9f1 100644 --- a/lightrag/kg/mongo_impl.py +++ b/lightrag/kg/mongo_impl.py @@ -117,6 +117,10 @@ class MongoKVStorage(BaseKVStorage): # Mongo handles persistence automatically pass + async def drop(self) -> None: + """Drop the collection""" + await self._data.drop() + @final @dataclass @@ -198,6 +202,10 @@ class MongoDocStatusStorage(DocStatusStorage): # Mongo handles persistence automatically pass + async def drop(self) -> None: + """Drop the collection""" + await self._data.drop() + @final @dataclass diff --git a/lightrag/kg/redis_impl.py b/lightrag/kg/redis_impl.py index 056fbc9e..3b8ebee4 100644 --- a/lightrag/kg/redis_impl.py +++ b/lightrag/kg/redis_impl.py @@ -61,3 +61,8 @@ class RedisKVStorage(BaseKVStorage): async def index_done_callback(self) -> None: # Redis handles persistence automatically pass + + async def drop(self) -> None: + keys = await self._redis.keys(f"{self.namespace}:*") + if keys: + await self._redis.delete(*keys)