From 618618876b52182aa60a820f7641048214da9d35 Mon Sep 17 00:00:00 2001
From: ZiTao-Li <zitao.l@alibaba-inc.com>
Date: Wed, 15 May 2024 10:31:54 +0800
Subject: [PATCH] move emb_model_config_name to knowledge_config

---
 ...fig_example.json => knowledge_config.json} |  3 +++
 .../rag_example.py                            |  2 +-
 src/agentscope/agents/rag_agents.py           | 24 +++++++------------
 src/agentscope/rag/knowledge_bank.py          |  2 +-
 4 files changed, 13 insertions(+), 18 deletions(-)
 rename examples/conversation_with_RAG_agents/configs/{detailed_rag_config_example.json => knowledge_config.json} (94%)

diff --git a/examples/conversation_with_RAG_agents/configs/detailed_rag_config_example.json b/examples/conversation_with_RAG_agents/configs/knowledge_config.json
similarity index 94%
rename from examples/conversation_with_RAG_agents/configs/detailed_rag_config_example.json
rename to examples/conversation_with_RAG_agents/configs/knowledge_config.json
index d4c9604f4..0935c3dc5 100644
--- a/examples/conversation_with_RAG_agents/configs/detailed_rag_config_example.json
+++ b/examples/conversation_with_RAG_agents/configs/knowledge_config.json
@@ -1,6 +1,7 @@
 [
   {
     "knowledge_id": "agentscope_code_rag",
+    "emb_model_config_name": "qwen_emb_config",
     "chunk_size": 2048,
     "chunk_overlap": 40,
     "data_processing": [
@@ -37,6 +38,7 @@
   },
   {
     "knowledge_id": "agentscope_api_rag",
+    "emb_model_config_name": "qwen_emb_config",
     "chunk_size": 2048,
     "chunk_overlap": 40,
     "data_processing": [
@@ -59,6 +61,7 @@
   },
   {
     "knowledge_id": "agentscope_global_rag",
+    "emb_model_config_name": "qwen_emb_config",
     "chunk_size": 2048,
     "chunk_overlap": 40,
     "data_processing": [
diff --git a/examples/conversation_with_RAG_agents/rag_example.py b/examples/conversation_with_RAG_agents/rag_example.py
index 1f706190c..85bff1299 100644
--- a/examples/conversation_with_RAG_agents/rag_example.py
+++ b/examples/conversation_with_RAG_agents/rag_example.py
@@ -67,7 +67,7 @@ def main() -> None:
 
     # the knowledge bank can be configured by loading config file
     with open(
-        "configs/detailed_rag_config_example.json",
+        "configs/knowledge_config.json",
         "r",
         encoding="utf-8",
     ) as f:
diff --git a/src/agentscope/agents/rag_agents.py b/src/agentscope/agents/rag_agents.py
index 4e4b1d898..4928a59cc 100644
--- a/src/agentscope/agents/rag_agents.py
+++ b/src/agentscope/agents/rag_agents.py
@@ -12,10 +12,17 @@
 
 from agentscope.agents.agent import AgentBase
 from agentscope.message import Msg
-from agentscope.models import load_model_by_config_name
 from agentscope.rag import KnowledgeBank
 
 
+CHECKING_PROMPT = """
+                Does the retrieved content is relevant to the query?
+                Retrieved content: {}
+                Query: {}
+                Only answer YES or NO.
+                """
+
+
 class RAGAgentBase(AgentBase, ABC):
     """
     Base class for RAG agents
@@ -26,7 +33,6 @@ def __init__(
         name: str,
         sys_prompt: str,
         model_config_name: str,
-        emb_model_config_name: str,
         memory_config: Optional[dict] = None,
         rag_config: Optional[dict] = None,
     ) -> None:
@@ -39,8 +45,6 @@ def __init__(
                 system prompt for the RAG agent.
             model_config_name (str):
                 language model for the agent.
-            emb_model_config_name (str):
-                embedding model for the agent.
             memory_config (dict):
                 memory configuration.
             rag_config (dict):
@@ -56,8 +60,6 @@ def __init__(
             use_memory=True,
             memory_config=memory_config,
         )
-        # setup embedding model used in RAG
-        self.emb_model = load_model_by_config_name(emb_model_config_name)
 
         # setup RAG configurations
         self.rag_config = rag_config or {}
@@ -176,7 +178,6 @@ def __init__(
         sys_prompt: str,
         model_config_name: str,
         knowledge_bank: Optional[KnowledgeBank],
-        emb_model_config_name: str = None,
         memory_config: Optional[dict] = None,
         rag_config: Optional[dict] = None,
         **kwargs: Any,
@@ -190,8 +191,6 @@ def __init__(
                 system prompt for the RAG agent
             model_config_name (str):
                 language model for the agent
-            emb_model_config_name (str):
-                embedding model for the agent
             memory_config (dict):
                 memory configuration
             rag_config (dict):
@@ -211,7 +210,6 @@ def __init__(
             name=name,
             sys_prompt=sys_prompt,
             model_config_name=model_config_name,
-            emb_model_config_name=emb_model_config_name,
             memory_config=memory_config,
             rag_config=rag_config,
         )
@@ -314,12 +312,6 @@ def reply(self, x: dict = None) -> dict:
                 # if the max score is lower than 0.4, then we let LLM
                 # decide whether the retrieved content is relevant
                 # to the user input.
-                CHECKING_PROMPT = """
-                Does the retrieved content is relevant to the query?
-                Retrieved content: {}
-                Query: {}
-                Only answer YES or NO.
-                """
                 msg = Msg(
                     name="user",
                     role="user",
diff --git a/src/agentscope/rag/knowledge_bank.py b/src/agentscope/rag/knowledge_bank.py
index 1deffa139..aae9e4279 100644
--- a/src/agentscope/rag/knowledge_bank.py
+++ b/src/agentscope/rag/knowledge_bank.py
@@ -50,7 +50,7 @@ def _init_knowledge(self) -> None:
         for config in self.configs:
             self.add_data_for_rag(
                 knowledge_id=config["knowledge_id"],
-                emb_model_name="qwen_emb_config",
+                emb_model_name=config["emb_model_config_name"],
                 index_config=config,
             )
         logger.info("knowledge bank initialization completed.\n ")