diff --git a/src/agentscope/rag/__init__.py b/src/agentscope/rag/__init__.py index 362f1de14..31f035615 100644 --- a/src/agentscope/rag/__init__.py +++ b/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", ] diff --git a/src/agentscope/rag/knowledge_bank.py b/src/agentscope/rag/knowledge_bank.py index 8f07d12b6..ae4cc57ce 100644 --- a/src/agentscope/rag/knowledge_bank.py +++ b/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,9 +128,11 @@ 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}.") @@ -135,7 +140,7 @@ 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: diff --git a/tests/knowledge_test.py b/tests/knowledge_test.py index dde7877bf..1fae3ed01 100644 --- a/tests/knowledge_test.py +++ b/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 = {