gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #729 from ArnoChenFx/add-namespace-prefix

Browse Source 
add namespace prefix to storage namespaces
This commit is contained in:
zrguo
2025-02-08 13:57:20 +08:00
committed by GitHub
parent 48aa80baa2 b2b8cf0aa4
commit e787d92a0c
7 changed files with 72 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
lightrag/api/lightrag_server.py
View File

@@ -40,7 +40,7 @@ from .ollama_api import (
from .ollama_api import ollama_server_infos from .ollama_api import ollama_server_infos
# Load environment variables # Load environment variables
load_dotenv() load_dotenv(override=True)
class RAGStorageConfig: class RAGStorageConfig:
@@ -532,6 +532,14 @@ def parse_args() -> argparse.Namespace:
help="Number of conversation history turns to include (default: from env or 3)", help="Number of conversation history turns to include (default: from env or 3)",
) )
# Namespace
parser.add_argument(
"--namespace-prefix",
type=str,
default=get_env_value("NAMESPACE_PREFIX", ""),
help="Prefix of the namespace",
)
args = parser.parse_args() args = parser.parse_args()
ollama_server_infos.LIGHTRAG_MODEL = args.simulated_model_name ollama_server_infos.LIGHTRAG_MODEL = args.simulated_model_name
@@ -861,6 +869,8 @@ def create_app(args):
"similarity_threshold": 0.95, "similarity_threshold": 0.95,
"use_llm_check": False, "use_llm_check": False,
}, },
log_level=args.log_level,
namespace_prefix=args.namespace_prefix,
) )
else: else:
rag = LightRAG( rag = LightRAG(
@@ -890,6 +900,8 @@ def create_app(args):
"similarity_threshold": 0.95, "similarity_threshold": 0.95,
"use_llm_check": False, "use_llm_check": False,
}, },
log_level=args.log_level,
namespace_prefix=args.namespace_prefix,
) )
async def index_file(file_path: Union[str, Path]) -> None: async def index_file(file_path: Union[str, Path]) -> None:

2
lightrag/api/ollama_api.py
View File

@@ -15,7 +15,7 @@ from dotenv import load_dotenv
# Load environment variables # Load environment variables
load_dotenv() load_dotenv(override=True)
class OllamaServerInfos: class OllamaServerInfos:

4
lightrag/kg/mongo_impl.py
View File

@@ -52,7 +52,7 @@ class MongoKVStorage(BaseKVStorage):
return set([s for s in data if s not in existing_ids]) return set([s for s in data if s not in existing_ids])
async def upsert(self, data: dict[str, dict]): 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 mode, items in data.items():
for k, v in tqdm_async(items.items(), desc="Upserting"): for k, v in tqdm_async(items.items(), desc="Upserting"):
key = f"{mode}_{k}" key = f"{mode}_{k}"
@@ -69,7 +69,7 @@ class MongoKVStorage(BaseKVStorage):
return data return data
async def get_by_mode_and_id(self, mode: str, id: str) -> Union[dict, None]: 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 = {} res = {}
v = self._data.find_one({"_id": mode + "_" + id}) v = self._data.find_one({"_id": mode + "_" + id})
if v: if v:

16
lightrag/kg/oracle_impl.py
View File

