Lightrag/support import existing kg

lazyllm/docs/tools.py

@@ -238,6 +238,7 @@
 Call the function processYml.
 ''')
 
+
 add_english_doc('rag.readers.ReaderBase', '''
 The base class of file readers, which inherits from the ModuleBase base class and has Callable capabilities. Subclasses that inherit from this class only need to implement the _load_data function, and its return parameter type is List[DocNode]. Generally, the input parameters of the _load_data function are file (Path), extra_info(Dict), and fs (AbstractFileSystem).
 

lazyllm/thirdparty/__init__.py

@@ -81,6 +81,6 @@ def __getattribute__(self, __name):
            'deepspeed', 'fire', 'numpy', 'peft', 'torch', 'transformers', 'collie', 'faiss', 'flash_attn', 'google',
            'lightllm', 'vllm', 'ChatTTS', 'wandb', 'funasr', 'sklearn', 'torchvision', 'scipy', 'pymilvus',
            'sentence_transformers', 'gradio', 'chromadb', 'nltk', 'PIL', 'httpx', 'bm25s', 'kubernetes', 'pymongo',
-           'rapidfuzz', 'FlagEmbedding']
+           'rapidfuzz', 'FlagEmbedding', 'nano_vectordb', 'networkx']
 for m in modules:
     vars()[m] = PackageWrapper(m)

lazyllm/tools/rag/__init__.py

@@ -1,4 +1,4 @@
-from .document import Document
+from .document import Document, GraphDocument
 from .retriever import Retriever
 from .rerank import Reranker, register_reranker
 from .transform import SentenceSplitter, LLMParser, NodeTransform, TransformArgs, AdaptiveTransform
@@ -16,6 +16,7 @@
 
 __all__ = [
     "Document",
+    "GraphDocument",
     "Reranker",
     "Retriever",
     "NodeTransform",

lazyllm/tools/rag/doc_impl.py

@@ -263,10 +263,10 @@ def create_node_group(self, name, transform: Union[str, Callable] = None, parent
                                         trans_node=trans_node, num_workers=num_workers, **kwargs)
 
     @classmethod
-    def register_global_reader(cls, pattern: str, func: Optional[Callable] = None):
+    def register_global_reader(cls, pattern: str, func: Optional[Callable] = None) -> Callable:
         if func is not None:
             cls._registered_file_reader[pattern] = func
-            return None
+            return func
 
         def decorator(klass):
             if callable(klass): cls._registered_file_reader[pattern] = klass
@@ -287,9 +287,10 @@ def register_index(self, index_type: str, index_cls: IndexBase, *args, **kwargs)
         else:
             self.index_pending_registrations.append((index_type, index_cls, args, kwargs))
 
-    def add_reader(self, pattern: str, func: Optional[Callable] = None):
+    def add_reader(self, pattern: str, func: Optional[Callable] = None) -> Callable:
         assert callable(func), 'func for reader should be callable'
         self._local_file_reader[pattern] = func
+        return func
 
     def worker(self):
         is_first_run = True