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

fix pre commit

Browse Source
This commit is contained in:
jin
2024-11-12 13:32:40 +08:00
parent cc45d802a6
commit 41599897fb
7 changed files with 345 additions and 310 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
lightrag/lightrag.py
View File

@@ -38,15 +38,11 @@ from .storage import (
JsonKVStorage,
NanoVectorDBStorage,
NetworkXStorage,
)
)
from .kg.neo4j_impl import Neo4JStorage
from .kg.oracle_impl import (
OracleKVStorage,
OracleGraphStorage,
OracleVectorDBStorage
)
from .kg.oracle_impl import OracleKVStorage, OracleGraphStorage, OracleVectorDBStorage
# future KG integrations
@@ -54,6 +50,7 @@ from .kg.oracle_impl import (
# GraphStorage as ArangoDBStorage
# )
def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
try:
return asyncio.get_event_loop()
@@ -72,7 +69,7 @@ class LightRAG:
default_factory=lambda: f"./lightrag_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}"
)
kv_storage : str = field(default="JsonKVStorage")
kv_storage: str = field(default="JsonKVStorage")
vector_storage: str = field(default="NanoVectorDBStorage")
graph_storage: str = field(default="NetworkXStorage")
@@ -115,7 +112,7 @@ class LightRAG:
# storage
vector_db_storage_cls_kwargs: dict = field(default_factory=dict)
enable_llm_cache: bool = True
# extension
@@ -134,18 +131,25 @@ class LightRAG:
# @TODO: should move all storage setup here to leverage initial start params attached to self.
self.key_string_value_json_storage_cls: Type[BaseKVStorage] = self._get_storage_class()[self.kv_storage]
self.vector_db_storage_cls: Type[BaseVectorStorage] = self._get_storage_class()[self.vector_storage]
self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[self.graph_storage]
self.key_string_value_json_storage_cls: Type[BaseKVStorage] = (
self._get_storage_class()[self.kv_storage]
)
self.vector_db_storage_cls: Type[BaseVectorStorage] = self._get_storage_class()[
self.vector_storage
]
self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[
self.graph_storage
]
if not os.path.exists(self.working_dir):
logger.info(f"Creating working directory {self.working_dir}")
os.makedirs(self.working_dir)
self.llm_response_cache = (
self.key_string_value_json_storage_cls(
namespace="llm_response_cache", global_config=asdict(self),embedding_func=None
namespace="llm_response_cache",
global_config=asdict(self),
embedding_func=None,
)
if self.enable_llm_cache
else None
@@ -159,13 +163,19 @@ class LightRAG:
# add embedding func by walter
####
self.full_docs = self.key_string_value_json_storage_cls(
namespace="full_docs", global_config=asdict(self), embedding_func=self.embedding_func
namespace="full_docs",
global_config=asdict(self),
embedding_func=self.embedding_func,
)
self.text_chunks = self.key_string_value_json_storage_cls(
namespace="text_chunks", global_config=asdict(self), embedding_func=self.embedding_func
namespace="text_chunks",
global_config=asdict(self),
embedding_func=self.embedding_func,
)
self.chunk_entity_relation_graph = self.graph_storage_cls(
namespace="chunk_entity_relation", global_config=asdict(self), embedding_func=self.embedding_func
namespace="chunk_entity_relation",
global_config=asdict(self),
embedding_func=self.embedding_func,
)
####
# add embedding func by walter over
@@ -200,13 +210,11 @@ class LightRAG:
def _get_storage_class(self) -> Type[BaseGraphStorage]:
return {
# kv storage
"JsonKVStorage":JsonKVStorage,
"OracleKVStorage":OracleKVStorage,
"JsonKVStorage": JsonKVStorage,
"OracleKVStorage": OracleKVStorage,
# vector storage
"NanoVectorDBStorage":NanoVectorDBStorage,
"OracleVectorDBStorage":OracleVectorDBStorage,
"NanoVectorDBStorage": NanoVectorDBStorage,
"OracleVectorDBStorage": OracleVectorDBStorage,
# graph storage
"NetworkXStorage": NetworkXStorage,
"Neo4JStorage": Neo4JStorage,