cleaned code

lightrag/kg/json_kv_impl.py

@@ -50,8 +50,6 @@ class JsonKVStorage(BaseKVStorage):
     async def drop(self) -> None:
         self._data = {}
 
-    async def get_by_status_and_ids(
+    async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
-        self, status: str
-    ) -> Union[list[dict[str, Any]], None]:
         result = [v for _, v in self._data.items() if v["status"] == status]
         return result if result else None

lightrag/kg/mongo_impl.py

@@ -12,7 +12,7 @@ if not pm.is_installed("motor"):
 
 from pymongo import MongoClient
 from motor.motor_asyncio import AsyncIOMotorClient
-from typing import Any, TypeVar, Union, List, Tuple
+from typing import Any, Union, List, Tuple
 
 from ..utils import logger
 from ..base import BaseKVStorage, BaseGraphStorage
@@ -77,9 +77,7 @@ class MongoKVStorage(BaseKVStorage):
         """Drop the collection"""
         await self._data.drop()
 
-    async def get_by_status_and_ids(
+    async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
-        self, status: str
-    ) -> Union[list[dict[str, Any]], None]:
         """Get documents by status and ids"""
         return self._data.find({"status": status})
 

lightrag/kg/oracle_impl.py

@@ -4,7 +4,7 @@ import asyncio
 # import html
 # import os
 from dataclasses import dataclass
-from typing import Any, TypeVar, Union
+from typing import Any, Union
 import numpy as np
 import array
 import pipmaster as pm
@@ -243,11 +243,7 @@ class OracleKVStorage(BaseKVStorage):
         """Specifically for llm_response_cache."""
         SQL = SQL_TEMPLATES["get_by_status_" + self.namespace]
         params = {"workspace": self.db.workspace, "status": status}
-        res = await self.db.query(SQL, params, multirows=True)
+        return await self.db.query(SQL, params, multirows=True)
-        if res:
-            return res
-        else:
-            return None
 
     async def filter_keys(self, keys: list[str]) -> set[str]:
         """Return keys that don't exist in storage"""

lightrag/kg/postgres_impl.py

@@ -238,9 +238,7 @@ class PGKVStorage(BaseKVStorage):
         else:
             return None
         
-    async def get_by_status_and_ids(
+    async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
-        self, status: str
-    ) -> Union[list[dict[str, Any]], None]:
         """Specifically for llm_response_cache."""
         SQL = SQL_TEMPLATES["get_by_status_" + self.namespace]
         params = {"workspace": self.db.workspace, "status": status}

lightrag/kg/redis_impl.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, TypeVar, Union
+from typing import Any, Union
 from tqdm.asyncio import tqdm as tqdm_async
 from dataclasses import dataclass
 import pipmaster as pm
@@ -59,9 +59,7 @@ class RedisKVStorage(BaseKVStorage):
         if keys:
             await self._redis.delete(*keys)
             
-    async def get_by_status_and_ids(
+    async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
-        self, status: str,
-    ) -> Union[list[dict[str, Any]], None]:
         pipe = self._redis.pipeline()
         for key in await self._redis.keys(f"{self.namespace}:*"):
             pipe.hgetall(key)

lightrag/kg/tidb_impl.py

@@ -1,7 +1,7 @@
 import asyncio
 import os
 from dataclasses import dataclass
-from typing import Any, TypeVar, Union
+from typing import Any, Union
 
 import numpy as np
 import pipmaster as pm
@@ -333,9 +333,7 @@ class TiDBVectorDBStorage(BaseVectorStorage):
                 merge_sql = SQL_TEMPLATES["insert_relationship"]
                 await self.db.execute(merge_sql, data)
 
-    async def get_by_status_and_ids(
+    async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
-        self, status: str
-    ) -> Union[list[dict], None]:
         SQL = SQL_TEMPLATES["get_by_status_" + self.namespace]
         params = {"workspace": self.db.workspace, "status": status}
         return await self.db.query(SQL, params, multirows=True)