davidlzs commited on
Commit
af3389f
·
1 Parent(s): 45562ea

fix extra kwargs error: keyword_extraction.

Browse files

add lazy_external_load to reduce external lib deps whenever it's not necessary for user.

Files changed (2) hide show
  1. lightrag/lightrag.py +24 -8
  2. lightrag/llm.py +2 -0
lightrag/lightrag.py CHANGED
@@ -40,14 +40,6 @@ from .storage import (
40
  NetworkXStorage,
41
  )
42
 
43
- from .kg.neo4j_impl import Neo4JStorage
44
-
45
- from .kg.oracle_impl import OracleKVStorage, OracleGraphStorage, OracleVectorDBStorage
46
-
47
- from .kg.milvus_impl import MilvusVectorDBStorge
48
-
49
- from .kg.mongo_impl import MongoKVStorage
50
-
51
  # future KG integrations
52
 
53
  # from .kg.ArangoDB_impl import (
@@ -55,6 +47,30 @@ from .kg.mongo_impl import MongoKVStorage
55
  # )
56
 
57
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
58
  def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
59
  """
60
  Ensure that there is always an event loop available.
 
40
  NetworkXStorage,
41
  )
42
 
 
 
 
 
 
 
 
 
43
  # future KG integrations
44
 
45
  # from .kg.ArangoDB_impl import (
 
47
  # )
48
 
49
 
50
+ def lazy_external_import(module_name: str, class_name: str):
51
+ """Lazily import an external module and return a class from it."""
52
+
53
+ def import_class():
54
+ import importlib
55
+
56
+ # Import the module using importlib
57
+ module = importlib.import_module(module_name)
58
+
59
+ # Get the class from the module
60
+ return getattr(module, class_name)
61
+
62
+ # Return the import_class function itself, not its result
63
+ return import_class
64
+
65
+
66
+ Neo4JStorage = lazy_external_import(".kg.neo4j_impl", "Neo4JStorage")
67
+ OracleKVStorage = lazy_external_import(".kg.oracle_impl", "OracleKVStorage")
68
+ OracleGraphStorage = lazy_external_import(".kg.oracle_impl", "OracleGraphStorage")
69
+ OracleVectorDBStorage = lazy_external_import(".kg.oracle_impl", "OracleVectorDBStorage")
70
+ MilvusVectorDBStorge = lazy_external_import(".kg.milvus_impl", "MilvusVectorDBStorge")
71
+ MongoKVStorage = lazy_external_import(".kg.mongo_impl", "MongoKVStorage")
72
+
73
+
74
  def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
75
  """
76
  Ensure that there is always an event loop available.
lightrag/llm.py CHANGED
@@ -1074,6 +1074,8 @@ class MultiModel:
1074
  self, prompt, system_prompt=None, history_messages=[], **kwargs
1075
  ) -> str:
1076
  kwargs.pop("model", None) # stop from overwriting the custom model name
 
 
1077
  next_model = self._next_model()
1078
  args = dict(
1079
  prompt=prompt,
 
1074
  self, prompt, system_prompt=None, history_messages=[], **kwargs
1075
  ) -> str:
1076
  kwargs.pop("model", None) # stop from overwriting the custom model name
1077
+ kwargs.pop("keyword_extraction", None)
1078
+ kwargs.pop("mode", None)
1079
  next_model = self._next_model()
1080
  args = dict(
1081
  prompt=prompt,