diff --git a/lightrag/kg/mongo_impl.py b/lightrag/kg/mongo_impl.py index c7b16a70..a6e6edfd 100644 --- a/lightrag/kg/mongo_impl.py +++ b/lightrag/kg/mongo_impl.py @@ -797,8 +797,8 @@ class MongoGraphStorage(BaseGraphStorage): @final @dataclass class MongoVectorDBStorage(BaseVectorStorage): - db: AsyncIOMotorDatabase = field(default=None) - _data: AsyncIOMotorCollection = field(default=None) + db: AsyncIOMotorDatabase | None = field(default=None) + _data: AsyncIOMotorCollection | None = field(default=None) def __post_init__(self): kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {}) diff --git a/lightrag/kg/oracle_impl.py b/lightrag/kg/oracle_impl.py index e6a8e8f4..0916f6b0 100644 --- a/lightrag/kg/oracle_impl.py +++ b/lightrag/kg/oracle_impl.py @@ -398,7 +398,7 @@ class OracleKVStorage(BaseKVStorage): @final @dataclass class OracleVectorDBStorage(BaseVectorStorage): - db: OracleDB = field(default=None) + db: OracleDB | None = field(default=None) def __post_init__(self): config = self.global_config.get("vector_db_storage_cls_kwargs", {}) diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py index e10c3832..3c88a05b 100644 --- a/lightrag/kg/postgres_impl.py +++ b/lightrag/kg/postgres_impl.py @@ -1,5 +1,4 @@ import asyncio -import inspect import json import os import time @@ -373,7 +372,7 @@ class PGKVStorage(BaseKVStorage): @final @dataclass class PGVectorStorage(BaseVectorStorage): - db: PostgreSQLDB = field(default=None) + db: PostgreSQLDB | None = field(default=None) def __post_init__(self): self._max_batch_size = self.global_config["embedding_batch_num"] @@ -394,10 +393,10 @@ class PGVectorStorage(BaseVectorStorage): await ClientManager.release_client(self.db) self.db = None - def _upsert_chunks(self, item: dict): + def _upsert_chunks(self, item: dict[str, Any]) -> tuple[str, dict[str, Any]]: try: upsert_sql = SQL_TEMPLATES["upsert_chunk"] - data = { + data: dict[str, Any] = { "workspace": self.db.workspace, "id": item["__id__"], "tokens": item["tokens"], @@ -412,9 +411,9 @@ class PGVectorStorage(BaseVectorStorage): return upsert_sql, data - def _upsert_entities(self, item: dict): + def _upsert_entities(self, item: dict[str, Any]) -> tuple[str, dict[str, Any]]: upsert_sql = SQL_TEMPLATES["upsert_entity"] - data = { + data: dict[str, Any] = { "workspace": self.db.workspace, "id": item["__id__"], "entity_name": item["entity_name"], @@ -423,9 +422,9 @@ class PGVectorStorage(BaseVectorStorage): } return upsert_sql, data - def _upsert_relationships(self, item: dict): + def _upsert_relationships(self, item: dict[str, Any]) -> tuple[str, dict[str, Any]]: upsert_sql = SQL_TEMPLATES["upsert_relationship"] - data = { + data: dict[str, Any] = { "workspace": self.db.workspace, "id": item["__id__"], "source_id": item["src_id"], @@ -881,12 +880,6 @@ class PGGraphStorage(BaseGraphStorage): $$) AS (node_exists bool)""" % (self.graph_name, entity_name_label) single_result = (await self._query(query))[0] - logger.debug( - "{%s}:query:{%s}:result:{%s}", - inspect.currentframe().f_code.co_name, - query, - single_result["node_exists"], - ) return single_result["node_exists"] @@ -904,12 +897,7 @@ class PGGraphStorage(BaseGraphStorage): ) single_result = (await self._query(query))[0] - logger.debug( - "{%s}:query:{%s}:result:{%s}", - inspect.currentframe().f_code.co_name, - query, - single_result["edge_exists"], - ) + return single_result["edge_exists"] async def get_node(self, node_id: str) -> dict[str, str] | None: @@ -922,12 +910,7 @@ class PGGraphStorage(BaseGraphStorage): if record: node = record[0] node_dict = node["n"] - logger.debug( - "{%s}: query: {%s}, result: {%s}", - inspect.currentframe().f_code.co_name, - query, - node_dict, - ) + return node_dict return None @@ -941,12 +924,7 @@ class PGGraphStorage(BaseGraphStorage): record = (await self._query(query))[0] if record: edge_count = int(record["total_edge_count"]) - logger.debug( - "{%s}:query:{%s}:result:{%s}", - inspect.currentframe().f_code.co_name, - query, - edge_count, - ) + return edge_count async def edge_degree(self, src_id: str, tgt_id: str) -> int: @@ -958,11 +936,7 @@ class PGGraphStorage(BaseGraphStorage): trg_degree = 0 if trg_degree is None else trg_degree degrees = int(src_degree) + int(trg_degree) - logger.debug( - "{%s}:query:src_Degree+trg_degree:result:{%s}", - inspect.currentframe().f_code.co_name, - degrees, - ) + return degrees async def get_edge( @@ -983,12 +957,7 @@ class PGGraphStorage(BaseGraphStorage): record = await self._query(query) if record and record[0] and record[0]["edge_properties"]: result = record[0]["edge_properties"] - logger.debug( - "{%s}:query:{%s}:result:{%s}", - inspect.currentframe().f_code.co_name, - query, - result, - ) + return result async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None: @@ -1055,13 +1024,9 @@ class PGGraphStorage(BaseGraphStorage): try: await self._query(query, readonly=False, upsert=True) - logger.debug( - "Upserted node with label '{%s}' and properties: {%s}", - label, - properties, - ) + except Exception as e: - logger.error("Error during upsert: {%s}", e) + logger.error("POSTGRES, Error during upsert: {%s}", e) raise @retry( @@ -1097,15 +1062,10 @@ class PGGraphStorage(BaseGraphStorage): tgt_label, self._format_properties(edge_properties), ) - # logger.info(f"-- inserting edge after formatted: {params}") + try: await self._query(query, readonly=False, upsert=True) - logger.debug( - "Upserted edge from '{%s}' to '{%s}' with properties: {%s}", - src_label, - tgt_label, - edge_properties, - ) + except Exception as e: logger.error("Error during edge upsert: {%s}", e) raise diff --git a/lightrag/kg/tidb_impl.py b/lightrag/kg/tidb_impl.py index e8aaaf66..ed9c8d4b 100644 --- a/lightrag/kg/tidb_impl.py +++ b/lightrag/kg/tidb_impl.py @@ -277,7 +277,7 @@ class TiDBKVStorage(BaseKVStorage): @final @dataclass class TiDBVectorDBStorage(BaseVectorStorage): - db: TiDB = field(default=None) + db: TiDB | None = field(default=None) def __post_init__(self): self._client_file_name = os.path.join(