Update oracle_impl.py

lightrag/kg/oracle_impl.py

@@ -114,7 +114,9 @@ class OracleDB:
 
         logger.info("Finished check all tables in Oracle database")
 
-    async def query(self, sql: str, params: dict = None, multirows: bool = False) -> Union[dict, None]:
+    async def query(
+        self, sql: str, params: dict = None, multirows: bool = False
+    ) -> Union[dict, None]:
         async with self.pool.acquire() as connection:
             connection.inputtypehandler = self.input_type_handler
             connection.outputtypehandler = self.output_type_handler
@@ -174,9 +176,9 @@ class OracleKVStorage(BaseKVStorage):
     async def get_by_id(self, id: str) -> Union[dict, None]:
         """根据 id 获取 doc_full 数据."""
         SQL = SQL_TEMPLATES["get_by_id_" + self.namespace]
-        params = {"workspace":self.db.workspace, "id":id}
+        params = {"workspace": self.db.workspace, "id": id}
         # print("get_by_id:"+SQL)
-        res = await self.db.query(SQL,params)
+        res = await self.db.query(SQL, params)
         if res:
             data = res  # {"data":res}
             # print (data)
@@ -187,11 +189,13 @@ class OracleKVStorage(BaseKVStorage):
     # Query by id
     async def get_by_ids(self, ids: list[str], fields=None) -> Union[list[dict], None]:
         """根据 id 获取 doc_chunks 数据"""
-        SQL = SQL_TEMPLATES["get_by_ids_" + self.namespace].format(ids=",".join([f"'{id}'" for id in ids]))
-        params = {"workspace":self.db.workspace}
-        #print("get_by_ids:"+SQL)
-        #print(params)
-        res = await self.db.query(SQL,params, multirows=True)
+        SQL = SQL_TEMPLATES["get_by_ids_" + self.namespace].format(
+            ids=",".join([f"'{id}'" for id in ids])
+        )
+        params = {"workspace": self.db.workspace}
+        # print("get_by_ids:"+SQL)
+        # print(params)
+        res = await self.db.query(SQL, params, multirows=True)
         if res:
             data = res  # [{"data":i} for i in res]
             # print(data)
@@ -201,16 +205,17 @@ class OracleKVStorage(BaseKVStorage):
 
     async def filter_keys(self, keys: list[str]) -> set[str]:
         """过滤掉重复内容"""
-        SQL = SQL_TEMPLATES["filter_keys"].format(table_name=N_T[self.namespace],
-                                                  ids=",".join([f"'{id}'" for id in keys]))
-        params = {"workspace":self.db.workspace}
+        SQL = SQL_TEMPLATES["filter_keys"].format(
+            table_name=N_T[self.namespace], ids=",".join([f"'{id}'" for id in keys])
+        )
+        params = {"workspace": self.db.workspace}
         try:
             await self.db.query(SQL, params)
         except Exception as e:
             logger.error(f"Oracle database error: {e}")
             print(SQL)
             print(params)
-        res = await self.db.query(SQL, params,multirows=True)
+        res = await self.db.query(SQL, params, multirows=True)
         data = None
         if res:
             exist_keys = [key["id"] for key in res]
@@ -248,15 +253,16 @@ class OracleKVStorage(BaseKVStorage):
             # print(list_data)
             for item in list_data:
                 merge_sql = SQL_TEMPLATES["merge_chunk"]
-                data = {"check_id":item["__id__"],
-                        "id":item["__id__"],
-                        "content":item["content"],
-                        "workspace":self.db.workspace,
-                        "tokens":item["tokens"],
-                        "chunk_order_index":item["chunk_order_index"],
-                        "full_doc_id":item["full_doc_id"],
-                        "content_vector":item["__vector__"]
-                    }
+                data = {
+                    "check_id": item["__id__"],
+                    "id": item["__id__"],
+                    "content": item["content"],
+                    "workspace": self.db.workspace,
+                    "tokens": item["tokens"],
+                    "chunk_order_index": item["chunk_order_index"],
+                    "full_doc_id": item["full_doc_id"],
+                    "content_vector": item["__vector__"],
+                }
                 # print(merge_sql)
                 await self.db.execute(merge_sql, data)
 
