diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index fa7ce2a8..8e2ed3a8 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -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