better handling of namespace
This commit is contained in:
ArnoChen
@@ -32,6 +32,7 @@ from ..base import (
|
||||
BaseGraphStorage,
|
||||
T,
|
||||
)
|
||||
from ..namespace import NameSpace, is_namespace
|
||||
|
||||
if sys.platform.startswith("win"):
|
||||
import asyncio.windows_events
|
||||
@@ -187,7 +188,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 self.namespace.endswith("llm_response_cache"):
|
||||
if is_namespace(self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE):
|
||||
array_res = await self.db.query(sql, params, multirows=True)
|
||||
res = {}
|
||||
for row in array_res:
|
||||
@@ -203,7 +204,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 self.namespace.endswith("llm_response_cache"):
|
||||
if is_namespace(self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE):
|
||||
array_res = await self.db.query(sql, params, multirows=True)
|
||||
res = {}
|
||||
for row in array_res:
|
||||
@@ -219,7 +220,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
ids=",".join([f"'{id}'" for id in ids])
|
||||
)
|
||||
params = {"workspace": self.db.workspace}
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
if is_namespace(self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE):
|
||||
array_res = await self.db.query(sql, params, multirows=True)
|
||||
modes = set()
|
||||
dict_res: dict[str, dict] = {}
|
||||
@@ -239,7 +240,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
return None
|
||||
|
||||
async def all_keys(self) -> list[dict]:
|
||||
if self.namespace.endswith("llm_response_cache"):
|
||||
if is_namespace(self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE):
|
||||
sql = "select workspace,mode,id from lightrag_llm_cache"
|
||||
res = await self.db.query(sql, multirows=True)
|
||||
return res
|
||||
@@ -251,7 +252,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
async def filter_keys(self, keys: List[str]) -> Set[str]:
|
||||
"""Filter out duplicated content"""
|
||||
sql = SQL_TEMPLATES["filter_keys"].format(
|
||||
table_name=NAMESPACE_TABLE_MAP[self.namespace],
|
||||
table_name=namespace_to_table_name(self.namespace),
|
||||
ids=",".join([f"'{id}'" for id in keys]),
|
||||
)
|
||||
params = {"workspace": self.db.workspace}
|
||||
@@ -270,9 +271,9 @@ class PGKVStorage(BaseKVStorage):
|
||||
|
||||
################ INSERT METHODS ################
|
||||
async def upsert(self, data: Dict[str, dict]):
|
||||
if self.namespace.endswith("text_chunks"):
|
||||
if is_namespace(self.namespace, NameSpace.KV_STORE_TEXT_CHUNKS):
|
||||
pass
|
||||
elif self.namespace.endswith("full_docs"):
|
||||
elif is_namespace(self.namespace, NameSpace.KV_STORE_FULL_DOCS):
|
||||
for k, v in data.items():
|
||||
upsert_sql = SQL_TEMPLATES["upsert_doc_full"]
|
||||
_data = {
|
||||
@@ -281,7 +282,7 @@ class PGKVStorage(BaseKVStorage):
|
||||
"workspace": self.db.workspace,
|
||||
}
|
||||
await self.db.execute(upsert_sql, _data)
|
||||
elif self.namespace.endswith("llm_response_cache"):
|
||||
elif is_namespace(self.namespace, NameSpace.KV_STORE_LLM_RESPONSE_CACHE):
|
||||
for mode, items in data.items():
|
||||
for k, v in items.items():
|
||||
upsert_sql = SQL_TEMPLATES["upsert_llm_response_cache"]
|
||||
@@ -296,12 +297,11 @@ class PGKVStorage(BaseKVStorage):
|
||||
await self.db.execute(upsert_sql, _data)
|
||||
|
||||
async def index_done_callback(self):
|
||||
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
|
||||
if is_namespace(
|
||||
self.namespace,
|
||||
(NameSpace.KV_STORE_FULL_DOCS, NameSpace.KV_STORE_TEXT_CHUNKS),
|
||||
):
|
||||
logger.info("full doc and chunk data had been saved into postgresql db!")
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -393,11 +393,11 @@ class PGVectorStorage(BaseVectorStorage):
|
||||
for i, d in enumerate(list_data):
|
||||
d["__vector__"] = embeddings[i]
|
||||
for item in list_data:
|
||||
if self.namespace.endswith("chunks"):
|
||||
if is_namespace(self.namespace, NameSpace.VECTOR_STORE_CHUNKS):
|
||||
upsert_sql, data = self._upsert_chunks(item)
|
||||
elif self.namespace.endswith("entities"):
|
||||
elif is_namespace(self.namespace, NameSpace.VECTOR_STORE_ENTITIES):
|
||||
upsert_sql, data = self._upsert_entities(item)
|
||||
elif self.namespace.endswith("relationships"):
|
||||
elif is_namespace(self.namespace, NameSpace.VECTOR_STORE_RELATIONSHIPS):
|
||||
upsert_sql, data = self._upsert_relationships(item)
|
||||
else:
|
||||
raise ValueError(f"{self.namespace} is not supported")
|
||||
@@ -1027,16 +1027,22 @@ class PGGraphStorage(BaseGraphStorage):
|
||||
|
||||
|
||||
NAMESPACE_TABLE_MAP = {
|
||||
"full_docs": "LIGHTRAG_DOC_FULL",
|
||||
"text_chunks": "LIGHTRAG_DOC_CHUNKS",
|
||||
"chunks": "LIGHTRAG_DOC_CHUNKS",
|
||||
"entities": "LIGHTRAG_VDB_ENTITY",
|
||||
"relationships": "LIGHTRAG_VDB_RELATION",
|
||||
"doc_status": "LIGHTRAG_DOC_STATUS",
|
||||
"llm_response_cache": "LIGHTRAG_LLM_CACHE",
|
||||
NameSpace.KV_STORE_FULL_DOCS: "LIGHTRAG_DOC_FULL",
|
||||
NameSpace.KV_STORE_TEXT_CHUNKS: "LIGHTRAG_DOC_CHUNKS",
|
||||
NameSpace.VECTOR_STORE_CHUNKS: "LIGHTRAG_DOC_CHUNKS",
|
||||
NameSpace.VECTOR_STORE_ENTITIES: "LIGHTRAG_VDB_ENTITY",
|
||||
NameSpace.VECTOR_STORE_RELATIONSHIPS: "LIGHTRAG_VDB_RELATION",
|
||||
NameSpace.DOC_STATUS: "LIGHTRAG_DOC_STATUS",
|
||||
NameSpace.KV_STORE_LLM_RESPONSE_CACHE: "LIGHTRAG_LLM_CACHE",
|
||||
}
|
||||
|
||||
|
||||
def namespace_to_table_name(namespace: str) -> str:
|
||||
for k, v in NAMESPACE_TABLE_MAP.items():
|
||||
if is_namespace(namespace, k):
|
||||
return v
|
||||
|
||||
|
||||
TABLES = {
|
||||
"LIGHTRAG_DOC_FULL": {
|
||||
"ddl": """CREATE TABLE LIGHTRAG_DOC_FULL (
|
||||
|
||||
Reference in New Issue
Block a user