lazy load llamaindex knownledge (#429)

src/agentscope/rag/__init__.py

@@ -1,11 +1,9 @@
 # -*- coding: utf-8 -*-
 """ Import all pipeline related modules in the package. """
 from .knowledge import Knowledge
-from .llama_index_knowledge import LlamaIndexKnowledge
 from .knowledge_bank import KnowledgeBank
 
 __all__ = [
     "Knowledge",
-    "LlamaIndexKnowledge",
     "KnowledgeBank",
 ]

src/agentscope/rag/knowledge_bank.py

@@ -7,8 +7,8 @@
 from typing import Optional, Union
 from loguru import logger
 from agentscope.agents import AgentBase
-from .llama_index_knowledge import LlamaIndexKnowledge
 from ..manager import ModelManager
+from .knowledge import Knowledge
 
 DEFAULT_INDEX_CONFIG = {
     "knowledge_id": "",
@@ -43,13 +43,14 @@ def __init__(
         configs: Union[dict, str],
     ) -> None:
         """initialize the knowledge bank"""
+
         if isinstance(configs, str):
             logger.info(f"Loading configs from {configs}")
             with open(configs, "r", encoding="utf-8") as fp:
                 self.configs = json.loads(fp.read())
         else:
             self.configs = configs
-        self.stored_knowledge: dict[str, LlamaIndexKnowledge] = {}
+        self.stored_knowledge: dict[str, Knowledge] = {}
         self._init_knowledge()
 
     def _init_knowledge(self) -> None:
@@ -104,6 +105,8 @@ def add_data_as_knowledge(
                 )
             ''
         """
+        from .llama_index_knowledge import LlamaIndexKnowledge
+
         if knowledge_id in self.stored_knowledge:
             raise ValueError(f"knowledge_id {knowledge_id} already exists.")
 
@@ -125,17 +128,19 @@ def add_data_as_knowledge(
             knowledge_id=knowledge_id,
             emb_model=model_manager.get_model_by_config_name(emb_model_name),
             knowledge_config=knowledge_config,
-            model=model_manager.get_model_by_config_name(model_name)
-            if model_name
-            else None,
+            model=(
+                model_manager.get_model_by_config_name(model_name)
+                if model_name
+                else None
+            ),
         )
         logger.info(f"data loaded for knowledge_id = {knowledge_id}.")
 
     def get_knowledge(
         self,
         knowledge_id: str,
         duplicate: bool = False,
-    ) -> LlamaIndexKnowledge:
+    ) -> Knowledge:
         """
         Get a Knowledge object from the knowledge bank.
         Args:
@@ -144,7 +149,7 @@ def get_knowledge(
             duplicate (bool):
                 whether return a copy of the Knowledge object.
         Returns:
-            LlamaIndexKnowledge:
+            Knowledge:
                 the Knowledge object defined with Llama-index
         """
         if knowledge_id not in self.stored_knowledge:

tests/knowledge_test.py

@@ -10,7 +10,6 @@
 
 import agentscope
 from agentscope.manager import ASManager
-from agentscope.rag import LlamaIndexKnowledge
 from agentscope.models import OpenAIEmbeddingWrapper, ModelResponse
 
 
@@ -59,6 +58,8 @@ def tearDown(self) -> None:
 
     def test_llamaindexknowledge(self) -> None:
         """test llamaindexknowledge"""
+        from agentscope.rag.llama_index_knowledge import LlamaIndexKnowledge
+
         dummy_model = DummyModel()
 
         knowledge_config = {