Skip to content

Commit

Permalink
fix extra kwargs error: keyword_extraction.
Browse files Browse the repository at this point in the history
add lazy_external_load to reduce external lib deps whenever it's not necessary for user.
  • Loading branch information
davidleon committed Dec 9, 2024
1 parent 97d1894 commit eb2508a
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 8 deletions.
31 changes: 23 additions & 8 deletions lightrag/lightrag.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,20 +40,35 @@
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 (
# GraphStorage as ArangoDBStorage
# )

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:
"""
Expand Down
2 changes: 2 additions & 0 deletions lightrag/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -1074,6 +1074,8 @@ async def llm_model_func(
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,
Expand Down

0 comments on commit eb2508a

Please sign in to comment.