diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 0eb1b27e..3e15a975 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -40,14 +40,6 @@ from .storage import ( NetworkXStorage, ) -from .kg.neo4j_impl import Neo4JStorage - -from .kg.oracle_impl import OracleKVStorage, OracleGraphStorage, OracleVectorDBStorage - -from .kg.milvus_impl import MilvusVectorDBStorge - -from .kg.mongo_impl import MongoKVStorage - # future KG integrations # from .kg.ArangoDB_impl import ( @@ -55,6 +47,30 @@ from .kg.mongo_impl import MongoKVStorage # ) +def lazy_external_import(module_name: str, class_name: str): + """Lazily import an external module and return a class from it.""" + + def import_class(): + import importlib + + # Import the module using importlib + module = importlib.import_module(module_name) + + # Get the class from the module + return getattr(module, class_name) + + # Return the import_class function itself, not its result + return import_class + + +Neo4JStorage = lazy_external_import(".kg.neo4j_impl", "Neo4JStorage") +OracleKVStorage = lazy_external_import(".kg.oracle_impl", "OracleKVStorage") +OracleGraphStorage = lazy_external_import(".kg.oracle_impl", "OracleGraphStorage") +OracleVectorDBStorage = lazy_external_import(".kg.oracle_impl", "OracleVectorDBStorage") +MilvusVectorDBStorge = lazy_external_import(".kg.milvus_impl", "MilvusVectorDBStorge") +MongoKVStorage = lazy_external_import(".kg.mongo_impl", "MongoKVStorage") + + def always_get_an_event_loop() -> asyncio.AbstractEventLoop: """ Ensure that there is always an event loop available. diff --git a/lightrag/llm.py b/lightrag/llm.py index 568fce04..6a64244a 100644 --- a/lightrag/llm.py +++ b/lightrag/llm.py @@ -919,6 +919,8 @@ class MultiModel: self, prompt, system_prompt=None, history_messages=[], **kwargs ) -> str: kwargs.pop("model", None) # stop from overwriting the custom model name + kwargs.pop("keyword_extraction", None) + kwargs.pop("mode", None) next_model = self._next_model() args = dict( prompt=prompt,