Merge pull request #823 from YanSte/cleanup-drop

Cleanup drop

lightrag/kg/json_doc_status_impl.py

@@ -68,6 +68,3 @@ class JsonDocStatusStorage(DocStatusStorage):
         for doc_id in doc_ids:
             self._data.pop(doc_id, None)
         await self.index_done_callback()
-
-    async def drop(self) -> None:
-        raise NotImplementedError

lightrag/kg/json_kv_impl.py

@@ -46,9 +46,6 @@ class JsonKVStorage(BaseKVStorage):
         left_data = {k: v for k, v in data.items() if k not in self._data}
         self._data.update(left_data)
 
-    async def drop(self) -> None:
-        self._data = {}
-
     async def delete(self, ids: list[str]) -> None:
         for doc_id in ids:
             self._data.pop(doc_id, None)

lightrag/kg/mongo_impl.py

@@ -117,10 +117,6 @@ class MongoKVStorage(BaseKVStorage):
         # Mongo handles persistence automatically
         pass
 
-    async def drop(self) -> None:
-        """Drop the collection"""
-        await self._data.drop()
-
 
 @final
 @dataclass
@@ -169,10 +165,6 @@ class MongoDocStatusStorage(DocStatusStorage):
             )
         await asyncio.gather(*update_tasks)
 
-    async def drop(self) -> None:
-        """Drop the collection"""
-        await self._data.drop()
-
     async def get_status_counts(self) -> dict[str, int]:
         """Get counts of documents in each status"""
         pipeline = [{"$group": {"_id": "$status", "count": {"$sum": 1}}}]

lightrag/kg/oracle_impl.py

@@ -320,9 +320,6 @@ class OracleKVStorage(BaseKVStorage):
         # Oracle handles persistence automatically
         pass
 
-    async def drop(self) -> None:
-        raise NotImplementedError
-
 
 @final
 @dataclass

lightrag/kg/postgres_impl.py

@@ -301,9 +301,6 @@ class PGKVStorage(BaseKVStorage):
         # PG handles persistence automatically
         pass
 
-    async def drop(self) -> None:
-        raise NotImplementedError
-
 
 @final
 @dataclass
@@ -534,9 +531,6 @@ class PGDocStatusStorage(DocStatusStorage):
             )
         return data
 
-    async def drop(self) -> None:
-        raise NotImplementedError
-
 
 class PGGraphQueryException(Exception):
     """Exception for the AGE queries."""

lightrag/kg/redis_impl.py

@@ -58,11 +58,6 @@ class RedisKVStorage(BaseKVStorage):
         for k in data:
             data[k]["_id"] = k
 
-    async def drop(self) -> None:
-        keys = await self._redis.keys(f"{self.namespace}:*")
-        if keys:
-            await self._redis.delete(*keys)
-
     async def index_done_callback(self) -> None:
         # Redis handles persistence automatically
         pass

lightrag/kg/tidb_impl.py

@@ -214,9 +214,6 @@ class TiDBKVStorage(BaseKVStorage):
         # Ti handles persistence automatically
         pass
 
-    async def drop(self) -> None:
-        raise NotImplementedError
-
 
 @final
 @dataclass

lightrag/operate.py

@@ -242,16 +242,10 @@ async def _merge_edges_then_upsert(
         # Handle the case where get_edge returns None or missing fields
         if already_edge:
             # Get weight with default 0.0 if missing
-            if "weight" in already_edge:
-                already_weights.append(already_edge["weight"])
-            else:
-                logger.warning(
-                    f"Edge between {src_id} and {tgt_id} missing weight field"
-                )
-                already_weights.append(0.0)
+            already_weights.append(already_edge.get("weight", 0.0))
 
             # Get source_id with empty string default if missing or None
-            if "source_id" in already_edge and already_edge["source_id"] is not None:
+            if already_edge.get("source_id") is not None:
                 already_source_ids.extend(
                     split_string_by_multi_markers(
                         already_edge["source_id"], [GRAPH_FIELD_SEP]
@@ -259,14 +253,11 @@ async def _merge_edges_then_upsert(
                 )
 
             # Get description with empty string default if missing or None
-            if (
-                "description" in already_edge
-                and already_edge["description"] is not None
-            ):
+            if already_edge.get("description") is not None:
                 already_description.append(already_edge["description"])
 
             # Get keywords with empty string default if missing or None
-            if "keywords" in already_edge and already_edge["keywords"] is not None:
+            if already_edge.get("keywords") is not None:
                 already_keywords.extend(
                     split_string_by_multi_markers(
                         already_edge["keywords"], [GRAPH_FIELD_SEP]