@@ -185,7 +185,7 @@ class OracleKVStorage(BaseKVStorage):
SQL = SQL_TEMPLATES["get_by_id_" + self.namespace] 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) # 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) array_res = await self.db.query(SQL, params, multirows=True)
res = {} res = {}
for row in array_res: for row in array_res:
@@ -201,7 +201,7 @@ class OracleKVStorage(BaseKVStorage):
"""Specifically for llm_response_cache.""" """Specifically for llm_response_cache."""
SQL = SQL_TEMPLATES["get_by_mode_id_" + self.namespace] SQL = SQL_TEMPLATES["get_by_mode_id_" + self.namespace]
params = {"workspace": self.db.workspace, "cache_mode": mode, "id": id} 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) array_res = await self.db.query(SQL, params, multirows=True)
res = {} res = {}
for row in array_res: for row in array_res:
@@ -218,7 +218,7 @@ class OracleKVStorage(BaseKVStorage):
params = {"workspace": self.db.workspace} params = {"workspace": self.db.workspace}
# print("get_by_ids:"+SQL) # print("get_by_ids:"+SQL)
res = await self.db.query(SQL, params, multirows=True) res = await self.db.query(SQL, params, multirows=True)
if "llm_response_cache" == self.namespace: if self.namespace.endswith("llm_response_cache"):
modes = set() modes = set()
dict_res: dict[str, dict] = {} dict_res: dict[str, dict] = {}
for row in res: for row in res:
@@ -269,7 +269,7 @@ class OracleKVStorage(BaseKVStorage):
################ INSERT METHODS ################ ################ INSERT METHODS ################
async def upsert(self, data: dict[str, dict]): async def upsert(self, data: dict[str, dict]):
if self.namespace == "text_chunks": if self.namespace.endswith("text_chunks"):
list_data = [ list_data = [
{ {
"id": k, "id": k,
@@ -302,7 +302,7 @@ class OracleKVStorage(BaseKVStorage):
"status": item["status"], "status": item["status"],
} }
await self.db.execute(merge_sql, _data) await self.db.execute(merge_sql, _data)
if self.namespace == "full_docs": if self.namespace.endswith("full_docs"):
for k, v in data.items(): for k, v in data.items():
# values.clear() # values.clear()
merge_sql = SQL_TEMPLATES["merge_doc_full"] merge_sql = SQL_TEMPLATES["merge_doc_full"]
@@ -313,7 +313,7 @@ class OracleKVStorage(BaseKVStorage):
} }
await self.db.execute(merge_sql, _data) 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 mode, items in data.items():
for k, v in items.items(): for k, v in items.items():
upsert_sql = SQL_TEMPLATES["upsert_llm_response_cache"] upsert_sql = SQL_TEMPLATES["upsert_llm_response_cache"]
@@ -334,8 +334,10 @@ class OracleKVStorage(BaseKVStorage):
await self.db.execute(SQL, params) await self.db.execute(SQL, params)
async def index_done_callback(self): async def index_done_callback(self):
if self.namespace in ["full_docs", "text_chunks"]: 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!") logger.info("full doc and chunk data had been saved into oracle db!")
break
@dataclass @dataclass

28
lightrag/kg/postgres_impl.py
View File

