Fix doc_status error

This commit is contained in:
yangdx
2025-02-11 10:17:51 +08:00
parent cddde8053d
commit 0e660d4000

View File

@@ -408,17 +408,8 @@ 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=make_namespace(
self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE
),
embedding_func=self.embedding_func,
)
# Initialize document status storage
self.doc_status_storage_cls = self._get_storage_class(self.doc_status_storage)
# Check if Oracle storage implementation is used
if (
@@ -528,7 +519,7 @@ class LightRAG:
self.kv_storage == "PGKVStorage"
or self.vector_storage == "PGVectorStorage"
or self.graph_storage == "PGGraphStorage"
or self.json_doc_status_storage == "PGDocStatusStorage"
or self.doc_status_storage == "PGDocStatusStorage"
):
# Read configuration file
config_parser = configparser.ConfigParser()
@@ -578,12 +569,16 @@ class LightRAG:
self.vector_db_storage_cls.db = postgres_db
if self.graph_storage == "PGGraphStorage":
self.graph_storage_cls.db = postgres_db
if self.json_doc_status_storage == "OracleGraphStorage":
self.json_doc_status_storage = postgres_db
if self.doc_status_storage == "OracleGraphStorage":
self.doc_status_storage_cls = postgres_db
self.llm_response_cache = self.key_string_value_json_storage_cls(
namespace=make_namespace(
self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE
),
embedding_func=self.embedding_func,
)
####
# Add embedding function by walter
####
self.full_docs: BaseKVStorage = self.key_string_value_json_storage_cls(
namespace=make_namespace(
self.namespace_prefix, NameSpace.KV_STORE_FULL_DOCS
@@ -602,9 +597,6 @@ class LightRAG:
),
embedding_func=self.embedding_func,
)
####
# End of adding embedding function by walter
####
self.entities_vdb = self.vector_db_storage_cls(
namespace=make_namespace(
@@ -627,6 +619,7 @@ class LightRAG:
embedding_func=self.embedding_func,
)
# What's for, Is this nessisary ?
if self.llm_response_cache and hasattr(
self.llm_response_cache, "global_config"
):
@@ -639,6 +632,17 @@ class LightRAG:
embedding_func=self.embedding_func,
)
# self.json_doc_status_storage = self.key_string_value_json_storage_cls(
# namespace=self.namespace_prefix + "json_doc_status_storage",
# embedding_func=None,
# )
self.doc_status: DocStatusStorage = self.doc_status_storage_cls(
namespace=make_namespace(self.namespace_prefix, NameSpace.DOC_STATUS),
global_config=global_config,
embedding_func=None,
)
self.llm_model_func = limit_async_func_call(self.llm_model_max_async)(
partial(
self.llm_model_func,
@@ -647,14 +651,6 @@ class LightRAG:
)
)
# Initialize document status storage
self.doc_status_storage_cls = self._get_storage_class(self.doc_status_storage)
self.doc_status: DocStatusStorage = self.doc_status_storage_cls(
namespace=make_namespace(self.namespace_prefix, NameSpace.DOC_STATUS),
global_config=global_config,
embedding_func=None,
)
async def get_graph_labels(self):
text = await self.chunk_entity_relation_graph.get_all_labels()
return text