add namespace prefix to storage namespaces
This commit is contained in:
@@ -52,7 +52,7 @@ class MongoKVStorage(BaseKVStorage):
|
||||
return set([s for s in data if s not in existing_ids])
|
||||
|
||||
async def upsert(self, data: dict[str, dict]):
|
||||
if self.namespace == "llm_response_cache":
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
for mode, items in data.items():
|
||||
for k, v in tqdm_async(items.items(), desc="Upserting"):
|
||||
key = f"{mode}_{k}"
|
||||
@@ -69,7 +69,7 @@ class MongoKVStorage(BaseKVStorage):
|
||||
return data
|
||||
|
||||
async def get_by_mode_and_id(self, mode: str, id: str) -> Union[dict, None]:
|
||||
if "llm_response_cache" == self.namespace:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
res = {}
|
||||
v = self._data.find_one({"_id": mode + "_" + id})
|
||||
if v:
|
||||
|
@@ -185,7 +185,7 @@ class OracleKVStorage(BaseKVStorage):
|
||||
SQL = SQL_TEMPLATES["get_by_id_" + self.namespace]
|
||||
params = {"workspace": self.db.workspace, "id": id}
|
||||
# print("get_by_id:"+SQL)
|
||||
if "llm_response_cache" == self.namespace:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
array_res = await self.db.query(SQL, params, multirows=True)
|
||||
res = {}
|
||||
for row in array_res:
|
||||
@@ -201,7 +201,7 @@ class OracleKVStorage(BaseKVStorage):
|
||||
"""Specifically for llm_response_cache."""
|
||||
SQL = SQL_TEMPLATES["get_by_mode_id_" + self.namespace]
|
||||
params = {"workspace": self.db.workspace, "cache_mode": mode, "id": id}
|
||||
if "llm_response_cache" == self.namespace:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
array_res = await self.db.query(SQL, params, multirows=True)
|
||||
res = {}
|
||||
for row in array_res:
|
||||
@@ -218,7 +218,7 @@ class OracleKVStorage(BaseKVStorage):
|
||||
params = {"workspace": self.db.workspace}
|
||||
# print("get_by_ids:"+SQL)
|
||||
res = await self.db.query(SQL, params, multirows=True)
|
||||
if "llm_response_cache" == self.namespace:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
modes = set()
|
||||
dict_res: dict[str, dict] = {}
|
||||
for row in res:
|
||||
@@ -269,7 +269,7 @@ class OracleKVStorage(BaseKVStorage):
|
||||
|
||||
################ INSERT METHODS ################
|
||||
async def upsert(self, data: dict[str, dict]):
|
||||
if self.namespace == "text_chunks":
|
||||
if self.namespace.endswith("text_chunks"):
|
||||
list_data = [
|
||||
{
|
||||
"id": k,
|
||||
@@ -302,7 +302,7 @@ class OracleKVStorage(BaseKVStorage):
|
||||
"status": item["status"],
|
||||
}
|
||||
await self.db.execute(merge_sql, _data)
|
||||
if self.namespace == "full_docs":
|
||||
if self.namespace.endswith("full_docs"):
|
||||
for k, v in data.items():
|
||||
# values.clear()
|
||||
merge_sql = SQL_TEMPLATES["merge_doc_full"]
|
||||
@@ -313,7 +313,7 @@ class OracleKVStorage(BaseKVStorage):
|
||||
}
|
||||
await self.db.execute(merge_sql, _data)
|
||||
|
||||
if self.namespace == "llm_response_cache":
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
for mode, items in data.items():
|
||||
for k, v in items.items():
|
||||
upsert_sql = SQL_TEMPLATES["upsert_llm_response_cache"]
|
||||
@@ -334,8 +334,10 @@ class OracleKVStorage(BaseKVStorage):
|
||||
await self.db.execute(SQL, params)
|
||||
|
||||
async def index_done_callback(self):
|
||||
if self.namespace in ["full_docs", "text_chunks"]:
|
||||
logger.info("full doc and chunk data had been saved into oracle db!")
|
||||
for n in ("full_docs", "text_chunks"):
|
||||
if self.namespace.endswith(n):
|
||||
logger.info("full doc and chunk data had been saved into oracle db!")
|
||||
break
|
||||
|
||||
|
||||
@dataclass
|
||||
|
@@ -187,7 +187,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
"""Get doc_full data by id."""
|
||||
sql = SQL_TEMPLATES["get_by_id_" + self.namespace]
|
||||
params = {"workspace": self.db.workspace, "id": id}
|
||||
if "llm_response_cache" == self.namespace:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
array_res = await self.db.query(sql, params, multirows=True)
|
||||
res = {}
|
||||
for row in array_res:
|
||||
@@ -203,7 +203,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
"""Specifically for llm_response_cache."""
|
||||
sql = SQL_TEMPLATES["get_by_mode_id_" + self.namespace]
|
||||
params = {"workspace": self.db.workspace, mode: mode, "id": id}
|
||||
if "llm_response_cache" == self.namespace:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
array_res = await self.db.query(sql, params, multirows=True)
|
||||
res = {}
|
||||
for row in array_res:
|
||||
@@ -219,7 +219,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
ids=",".join([f"'{id}'" for id in ids])
|
||||
)
|
||||
params = {"workspace": self.db.workspace}
|
||||
if "llm_response_cache" == self.namespace:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
array_res = await self.db.query(sql, params, multirows=True)
|
||||
modes = set()
|
||||
dict_res: dict[str, dict] = {}
|
||||
@@ -239,7 +239,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
return None
|
||||
|
||||
async def all_keys(self) -> list[dict]:
|
||||
if "llm_response_cache" == self.namespace:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
sql = "select workspace,mode,id from lightrag_llm_cache"
|
||||
res = await self.db.query(sql, multirows=True)
|
||||
return res
|
||||
@@ -270,9 +270,9 @@ class PGKVStorage(BaseKVStorage):
|
||||
|
||||
################ INSERT METHODS ################
|
||||
async def upsert(self, data: Dict[str, dict]):
|
||||
if self.namespace == "text_chunks":
|
||||
if self.namespace.endswith("text_chunks"):
|
||||
pass
|
||||
elif self.namespace == "full_docs":
|
||||
elif self.namespace.endswith("full_docs"):
|
||||
for k, v in data.items():
|
||||
upsert_sql = SQL_TEMPLATES["upsert_doc_full"]
|
||||
_data = {
|
||||
@@ -281,7 +281,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
"workspace": self.db.workspace,
|
||||
}
|
||||
await self.db.execute(upsert_sql, _data)
|
||||
elif self.namespace == "llm_response_cache":
|
||||
elif self.namespace.endswith("llm_response_cache"):
|
||||
for mode, items in data.items():
|
||||
for k, v in items.items():
|
||||
upsert_sql = SQL_TEMPLATES["upsert_llm_response_cache"]
|
||||
@@ -296,8 +296,10 @@ class PGKVStorage(BaseKVStorage):
|
||||
await self.db.execute(upsert_sql, _data)
|
||||
|
||||
async def index_done_callback(self):
|
||||
if self.namespace in ["full_docs", "text_chunks"]:
|
||||
logger.info("full doc and chunk data had been saved into postgresql db!")
|
||||
for n in ("full_docs", "text_chunks"):
|
||||
if self.namespace.endswith(n):
|
||||
logger.info("full doc and chunk data had been saved into postgresql db!")
|
||||
break
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -389,11 +391,11 @@ class PGVectorStorage(BaseVectorStorage):
|
||||
for i, d in enumerate(list_data):
|
||||
d["__vector__"] = embeddings[i]
|
||||
for item in list_data:
|
||||
if self.namespace == "chunks":
|
||||
if self.namespace.endswith("chunks"):
|
||||
upsert_sql, data = self._upsert_chunks(item)
|
||||
elif self.namespace == "entities":
|
||||
elif self.namespace.endswith("entities"):
|
||||
upsert_sql, data = self._upsert_entities(item)
|
||||
elif self.namespace == "relationships":
|
||||
elif self.namespace.endswith("relationships"):
|
||||
upsert_sql, data = self._upsert_relationships(item)
|
||||
else:
|
||||
raise ValueError(f"{self.namespace} is not supported")
|
||||
|
@@ -160,7 +160,7 @@ class TiDBKVStorage(BaseKVStorage):
|
||||
async def upsert(self, data: dict[str, dict]):
|
||||
left_data = {k: v for k, v in data.items() if k not in self._data}
|
||||
self._data.update(left_data)
|
||||
if self.namespace == "text_chunks":
|
||||
if self.namespace.endswith("text_chunks"):
|
||||
list_data = [
|
||||
{
|
||||
"__id__": k,
|
||||
@@ -196,7 +196,7 @@ class TiDBKVStorage(BaseKVStorage):
|
||||
)
|
||||
await self.db.execute(merge_sql, data)
|
||||
|
||||
if self.namespace == "full_docs":
|
||||
if self.namespace.endswith("full_docs"):
|
||||
merge_sql = SQL_TEMPLATES["upsert_doc_full"]
|
||||
data = []
|
||||
for k, v in self._data.items():
|
||||
@@ -211,8 +211,10 @@ class TiDBKVStorage(BaseKVStorage):
|
||||
return left_data
|
||||
|
||||
async def index_done_callback(self):
|
||||
if self.namespace in ["full_docs", "text_chunks"]:
|
||||
logger.info("full doc and chunk data had been saved into TiDB db!")
|
||||
for n in ("full_docs", "text_chunks"):
|
||||
if self.namespace.endswith(n):
|
||||
logger.info("full doc and chunk data had been saved into TiDB db!")
|
||||
break
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -258,7 +260,7 @@ class TiDBVectorDBStorage(BaseVectorStorage):
|
||||
if not len(data):
|
||||
logger.warning("You insert an empty data to vector DB")
|
||||
return []
|
||||
if self.namespace == "chunks":
|
||||
if self.namespace.endswith("chunks"):
|
||||
return []
|
||||
logger.info(f"Inserting {len(data)} vectors to {self.namespace}")
|
||||
|
||||
@@ -288,7 +290,7 @@ class TiDBVectorDBStorage(BaseVectorStorage):
|
||||
for i, d in enumerate(list_data):
|
||||
d["content_vector"] = embeddings[i]
|
||||
|
||||
if self.namespace == "entities":
|
||||
if self.namespace.endswith("entities"):
|
||||
data = []
|
||||
for item in list_data:
|
||||
param = {
|
||||
@@ -309,7 +311,7 @@ class TiDBVectorDBStorage(BaseVectorStorage):
|
||||
merge_sql = SQL_TEMPLATES["insert_entity"]
|
||||
await self.db.execute(merge_sql, data)
|
||||
|
||||
elif self.namespace == "relationships":
|
||||
elif self.namespace.endswith("relationships"):
|
||||
data = []
|
||||
for item in list_data:
|
||||
param = {
|
||||
|
Reference in New Issue
Block a user