@@ -265,11 +271,11 @@ class OracleKVStorage(BaseKVStorage):
                 # values.clear()
                 merge_sql = SQL_TEMPLATES["merge_doc_full"]
                 data = {
-                    "check_id":k,
-                    "id":k,
-                    "content":v["content"],
-                    "workspace":self.db.workspace
-                    }
+                    "check_id": k,
+                    "id": k,
+                    "content": v["content"],
+                    "workspace": self.db.workspace,
+                }
                 # print(merge_sql)
                 await self.db.execute(merge_sql, data)
         return left_data
@@ -301,7 +307,7 @@ class OracleVectorDBStorage(BaseVectorStorage):
         # 转换精度
         dtype = str(embedding.dtype).upper()
         dimension = embedding.shape[0]
-        embedding_string = "["+", ".join(map(str, embedding.tolist()))+"]"
+        embedding_string = "[" + ", ".join(map(str, embedding.tolist())) + "]"
 
         SQL = SQL_TEMPLATES[self.namespace].format(dimension=dimension, dtype=dtype)
         params = {
@@ -309,9 +315,9 @@ class OracleVectorDBStorage(BaseVectorStorage):
             "workspace": self.db.workspace,
             "top_k": top_k,
             "better_than_threshold": self.cosine_better_than_threshold,
-            }
+        }
         # print(SQL)
-        results = await self.db.query(SQL,params=params, multirows=True)
+        results = await self.db.query(SQL, params=params, multirows=True)
         # print("vector search result:",results)
         return results
 
@@ -346,18 +352,18 @@ class OracleGraphStorage(BaseGraphStorage):
         )
         embeddings = np.concatenate(embeddings_list)
         content_vector = embeddings[0]
-        merge_sql = SQL_TEMPLATES["merge_node"]        
+        merge_sql = SQL_TEMPLATES["merge_node"]
         data = {
-            "workspace":self.db.workspace,
-            "name":entity_name,
-            "entity_type":entity_type,
-            "description":description,
-            "source_chunk_id":source_id,
-            "content":content,
-            "content_vector":content_vector
-            }        
+            "workspace": self.db.workspace,
+            "name": entity_name,
+            "entity_type": entity_type,
+            "description": description,
+            "source_chunk_id": source_id,
+            "content": content,
+            "content_vector": content_vector,
+        }
         # print(merge_sql)
-        await self.db.execute(merge_sql,data)
+        await self.db.execute(merge_sql, data)
         # self._graph.add_node(node_id, **node_data)
 
     async def upsert_edge(
@@ -371,7 +377,9 @@ class OracleGraphStorage(BaseGraphStorage):
         keywords = edge_data["keywords"]
         description = edge_data["description"]
         source_chunk_id = edge_data["source_id"]
-        logger.debug(f"source_name:{source_name}, target_name:{target_name}, keywords: {keywords}")
+        logger.debug(
+            f"source_name:{source_name}, target_name:{target_name}, keywords: {keywords}"
+        )
 
         content = keywords + source_name + target_name + description
         contents = [content]
@@ -384,20 +392,20 @@ class OracleGraphStorage(BaseGraphStorage):
         )
         embeddings = np.concatenate(embeddings_list)
         content_vector = embeddings[0]
-        merge_sql = SQL_TEMPLATES["merge_edge"]        
+        merge_sql = SQL_TEMPLATES["merge_edge"]
         data = {
-            "workspace":self.db.workspace,
-            "source_name":source_name,
-            "target_name":target_name,
-            "weight":weight,
-            "keywords":keywords,
-            "description":description,
-            "source_chunk_id":source_chunk_id,
-            "content":content,
-            "content_vector":content_vector
-             }
+            "workspace": self.db.workspace,
+            "source_name": source_name,
+            "target_name": target_name,
+            "weight": weight,
+            "keywords": keywords,
+            "description": description,
+            "source_chunk_id": source_chunk_id,
+            "content": content,
+            "content_vector": content_vector,
+        }
         # print(merge_sql)
