move emb_model_config_name to knowledge_config

FredericW · May 15, 2024 · 6186188 · 6186188
1 parent 55c34fa
commit 6186188
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 18 deletions.
diff --git a/.../configs/detailed_rag_config_example.json → ..._RAG_agents/configs/knowledge_config.json b/.../configs/detailed_rag_config_example.json → ..._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
@@ -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
@@ -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
@@ -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 ")