cleaned import

This commit is contained in:
Yannick Stephan
2025-02-09 11:24:08 +01:00
parent 61fd3e6127
commit 4cce14e65e
8 changed files with 62 additions and 70 deletions

View File

@@ -51,8 +51,6 @@ class JsonKVStorage(BaseKVStorage):
async def drop(self) -> None:
self._data = {}
async def get_by_status(
self, status: str
) -> Union[list[dict[str, Any]], None]:
async def get_by_status(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

View File

@@ -77,9 +77,7 @@ class MongoKVStorage(BaseKVStorage):
"""Drop the collection"""
await self._data.drop()
async def get_by_status(
self, status: str
) -> Union[list[dict[str, Any]], None]:
async def get_by_status(self, status: str) -> Union[list[dict[str, Any]], None]:
"""Get documents by status and ids"""
return self._data.find({"status": status})

View File

@@ -229,9 +229,7 @@ class OracleKVStorage(BaseKVStorage):
res = [{k: v} for k, v in dict_res.items()]
return res
async def get_by_status(
self, status: str
) -> Union[list[dict[str, Any]], None]:
async def get_by_status(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}

View File

@@ -231,9 +231,7 @@ class PGKVStorage(BaseKVStorage):
else:
return await self.db.query(sql, params, multirows=True)
async def get_by_status(
self, status: str
) -> Union[list[dict[str, Any]], None]:
async def get_by_status(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}

View File

@@ -59,9 +59,7 @@ class RedisKVStorage(BaseKVStorage):
if keys:
await self._redis.delete(*keys)
async def get_by_status(
self, status: str
) -> Union[list[dict[str, Any]], None]:
async def get_by_status(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)

View File

@@ -322,9 +322,7 @@ class TiDBVectorDBStorage(BaseVectorStorage):
merge_sql = SQL_TEMPLATES["insert_relationship"]
await self.db.execute(merge_sql, data)
async def get_by_status(
self, status: str
) -> Union[list[dict[str, Any]], None]:
async def get_by_status(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)