better handling of namespace

This commit is contained in:
ArnoChen
2025-02-08 16:05:59 +08:00
parent e787d92a0c
commit 3f845e9e53
8 changed files with 156 additions and 93 deletions

View File

@@ -35,6 +35,8 @@ from .base import (
DocStatus,
)
from .namespace import NameSpace, make_namespace
from .prompt import GRAPH_FIELD_SEP
STORAGES = {
@@ -228,8 +230,13 @@ class LightRAG:
self.graph_storage_cls, global_config=global_config
)
self.json_doc_status_storage = self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "json_doc_status_storage",
embedding_func=None,
)
self.llm_response_cache = self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
embedding_func=self.embedding_func,
)
@@ -237,34 +244,33 @@ class LightRAG:
# add embedding func by walter
####
self.full_docs = self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "full_docs",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_FULL_DOCS),
embedding_func=self.embedding_func,
)
self.text_chunks = self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "text_chunks",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_TEXT_CHUNKS),
embedding_func=self.embedding_func,
)
self.chunk_entity_relation_graph = self.graph_storage_cls(
namespace=self.namespace_prefix + "chunk_entity_relation",
namespace=make_namespace(self.namespace_prefix, NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION),
embedding_func=self.embedding_func,
)
####
# add embedding func by walter over
####
self.entities_vdb = self.vector_db_storage_cls(
namespace=self.namespace_prefix + "entities",
namespace=make_namespace(self.namespace_prefix, NameSpace.VECTOR_STORE_ENTITIES),
embedding_func=self.embedding_func,
meta_fields={"entity_name"},
)
self.relationships_vdb = self.vector_db_storage_cls(
namespace=self.namespace_prefix + "relationships",
namespace=make_namespace(self.namespace_prefix, NameSpace.VECTOR_STORE_RELATIONSHIPS),
embedding_func=self.embedding_func,
meta_fields={"src_id", "tgt_id"},
)
self.chunks_vdb = self.vector_db_storage_cls(
namespace=self.namespace_prefix + "chunks",
namespace=make_namespace(self.namespace_prefix, NameSpace.VECTOR_STORE_CHUNKS),
embedding_func=self.embedding_func,
)
@@ -274,7 +280,7 @@ class LightRAG:
hashing_kv = self.llm_response_cache
else:
hashing_kv = self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
embedding_func=self.embedding_func,
)
@@ -289,7 +295,7 @@ class LightRAG:
# Initialize document status storage
self.doc_status_storage_cls = self._get_storage_class(self.doc_status_storage)
self.doc_status = self.doc_status_storage_cls(
namespace=self.namespace_prefix + "doc_status",
namespace=make_namespace(self.namespace_prefix, NameSpace.DOC_STATUS),
global_config=global_config,
embedding_func=None,
)
@@ -925,7 +931,7 @@ class LightRAG:
if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
global_config=asdict(self),
embedding_func=self.embedding_func,
),
@@ -942,7 +948,7 @@ class LightRAG:
if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
global_config=asdict(self),
embedding_func=self.embedding_func,
),
@@ -961,7 +967,7 @@ class LightRAG:
if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
global_config=asdict(self),
embedding_func=self.embedding_func,
),
@@ -1002,7 +1008,7 @@ class LightRAG:
global_config=asdict(self),
hashing_kv=self.llm_response_cache
or self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
global_config=asdict(self),
embedding_func=self.embedding_func,
),
@@ -1033,7 +1039,7 @@ class LightRAG:
if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
global_config=asdict(self),
embedding_func=self.embedding_funcne,
),
@@ -1049,7 +1055,7 @@ class LightRAG:
if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
global_config=asdict(self),
embedding_func=self.embedding_func,
),
@@ -1068,7 +1074,7 @@ class LightRAG:
if self.llm_response_cache
and hasattr(self.llm_response_cache, "global_config")
else self.key_string_value_json_storage_cls(
namespace=self.namespace_prefix + "llm_response_cache",
namespace=make_namespace(self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE),
global_config=asdict(self),
embedding_func=self.embedding_func,
),