-        await self.db.execute(merge_sql,data)
+        await self.db.execute(merge_sql, data)
         # self._graph.add_edge(source_node_id, target_node_id, **edge_data)
 
     async def embed_nodes(self, algorithm: str) -> tuple[np.ndarray, list[str]]:
@@ -428,13 +436,10 @@ class OracleGraphStorage(BaseGraphStorage):
     async def has_node(self, node_id: str) -> bool:
         """根据节点id检查节点是否存在"""
         SQL = SQL_TEMPLATES["has_node"]
-        params = {
-            "workspace":self.db.workspace,
-            "node_id":node_id
-            }
+        params = {"workspace": self.db.workspace, "node_id": node_id}
         # print(SQL)
         # print(self.db.workspace, node_id)
-        res = await self.db.query(SQL,params)
+        res = await self.db.query(SQL, params)
         if res:
             # print("Node exist!",res)
             return True
@@ -446,12 +451,12 @@ class OracleGraphStorage(BaseGraphStorage):
         """根据源和目标节点id检查边是否存在"""
         SQL = SQL_TEMPLATES["has_edge"]
         params = {
-            "workspace":self.db.workspace,
-            "source_node_id":source_node_id,
-            "target_node_id":target_node_id
-            }
+            "workspace": self.db.workspace,
+            "source_node_id": source_node_id,
+            "target_node_id": target_node_id,
+        }
         # print(SQL)
-        res = await self.db.query(SQL,params)
+        res = await self.db.query(SQL, params)
         if res:
             # print("Edge exist!",res)
             return True
@@ -462,12 +467,9 @@ class OracleGraphStorage(BaseGraphStorage):
     async def node_degree(self, node_id: str) -> int:
         """根据节���id获取节点的度"""
         SQL = SQL_TEMPLATES["node_degree"]
-        params = {
-            "workspace":self.db.workspace,
-            "node_id":node_id
-            }
+        params = {"workspace": self.db.workspace, "node_id": node_id}
         # print(SQL)
-        res = await self.db.query(SQL,params)
+        res = await self.db.query(SQL, params)
         if res:
             # print("Node degree",res["degree"])
             return res["degree"]
@@ -484,13 +486,10 @@ class OracleGraphStorage(BaseGraphStorage):
     async def get_node(self, node_id: str) -> Union[dict, None]:
         """根据节点id获取节点数据"""
         SQL = SQL_TEMPLATES["get_node"]
-        params = {
-            "workspace":self.db.workspace,
-            "node_id":node_id
-            }
+        params = {"workspace": self.db.workspace, "node_id": node_id}
         # print(self.db.workspace, node_id)
         # print(SQL)
-        res = await self.db.query(SQL,params)
+        res = await self.db.query(SQL, params)
         if res:
             # print("Get node!",self.db.workspace, node_id,res)
             return res
@@ -504,11 +503,11 @@ class OracleGraphStorage(BaseGraphStorage):
         """根据源和目标节点id获取边"""
         SQL = SQL_TEMPLATES["get_edge"]
         params = {
-            "workspace":self.db.workspace,
-            "source_node_id":source_node_id,
-            "target_node_id":target_node_id
-            }
-        res = await self.db.query(SQL,params)
+            "workspace": self.db.workspace,
+            "source_node_id": source_node_id,
+            "target_node_id": target_node_id,
+        }
+        res = await self.db.query(SQL, params)
         if res:
             # print("Get edge!",self.db.workspace, source_node_id, target_node_id,res[0])
             return res
@@ -520,10 +519,7 @@ class OracleGraphStorage(BaseGraphStorage):
         """根据节点id获取节点的所有边"""
         if await self.has_node(source_node_id):
             SQL = SQL_TEMPLATES["get_node_edges"]
-            params = {
-                "workspace":self.db.workspace,
-                "source_node_id":source_node_id
-                }            
+            params = {"workspace": self.db.workspace, "source_node_id": source_node_id}
             res = await self.db.query(sql=SQL, params=params, multirows=True)
             if res:
                 data = [(i["source_name"], i["target_name"]) for i in res]
