From 5a082a0052cc008bb51a30fbc25f72d0aaee683e Mon Sep 17 00:00:00 2001
From: Yannick Stephan <stephan.yannick@me.com>
Date: Sat, 8 Feb 2025 23:20:37 +0100
Subject: [PATCH] cleaned code

---
 lightrag/kg/json_kv_impl.py  | 4 +---
 lightrag/kg/mongo_impl.py    | 6 ++----
 lightrag/kg/oracle_impl.py   | 8 ++------
 lightrag/kg/postgres_impl.py | 4 +---
 lightrag/kg/redis_impl.py    | 6 ++----
 lightrag/kg/tidb_impl.py     | 6 ++----
 6 files changed, 10 insertions(+), 24 deletions(-)

diff --git a/lightrag/kg/json_kv_impl.py b/lightrag/kg/json_kv_impl.py
index 6ee49f7c..927ffe32 100644
--- a/lightrag/kg/json_kv_impl.py
+++ b/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(
-        self, status: str
-    ) -> Union[list[dict[str, Any]], None]:
+    async def get_by_status_and_ids(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
diff --git a/lightrag/kg/mongo_impl.py b/lightrag/kg/mongo_impl.py
index d0598ca4..7cfdb994 100644
--- a/lightrag/kg/mongo_impl.py
+++ b/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(
-        self, status: str
-    ) -> Union[list[dict[str, Any]], None]:
+    async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
         """Get documents by status and ids"""
         return self._data.find({"status": status})
 
diff --git a/lightrag/kg/oracle_impl.py b/lightrag/kg/oracle_impl.py
index 9438c323..f51a5eb8 100644
--- a/lightrag/kg/oracle_impl.py
+++ b/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)
-        if res:
-            return res
-        else:
-            return None
+        return await self.db.query(SQL, params, multirows=True)
 
     async def filter_keys(self, keys: list[str]) -> set[str]:
         """Return keys that don't exist in storage"""
diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py
index dccb2d7f..c6757765 100644
--- a/lightrag/kg/postgres_impl.py
+++ b/lightrag/kg/postgres_impl.py
@@ -238,9 +238,7 @@ class PGKVStorage(BaseKVStorage):
         else:
             return None
         
-    async def get_by_status_and_ids(
-        self, status: str
-    ) -> Union[list[dict[str, Any]], None]:
+    async def get_by_status_and_ids(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}
diff --git a/lightrag/kg/redis_impl.py b/lightrag/kg/redis_impl.py
index 15faa843..b3ff890f 100644
--- a/lightrag/kg/redis_impl.py
+++ b/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(
-        self, status: str,
-    ) -> Union[list[dict[str, Any]], None]:
+    async def get_by_status_and_ids(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)
diff --git a/lightrag/kg/tidb_impl.py b/lightrag/kg/tidb_impl.py
index 97d5794f..81450e87 100644
--- a/lightrag/kg/tidb_impl.py
+++ b/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(
-        self, status: str
-    ) -> Union[list[dict], None]:
+    async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], 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)