fix event loop conflict
This commit is contained in:
@@ -1,6 +1,5 @@
|
||||
import asyncio
|
||||
import os
|
||||
import importlib
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from functools import partial
|
||||
@@ -24,18 +23,15 @@ from .storage import (
|
||||
NanoVectorDBStorage,
|
||||
NetworkXStorage,
|
||||
)
|
||||
|
||||
from .kg.neo4j_impl import (
|
||||
Neo4JStorage
|
||||
)
|
||||
#future KG integrations
|
||||
|
||||
from .kg.neo4j_impl import Neo4JStorage
|
||||
# future KG integrations
|
||||
|
||||
# from .kg.ArangoDB_impl import (
|
||||
# GraphStorage as ArangoDBStorage
|
||||
# )
|
||||
|
||||
|
||||
|
||||
from .utils import (
|
||||
EmbeddingFunc,
|
||||
compute_mdhash_id,
|
||||
@@ -52,6 +48,7 @@ from .base import (
|
||||
QueryParam,
|
||||
)
|
||||
|
||||
|
||||
def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
@@ -64,7 +61,6 @@ def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
|
||||
|
||||
@dataclass
|
||||
class LightRAG:
|
||||
|
||||
working_dir: str = field(
|
||||
default_factory=lambda: f"./lightrag_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}"
|
||||
)
|
||||
@@ -74,8 +70,6 @@ class LightRAG:
|
||||
current_log_level = logger.level
|
||||
log_level: str = field(default=current_log_level)
|
||||
|
||||
|
||||
|
||||
# text chunking
|
||||
chunk_token_size: int = 1200
|
||||
chunk_overlap_token_size: int = 100
|
||||
@@ -130,8 +124,10 @@ class LightRAG:
|
||||
_print_config = ",\n ".join([f"{k} = {v}" for k, v in asdict(self).items()])
|
||||
logger.debug(f"LightRAG init with param:\n {_print_config}\n")
|
||||
|
||||
#@TODO: should move all storage setup here to leverage initial start params attached to self.
|
||||
self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[self.kg]
|
||||
# @TODO: should move all storage setup here to leverage initial start params attached to self.
|
||||
self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[
|
||||
self.kg
|
||||
]
|
||||
|
||||
if not os.path.exists(self.working_dir):
|
||||
logger.info(f"Creating working directory {self.working_dir}")
|
||||
@@ -185,6 +181,7 @@ class LightRAG:
|
||||
**self.llm_model_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
def _get_storage_class(self) -> Type[BaseGraphStorage]:
|
||||
return {
|
||||
"Neo4JStorage": Neo4JStorage,
|
||||
@@ -328,4 +325,4 @@ class LightRAG:
|
||||
if storage_inst is None:
|
||||
continue
|
||||
tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
|
||||
await asyncio.gather(*tasks)
|
||||
await asyncio.gather(*tasks)
|
||||
|
Reference in New Issue
Block a user