@@ -532,30 +528,31 @@ class OracleGraphStorage(BaseGraphStorage):
             else:
                 # print("Node Edge not exist!",self.db.workspace, source_node_id)
                 return []
-    
+
     async def get_all_nodes(self, limit: int):
         """查询所有节点"""
         SQL = SQL_TEMPLATES["get_all_nodes"]
-        params = {"workspace":self.db.workspace, "limit":str(limit)}
-        res = await self.db.query(sql=SQL,params=params, multirows=True)
+        params = {"workspace": self.db.workspace, "limit": str(limit)}
+        res = await self.db.query(sql=SQL, params=params, multirows=True)
         if res:
             return res
 
     async def get_all_edges(self, limit: int):
         """查询所有边"""
         SQL = SQL_TEMPLATES["get_all_edges"]
-        params = {"workspace":self.db.workspace, "limit":str(limit)}
-        res = await self.db.query(sql=SQL,params=params, multirows=True)
+        params = {"workspace": self.db.workspace, "limit": str(limit)}
+        res = await self.db.query(sql=SQL, params=params, multirows=True)
         if res:
             return res
-    
+
     async def get_statistics(self):
         SQL = SQL_TEMPLATES["get_statistics"]
-        params = {"workspace":self.db.workspace}
-        res = await self.db.query(sql=SQL,params=params, multirows=True)
+        params = {"workspace": self.db.workspace}
+        res = await self.db.query(sql=SQL, params=params, multirows=True)
         if res:
             return res
 
+
 N_T = {
     "full_docs": "LIGHTRAG_DOC_FULL",
     "text_chunks": "LIGHTRAG_DOC_CHUNKS",
@@ -727,7 +724,7 @@ SQL_TEMPLATES = {
                 WHEN NOT MATCHED THEN
                     INSERT(workspace,source_name,target_name,weight,keywords,description,source_chunk_id,content,content_vector)
                     values (:workspace,:source_name,:target_name,:weight,:keywords,:description,:source_chunk_id,:content,:content_vector) """,
-    "get_all_nodes":"""WITH t0 AS (
+    "get_all_nodes": """WITH t0 AS (
                         SELECT name AS id, entity_type AS label, entity_type, description,
                             '["' || replace(source_chunk_id, '<SEP>', '","') || '"]'     source_chunk_ids
                         FROM lightrag_graph_nodes
@@ -743,20 +740,20 @@ SQL_TEMPLATES = {
                     )
                     SELECT t0.id, label, entity_type, description, t2.content
                     FROM t0 LEFT JOIN t2 ON t0.id = t2.id""",
-    "get_all_edges":"""SELECT t1.id,t1.keywords as label,t1.keywords, t1.source_name as source, t1.target_name as target,
+    "get_all_edges": """SELECT t1.id,t1.keywords as label,t1.keywords, t1.source_name as source, t1.target_name as target,
                 t1.weight,t1.DESCRIPTION,t2.content
                 FROM LIGHTRAG_GRAPH_EDGES t1
                 LEFT JOIN LIGHTRAG_DOC_CHUNKS t2 on t1.source_chunk_id=t2.id
                 WHERE t1.workspace=:workspace
                 order by t1.CREATETIME DESC
                 fetch first :limit rows only""",
-    "get_statistics":"""select  count(distinct CASE WHEN type='node' THEN id END) as nodes_count,
+    "get_statistics": """select  count(distinct CASE WHEN type='node' THEN id END) as nodes_count,
                 count(distinct CASE WHEN type='edge' THEN id END) as edges_count
                 FROM (
-                select 'node' as type, id FROM GRAPH_TABLE (lightrag_graph 
+                select 'node' as type, id FROM GRAPH_TABLE (lightrag_graph
                     MATCH (a) WHERE a.workspace=:workspace columns(a.name as id))
                 UNION
-                select 'edge' as type, TO_CHAR(id) id FROM GRAPH_TABLE (lightrag_graph 
+                select 'edge' as type, TO_CHAR(id) id FROM GRAPH_TABLE (lightrag_graph
                     MATCH (a)-[e]->(b) WHERE e.workspace=:workspace columns(e.id))
                 )""",
 }