From 3b302ae9eee2cc89c4402dc316c5f25ed4e382a6 Mon Sep 17 00:00:00 2001
From: "weirui.kwr@alibaba-inc.com" <weirui.kwr@alibaba-inc.com>
Date: Wed, 18 Sep 2024 19:55:56 +0800
Subject: [PATCH] fix

---
 src/agentscope/message/msg.py        |  2 +-
 src/agentscope/rag/__init__.py       |  2 --
 src/agentscope/rag/knowledge_bank.py | 18 +++++++++++-------
 tests/knowledge_test.py              |  3 ++-
 4 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/agentscope/message/msg.py b/src/agentscope/message/msg.py
index 1f3e99dd3..284ed724e 100644
--- a/src/agentscope/message/msg.py
+++ b/src/agentscope/message/msg.py
@@ -47,7 +47,7 @@ def __init__(
         self,
         name: str,
         content: Any,
-        role: Union[str, Literal["system", "user", "assistant"]],
+        role: Union[str, Literal["system", "user", "assistant"]] = "system",
         url: Optional[Union[str, List[str]]] = None,
         metadata: Optional[Union[dict, str]] = None,
         echo: bool = False,
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..69de842f1 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": "",
@@ -49,7 +49,7 @@ def __init__(
                 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 +104,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 +127,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 +139,7 @@ def get_knowledge(
         self,
         knowledge_id: str,
         duplicate: bool = False,
-    ) -> LlamaIndexKnowledge:
+    ) -> Knowledge:
         """
         Get a Knowledge object from the knowledge bank.
         Args:
@@ -144,7 +148,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 = {