cleanup code
Browse files- lightrag/kg/oracle_impl.py +0 -17
- lightrag/kg/postgres_impl.py +15 -15
- lightrag/kg/tidb_impl.py +3 -8
- lightrag/llm/openai.py +15 -15
lightrag/kg/oracle_impl.py
CHANGED
@@ -140,8 +140,6 @@ class OracleDB:
|
|
140 |
await cursor.execute(sql, params)
|
141 |
except Exception as e:
|
142 |
logger.error(f"Oracle database error: {e}")
|
143 |
-
print(sql)
|
144 |
-
print(params)
|
145 |
raise
|
146 |
columns = [column[0].lower() for column in cursor.description]
|
147 |
if multirows:
|
@@ -172,8 +170,6 @@ class OracleDB:
|
|
172 |
await connection.commit()
|
173 |
except Exception as e:
|
174 |
logger.error(f"Oracle database error: {e}")
|
175 |
-
print(sql)
|
176 |
-
print(data)
|
177 |
raise
|
178 |
|
179 |
|
@@ -349,9 +345,7 @@ class OracleVectorDBStorage(BaseVectorStorage):
|
|
349 |
"top_k": top_k,
|
350 |
"better_than_threshold": self.cosine_better_than_threshold,
|
351 |
}
|
352 |
-
# print(SQL)
|
353 |
results = await self.db.query(SQL, params=params, multirows=True)
|
354 |
-
# print("vector search result:",results)
|
355 |
return results
|
356 |
|
357 |
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
@@ -477,8 +471,6 @@ class OracleGraphStorage(BaseGraphStorage):
|
|
477 |
"""根据节点id检查节点是否存在"""
|
478 |
SQL = SQL_TEMPLATES["has_node"]
|
479 |
params = {"workspace": self.db.workspace, "node_id": node_id}
|
480 |
-
# print(SQL)
|
481 |
-
# print(self.db.workspace, node_id)
|
482 |
res = await self.db.query(SQL, params)
|
483 |
if res:
|
484 |
# print("Node exist!",res)
|
@@ -494,7 +486,6 @@ class OracleGraphStorage(BaseGraphStorage):
|
|
494 |
"source_node_id": source_node_id,
|
495 |
"target_node_id": target_node_id,
|
496 |
}
|
497 |
-
# print(SQL)
|
498 |
res = await self.db.query(SQL, params)
|
499 |
if res:
|
500 |
# print("Edge exist!",res)
|
@@ -506,33 +497,25 @@ class OracleGraphStorage(BaseGraphStorage):
|
|
506 |
async def node_degree(self, node_id: str) -> int:
|
507 |
SQL = SQL_TEMPLATES["node_degree"]
|
508 |
params = {"workspace": self.db.workspace, "node_id": node_id}
|
509 |
-
# print(SQL)
|
510 |
res = await self.db.query(SQL, params)
|
511 |
if res:
|
512 |
-
# print("Node degree",res["degree"])
|
513 |
return res["degree"]
|
514 |
else:
|
515 |
-
# print("Edge not exist!")
|
516 |
return 0
|
517 |
|
518 |
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
|
519 |
"""根据源和目标节点id获取边的度"""
|
520 |
degree = await self.node_degree(src_id) + await self.node_degree(tgt_id)
|
521 |
-
# print("Edge degree",degree)
|
522 |
return degree
|
523 |
|
524 |
async def get_node(self, node_id: str) -> dict[str, str] | None:
|
525 |
"""根据节点id获取节点数据"""
|
526 |
SQL = SQL_TEMPLATES["get_node"]
|
527 |
params = {"workspace": self.db.workspace, "node_id": node_id}
|
528 |
-
# print(self.db.workspace, node_id)
|
529 |
-
# print(SQL)
|
530 |
res = await self.db.query(SQL, params)
|
531 |
if res:
|
532 |
-
# print("Get node!",self.db.workspace, node_id,res)
|
533 |
return res
|
534 |
else:
|
535 |
-
# print("Can't get node!",self.db.workspace, node_id)
|
536 |
return None
|
537 |
|
538 |
async def get_edge(
|
|
|
140 |
await cursor.execute(sql, params)
|
141 |
except Exception as e:
|
142 |
logger.error(f"Oracle database error: {e}")
|
|
|
|
|
143 |
raise
|
144 |
columns = [column[0].lower() for column in cursor.description]
|
145 |
if multirows:
|
|
|
170 |
await connection.commit()
|
171 |
except Exception as e:
|
172 |
logger.error(f"Oracle database error: {e}")
|
|
|
|
|
173 |
raise
|
174 |
|
175 |
|
|
|
345 |
"top_k": top_k,
|
346 |
"better_than_threshold": self.cosine_better_than_threshold,
|
347 |
}
|
|
|
348 |
results = await self.db.query(SQL, params=params, multirows=True)
|
|
|
349 |
return results
|
350 |
|
351 |
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
|
|
471 |
"""根据节点id检查节点是否存在"""
|
472 |
SQL = SQL_TEMPLATES["has_node"]
|
473 |
params = {"workspace": self.db.workspace, "node_id": node_id}
|
|
|
|
|
474 |
res = await self.db.query(SQL, params)
|
475 |
if res:
|
476 |
# print("Node exist!",res)
|
|
|
486 |
"source_node_id": source_node_id,
|
487 |
"target_node_id": target_node_id,
|
488 |
}
|
|
|
489 |
res = await self.db.query(SQL, params)
|
490 |
if res:
|
491 |
# print("Edge exist!",res)
|
|
|
497 |
async def node_degree(self, node_id: str) -> int:
|
498 |
SQL = SQL_TEMPLATES["node_degree"]
|
499 |
params = {"workspace": self.db.workspace, "node_id": node_id}
|
|
|
500 |
res = await self.db.query(SQL, params)
|
501 |
if res:
|
|
|
502 |
return res["degree"]
|
503 |
else:
|
|
|
504 |
return 0
|
505 |
|
506 |
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
|
507 |
"""根据源和目标节点id获取边的度"""
|
508 |
degree = await self.node_degree(src_id) + await self.node_degree(tgt_id)
|
|
|
509 |
return degree
|
510 |
|
511 |
async def get_node(self, node_id: str) -> dict[str, str] | None:
|
512 |
"""根据节点id获取节点数据"""
|
513 |
SQL = SQL_TEMPLATES["get_node"]
|
514 |
params = {"workspace": self.db.workspace, "node_id": node_id}
|
|
|
|
|
515 |
res = await self.db.query(SQL, params)
|
516 |
if res:
|
|
|
517 |
return res
|
518 |
else:
|
|
|
519 |
return None
|
520 |
|
521 |
async def get_edge(
|
lightrag/kg/postgres_impl.py
CHANGED
@@ -136,9 +136,9 @@ class PostgreSQLDB:
|
|
136 |
data = None
|
137 |
return data
|
138 |
except Exception as e:
|
139 |
-
logger.error(
|
140 |
-
|
141 |
-
|
142 |
raise
|
143 |
|
144 |
async def execute(
|
@@ -167,9 +167,7 @@ class PostgreSQLDB:
|
|
167 |
else:
|
168 |
logger.error(f"Upsert error: {e}")
|
169 |
except Exception as e:
|
170 |
-
logger.error(f"PostgreSQL database
|
171 |
-
print(sql)
|
172 |
-
print(data)
|
173 |
raise
|
174 |
|
175 |
@staticmethod
|
@@ -266,9 +264,10 @@ class PGKVStorage(BaseKVStorage):
|
|
266 |
new_keys = set([s for s in keys if s not in exist_keys])
|
267 |
return new_keys
|
268 |
except Exception as e:
|
269 |
-
logger.error(
|
270 |
-
|
271 |
-
|
|
|
272 |
|
273 |
################ INSERT METHODS ################
|
274 |
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
@@ -333,9 +332,9 @@ class PGVectorStorage(BaseVectorStorage):
|
|
333 |
"content_vector": json.dumps(item["__vector__"].tolist()),
|
334 |
}
|
335 |
except Exception as e:
|
336 |
-
logger.error(f"Error to prepare upsert
|
337 |
-
|
338 |
-
|
339 |
return upsert_sql, data
|
340 |
|
341 |
def _upsert_entities(self, item: dict):
|
@@ -454,9 +453,10 @@ class PGDocStatusStorage(DocStatusStorage):
|
|
454 |
print(f"new_keys: {new_keys}")
|
455 |
return new_keys
|
456 |
except Exception as e:
|
457 |
-
logger.error(
|
458 |
-
|
459 |
-
|
|
|
460 |
|
461 |
async def get_by_id(self, id: str) -> Union[dict[str, Any], None]:
|
462 |
sql = "select * from LIGHTRAG_DOC_STATUS where workspace=$1 and id=$2"
|
|
|
136 |
data = None
|
137 |
return data
|
138 |
except Exception as e:
|
139 |
+
logger.error(
|
140 |
+
f"PostgreSQL database,\nsql:{sql},\nparams:{params},\nerror:{e}"
|
141 |
+
)
|
142 |
raise
|
143 |
|
144 |
async def execute(
|
|
|
167 |
else:
|
168 |
logger.error(f"Upsert error: {e}")
|
169 |
except Exception as e:
|
170 |
+
logger.error(f"PostgreSQL database,\nsql:{sql},\ndata:{data},\nerror:{e}")
|
|
|
|
|
171 |
raise
|
172 |
|
173 |
@staticmethod
|
|
|
264 |
new_keys = set([s for s in keys if s not in exist_keys])
|
265 |
return new_keys
|
266 |
except Exception as e:
|
267 |
+
logger.error(
|
268 |
+
f"PostgreSQL database,\nsql:{sql},\nparams:{params},\nerror:{e}"
|
269 |
+
)
|
270 |
+
raise
|
271 |
|
272 |
################ INSERT METHODS ################
|
273 |
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
|
|
332 |
"content_vector": json.dumps(item["__vector__"].tolist()),
|
333 |
}
|
334 |
except Exception as e:
|
335 |
+
logger.error(f"Error to prepare upsert,\nsql: {e}\nitem: {item}")
|
336 |
+
raise
|
337 |
+
|
338 |
return upsert_sql, data
|
339 |
|
340 |
def _upsert_entities(self, item: dict):
|
|
|
453 |
print(f"new_keys: {new_keys}")
|
454 |
return new_keys
|
455 |
except Exception as e:
|
456 |
+
logger.error(
|
457 |
+
f"PostgreSQL database,\nsql:{sql},\nparams:{params},\nerror:{e}"
|
458 |
+
)
|
459 |
+
raise
|
460 |
|
461 |
async def get_by_id(self, id: str) -> Union[dict[str, Any], None]:
|
462 |
sql = "select * from LIGHTRAG_DOC_STATUS where workspace=$1 and id=$2"
|
lightrag/kg/tidb_impl.py
CHANGED
@@ -76,9 +76,7 @@ class TiDB:
|
|
76 |
try:
|
77 |
result = conn.execute(text(sql), params)
|
78 |
except Exception as e:
|
79 |
-
logger.error(f"Tidb database
|
80 |
-
print(sql)
|
81 |
-
print(params)
|
82 |
raise
|
83 |
if multirows:
|
84 |
rows = result.all()
|
@@ -103,9 +101,7 @@ class TiDB:
|
|
103 |
else:
|
104 |
conn.execute(text(sql), parameters=data)
|
105 |
except Exception as e:
|
106 |
-
logger.error(f"
|
107 |
-
print(sql)
|
108 |
-
print(data)
|
109 |
raise
|
110 |
|
111 |
|
@@ -145,8 +141,7 @@ class TiDBKVStorage(BaseKVStorage):
|
|
145 |
try:
|
146 |
await self.db.query(SQL)
|
147 |
except Exception as e:
|
148 |
-
logger.error(f"Tidb database
|
149 |
-
print(SQL)
|
150 |
res = await self.db.query(SQL, multirows=True)
|
151 |
if res:
|
152 |
exist_keys = [key["id"] for key in res]
|
|
|
76 |
try:
|
77 |
result = conn.execute(text(sql), params)
|
78 |
except Exception as e:
|
79 |
+
logger.error(f"Tidb database,\nsql:{sql},\nparams:{params},\nerror:{e}")
|
|
|
|
|
80 |
raise
|
81 |
if multirows:
|
82 |
rows = result.all()
|
|
|
101 |
else:
|
102 |
conn.execute(text(sql), parameters=data)
|
103 |
except Exception as e:
|
104 |
+
logger.error(f"Tidb database,\nsql:{sql},\ndata:{data},\nerror:{e}")
|
|
|
|
|
105 |
raise
|
106 |
|
107 |
|
|
|
141 |
try:
|
142 |
await self.db.query(SQL)
|
143 |
except Exception as e:
|
144 |
+
logger.error(f"Tidb database,\nsql:{SQL},\nkeys:{keys},\nerror:{e}")
|
|
|
145 |
res = await self.db.query(SQL, multirows=True)
|
146 |
if res:
|
147 |
exist_keys = [key["id"] for key in res]
|
lightrag/llm/openai.py
CHANGED
@@ -77,7 +77,7 @@ from lightrag.types import GPTKeywordExtractionFormat
|
|
77 |
from lightrag.api import __api_version__
|
78 |
|
79 |
import numpy as np
|
80 |
-
from typing import Union
|
81 |
|
82 |
|
83 |
class InvalidResponseError(Exception):
|
@@ -94,13 +94,13 @@ class InvalidResponseError(Exception):
|
|
94 |
),
|
95 |
)
|
96 |
async def openai_complete_if_cache(
|
97 |
-
model,
|
98 |
-
prompt,
|
99 |
-
system_prompt=None,
|
100 |
-
history_messages=None,
|
101 |
-
base_url=None,
|
102 |
-
api_key=None,
|
103 |
-
**kwargs,
|
104 |
) -> str:
|
105 |
if history_messages is None:
|
106 |
history_messages = []
|
@@ -125,7 +125,7 @@ async def openai_complete_if_cache(
|
|
125 |
)
|
126 |
kwargs.pop("hashing_kv", None)
|
127 |
kwargs.pop("keyword_extraction", None)
|
128 |
-
messages = []
|
129 |
if system_prompt:
|
130 |
messages.append({"role": "system", "content": system_prompt})
|
131 |
messages.extend(history_messages)
|
@@ -147,18 +147,18 @@ async def openai_complete_if_cache(
|
|
147 |
model=model, messages=messages, **kwargs
|
148 |
)
|
149 |
except APIConnectionError as e:
|
150 |
-
logger.error(f"OpenAI API Connection Error: {
|
151 |
raise
|
152 |
except RateLimitError as e:
|
153 |
-
logger.error(f"OpenAI API Rate Limit Error: {
|
154 |
raise
|
155 |
except APITimeoutError as e:
|
156 |
-
logger.error(f"OpenAI API Timeout Error: {
|
157 |
raise
|
158 |
except Exception as e:
|
159 |
-
logger.error(
|
160 |
-
|
161 |
-
|
162 |
raise
|
163 |
|
164 |
if hasattr(response, "__aiter__"):
|
|
|
77 |
from lightrag.api import __api_version__
|
78 |
|
79 |
import numpy as np
|
80 |
+
from typing import Any, Union
|
81 |
|
82 |
|
83 |
class InvalidResponseError(Exception):
|
|
|
94 |
),
|
95 |
)
|
96 |
async def openai_complete_if_cache(
|
97 |
+
model: str,
|
98 |
+
prompt: str,
|
99 |
+
system_prompt: str | None = None,
|
100 |
+
history_messages: list[dict[str, Any]] | None = None,
|
101 |
+
base_url: str | None = None,
|
102 |
+
api_key: str | None = None,
|
103 |
+
**kwargs: Any,
|
104 |
) -> str:
|
105 |
if history_messages is None:
|
106 |
history_messages = []
|
|
|
125 |
)
|
126 |
kwargs.pop("hashing_kv", None)
|
127 |
kwargs.pop("keyword_extraction", None)
|
128 |
+
messages: list[dict[str, Any]] = []
|
129 |
if system_prompt:
|
130 |
messages.append({"role": "system", "content": system_prompt})
|
131 |
messages.extend(history_messages)
|
|
|
147 |
model=model, messages=messages, **kwargs
|
148 |
)
|
149 |
except APIConnectionError as e:
|
150 |
+
logger.error(f"OpenAI API Connection Error: {e}")
|
151 |
raise
|
152 |
except RateLimitError as e:
|
153 |
+
logger.error(f"OpenAI API Rate Limit Error: {e}")
|
154 |
raise
|
155 |
except APITimeoutError as e:
|
156 |
+
logger.error(f"OpenAI API Timeout Error: {e}")
|
157 |
raise
|
158 |
except Exception as e:
|
159 |
+
logger.error(
|
160 |
+
f"OpenAI API Call Failed,\nModel: {model},\nParams: {kwargs}, Got: {e}"
|
161 |
+
)
|
162 |
raise
|
163 |
|
164 |
if hasattr(response, "__aiter__"):
|