@@ -187,7 +187,7 @@ class PGKVStorage(BaseKVStorage):
"""Get doc_full data by id.""" """Get doc_full data by id."""
sql = SQL_TEMPLATES["get_by_id_" + self.namespace] sql = SQL_TEMPLATES["get_by_id_" + self.namespace]
params = {"workspace": self.db.workspace, "id": id} 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) array_res = await self.db.query(sql, params, multirows=True)
res = {} res = {}
for row in array_res: for row in array_res:
@@ -203,7 +203,7 @@ class PGKVStorage(BaseKVStorage):
"""Specifically for llm_response_cache.""" """Specifically for llm_response_cache."""
sql = SQL_TEMPLATES["get_by_mode_id_" + self.namespace] sql = SQL_TEMPLATES["get_by_mode_id_" + self.namespace]
params = {"workspace": self.db.workspace, mode: mode, "id": id} 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) array_res = await self.db.query(sql, params, multirows=True)
res = {} res = {}
for row in array_res: for row in array_res:
@@ -219,7 +219,7 @@ class PGKVStorage(BaseKVStorage):
ids=",".join([f"'{id}'" for id in ids]) ids=",".join([f"'{id}'" for id in ids])
) )
params = {"workspace": self.db.workspace} 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) array_res = await self.db.query(sql, params, multirows=True)
modes = set() modes = set()
dict_res: dict[str, dict] = {} dict_res: dict[str, dict] = {}
@@ -239,7 +239,7 @@ class PGKVStorage(BaseKVStorage):
return None return None
async def all_keys(self) -> list[dict]: 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" sql = "select workspace,mode,id from lightrag_llm_cache"
res = await self.db.query(sql, multirows=True) res = await self.db.query(sql, multirows=True)
return res return res
@@ -270,9 +270,9 @@ class PGKVStorage(BaseKVStorage):
################ INSERT METHODS ################ ################ INSERT METHODS ################
async def upsert(self, data: Dict[str, dict]): async def upsert(self, data: Dict[str, dict]):
if self.namespace == "text_chunks": if self.namespace.endswith("text_chunks"):
pass pass
elif self.namespace == "full_docs": elif self.namespace.endswith("full_docs"):
for k, v in data.items(): for k, v in data.items():
upsert_sql = SQL_TEMPLATES["upsert_doc_full"] upsert_sql = SQL_TEMPLATES["upsert_doc_full"]
_data = { _data = {
@@ -281,7 +281,7 @@ class PGKVStorage(BaseKVStorage):
"workspace": self.db.workspace, "workspace": self.db.workspace,
} }
await self.db.execute(upsert_sql, _data) 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 mode, items in data.items():
for k, v in items.items(): for k, v in items.items():
upsert_sql = SQL_TEMPLATES["upsert_llm_response_cache"] upsert_sql = SQL_TEMPLATES["upsert_llm_response_cache"]
@@ -296,8 +296,12 @@ class PGKVStorage(BaseKVStorage):
await self.db.execute(upsert_sql, _data) await self.db.execute(upsert_sql, _data)
async def index_done_callback(self): async def index_done_callback(self):
if self.namespace in ["full_docs", "text_chunks"]: for n in ("full_docs", "text_chunks"):
logger.info("full doc and chunk data had been saved into postgresql db!") if self.namespace.endswith(n):
logger.info(
"full doc and chunk data had been saved into postgresql db!"
)
break
@dataclass @dataclass
@@ -389,11 +393,11 @@ class PGVectorStorage(BaseVectorStorage):
for i, d in enumerate(list_data): for i, d in enumerate(list_data):
d["__vector__"] = embeddings[i] d["__vector__"] = embeddings[i]
for item in list_data: for item in list_data:
if self.namespace == "chunks": if self.namespace.endswith("chunks"):
upsert_sql, data = self._upsert_chunks(item) upsert_sql, data = self._upsert_chunks(item)
elif self.namespace == "entities": elif self.namespace.endswith("entities"):
upsert_sql, data = self._upsert_entities(item) upsert_sql, data = self._upsert_entities(item)
elif self.namespace == "relationships": elif self.namespace.endswith("relationships"):
upsert_sql, data = self._upsert_relationships(item) upsert_sql, data = self._upsert_relationships(item)
else: else:
raise ValueError(f"{self.namespace} is not supported") raise ValueError(f"{self.namespace} is not supported")

20
lightrag/kg/tidb_impl.py
View File

@@ -160,7 +160,7 @@ class TiDBKVStorage(BaseKVStorage):
async def upsert(self, data: dict[str, dict]): async def upsert(self, data: dict[str, dict]):
left_data = {k: v for k, v in data.items() if k not in self._data} left_data = {k: v for k, v in data.items() if k not in self._data}
self._data.update(left_data) self._data.update(left_data)
if self.namespace == "text_chunks": if self.namespace.endswith("text_chunks"):
list_data = [ list_data = [
{ {
"__id__": k, "__id__": k,
@@ -190,13 +190,13 @@ class TiDBKVStorage(BaseKVStorage):
"tokens": item["tokens"], "tokens": item["tokens"],
"chunk_order_index": item["chunk_order_index"], "chunk_order_index": item["chunk_order_index"],
"full_doc_id": item["full_doc_id"], "full_doc_id": item["full_doc_id"],
"content_vector": f"{item["__vector__"].tolist()}", "content_vector": f"{item['__vector__'].tolist()}",
"workspace": self.db.workspace, "workspace": self.db.workspace,
} }
) )
await self.db.execute(merge_sql, data) 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"] merge_sql = SQL_TEMPLATES["upsert_doc_full"]
data = [] data = []
for k, v in self._data.items(): for k, v in self._data.items():
@@ -211,8 +211,10 @@ class TiDBKVStorage(BaseKVStorage):
return left_data return left_data
async def index_done_callback(self): async def index_done_callback(self):
if self.namespace in ["full_docs", "text_chunks"]: 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!") logger.info("full doc and chunk data had been saved into TiDB db!")
break
@dataclass @dataclass
@@ -258,7 +260,7 @@ class TiDBVectorDBStorage(BaseVectorStorage):
if not len(data): if not len(data):
logger.warning("You insert an empty data to vector DB") logger.warning("You insert an empty data to vector DB")
return [] return []
if self.namespace == "chunks": if self.namespace.endswith("chunks"):
return [] return []
logger.info(f"Inserting {len(data)} vectors to {self.namespace}") logger.info(f"Inserting {len(data)} vectors to {self.namespace}")
@@ -288,14 +290,14 @@ class TiDBVectorDBStorage(BaseVectorStorage):
for i, d in enumerate(list_data): for i, d in enumerate(list_data):
d["content_vector"] = embeddings[i] d["content_vector"] = embeddings[i]
if self.namespace == "entities": if self.namespace.endswith("entities"):
data = [] data = []
for item in list_data: for item in list_data:
param = { param = {
"id": item["id"], "id": item["id"],
"name": item["entity_name"], "name": item["entity_name"],
"content": item["content"], "content": item["content"],
"content_vector": f"{item["content_vector"].tolist()}", "content_vector": f"{item['content_vector'].tolist()}",
"workspace": self.db.workspace, "workspace": self.db.workspace,
} }
# update entity_id if node inserted by graph_storage_instance before # update entity_id if node inserted by graph_storage_instance before
@@ -309,7 +311,7 @@ class TiDBVectorDBStorage(BaseVectorStorage):
merge_sql = SQL_TEMPLATES["insert_entity"] merge_sql = SQL_TEMPLATES["insert_entity"]
await self.db.execute(merge_sql, data) await self.db.execute(merge_sql, data)
elif self.namespace == "relationships": elif self.namespace.endswith("relationships"):
data = [] data = []
for item in list_data: for item in list_data:
param = { param = {
@@ -317,7 +319,7 @@ class TiDBVectorDBStorage(BaseVectorStorage):
"source_name": item["src_id"], "source_name": item["src_id"],
"target_name": item["tgt_id"], "target_name": item["tgt_id"],
"content": item["content"], "content": item["content"],
"content_vector": f"{item["content_vector"].tolist()}", "content_vector": f"{item['content_vector'].tolist()}",
"workspace": self.db.workspace, "workspace": self.db.workspace,
} }
# update relation_id if node inserted by graph_storage_instance before # update relation_id if node inserted by graph_storage_instance before

39
lightrag/lightrag.py
View File

@@ -167,6 +167,7 @@ class LightRAG:
# storage # storage
vector_db_storage_cls_kwargs: dict = field(default_factory=dict) vector_db_storage_cls_kwargs: dict = field(default_factory=dict)
namespace_prefix: str = field(default="")
enable_llm_cache: bool = True enable_llm_cache: bool = True
# Sometimes there are some reason the LLM failed at Extracting Entities, and we want to continue without LLM cost, we can use this flag # Sometimes there are some reason the LLM failed at Extracting Entities, and we want to continue without LLM cost, we can use this flag
@@ -227,13 +228,8 @@ class LightRAG:
self.graph_storage_cls, global_config=global_config self.graph_storage_cls, global_config=global_config
) )
self.json_doc_status_storage = self.key_string_value_json_storage_cls(
namespace="json_doc_status_storage",
embedding_func=None,
)
self.llm_response_cache = self.key_string_value_json_storage_cls( self.llm_response_cache = self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
) )
@@ -241,33 +237,34 @@ class LightRAG:
# add embedding func by walter # add embedding func by walter
#### ####
self.full_docs = self.key_string_value_json_storage_cls( self.full_docs = self.key_string_value_json_storage_cls(
namespace="full_docs", namespace=self.namespace_prefix + "full_docs",
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
) )
self.text_chunks = self.key_string_value_json_storage_cls( self.text_chunks = self.key_string_value_json_storage_cls(
namespace="text_chunks", namespace=self.namespace_prefix + "text_chunks",
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
) )
self.chunk_entity_relation_graph = self.graph_storage_cls( self.chunk_entity_relation_graph = self.graph_storage_cls(
namespace="chunk_entity_relation", namespace=self.namespace_prefix + "chunk_entity_relation",
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
) )
#### ####
# add embedding func by walter over # add embedding func by walter over
#### ####
self.entities_vdb = self.vector_db_storage_cls( self.entities_vdb = self.vector_db_storage_cls(
namespace="entities", namespace=self.namespace_prefix + "entities",
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
meta_fields={"entity_name"}, meta_fields={"entity_name"},
) )
self.relationships_vdb = self.vector_db_storage_cls( self.relationships_vdb = self.vector_db_storage_cls(
namespace="relationships", namespace=self.namespace_prefix + "relationships",
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
meta_fields={"src_id", "tgt_id"}, meta_fields={"src_id", "tgt_id"},
) )
self.chunks_vdb = self.vector_db_storage_cls( self.chunks_vdb = self.vector_db_storage_cls(
namespace="chunks", namespace=self.namespace_prefix + "chunks",
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
) )
@@ -277,7 +274,7 @@ class LightRAG:
hashing_kv = self.llm_response_cache hashing_kv = self.llm_response_cache
else: else:
hashing_kv = self.key_string_value_json_storage_cls( hashing_kv = self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
) )
@@ -292,7 +289,7 @@ class LightRAG:
# Initialize document status storage # Initialize document status storage
self.doc_status_storage_cls = self._get_storage_class(self.doc_status_storage) self.doc_status_storage_cls = self._get_storage_class(self.doc_status_storage)
self.doc_status = self.doc_status_storage_cls( self.doc_status = self.doc_status_storage_cls(
namespace="doc_status", namespace=self.namespace_prefix + "doc_status",
global_config=global_config, global_config=global_config,
embedding_func=None, embedding_func=None,
) )
@@ -928,7 +925,7 @@ class LightRAG:
if self.llm_response_cache if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config") and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls( else self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
global_config=asdict(self), global_config=asdict(self),
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
), ),
@@ -945,7 +942,7 @@ class LightRAG:
if self.llm_response_cache if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config") and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls( else self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
global_config=asdict(self), global_config=asdict(self),
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
), ),
@@ -964,7 +961,7 @@ class LightRAG:
if self.llm_response_cache if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config") and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls( else self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
global_config=asdict(self), global_config=asdict(self),
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
), ),
@@ -1005,7 +1002,7 @@ class LightRAG:
global_config=asdict(self), global_config=asdict(self),
hashing_kv=self.llm_response_cache hashing_kv=self.llm_response_cache
or self.key_string_value_json_storage_cls( or self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
global_config=asdict(self), global_config=asdict(self),
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
), ),
@@ -1036,7 +1033,7 @@ class LightRAG:
if self.llm_response_cache if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config") and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls( else self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
global_config=asdict(self), global_config=asdict(self),
embedding_func=self.embedding_funcne, embedding_func=self.embedding_funcne,
), ),
@@ -1052,7 +1049,7 @@ class LightRAG:
if self.llm_response_cache if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config") and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls( else self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
global_config=asdict(self), global_config=asdict(self),
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
), ),
@@ -1071,7 +1068,7 @@ class LightRAG:
if self.llm_response_cache if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config") and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls( else self.key_string_value_json_storage_cls(
namespace="llm_response_cache", namespace=self.namespace_prefix + "llm_response_cache",
global_config=asdict(self), global_config=asdict(self),
embedding_func=self.embedding_func, embedding_func=self.embedding_func,
), ),