update tool protocols

braisedpork1964 · Dec 4, 2024 · 22d89eb · 22d89eb
1 parent ca08c0d
commit 22d89eb
Showing 1 changed file with 24 additions and 28 deletions.
diff --git a/mindsearch/agent/__init__.py b/mindsearch/agent/__init__.py
@@ -25,62 +25,58 @@
 LLM = {}
 
 
-def init_agent(
-    lang="cn", model_format="internlm_server", search_engine="BingSearch", use_async=False
-):
+def init_agent(lang="cn", model_format="internlm_server", search_engine="BingSearch", use_async=False):
     mode = "async" if use_async else "sync"
     llm = LLM.get(model_format, {}).get(mode)
     if llm is None:
         llm_cfg = deepcopy(getattr(llm_factory, model_format))
         if llm_cfg is None:
             raise NotImplementedError
         if use_async:
-            cls_name = (
-                llm_cfg["type"].split(".")[-1]
-                if isinstance(llm_cfg["type"], str)
-                else llm_cfg["type"].__name__
-            )
+            cls_name = llm_cfg["type"].split(".")[-1] if isinstance(llm_cfg["type"], str) else llm_cfg["type"].__name__
             llm_cfg["type"] = f"lagent.llms.Async{cls_name}"
         llm = create_object(llm_cfg)
         LLM.setdefault(model_format, {}).setdefault(mode, llm)
 
     date = datetime.now().strftime("The current date is %Y-%m-%d.")
     plugins = [
-        dict(
-            type=AsyncWebBrowser if use_async else WebBrowser,
-            searcher_type=search_engine,
-            topk=6,
-            secret_id=os.getenv("TENCENT_SEARCH_SECRET_ID"),
-            secret_key=os.getenv("TENCENT_SEARCH_SECRET_KEY"),
-        )
-        if search_engine == "TencentSearch"
-        else dict(
-            type=AsyncWebBrowser if use_async else WebBrowser,
-            searcher_type=search_engine,
-            topk=6,
-            api_key=os.getenv("WEB_SEARCH_API_KEY"),
+        (
+            dict(
+                type=AsyncWebBrowser if use_async else WebBrowser,
+                searcher_type=search_engine,
+                topk=6,
+                secret_id=os.getenv("TENCENT_SEARCH_SECRET_ID"),
+                secret_key=os.getenv("TENCENT_SEARCH_SECRET_KEY"),
+            )
+            if search_engine == "TencentSearch"
+            else dict(
+                type=AsyncWebBrowser if use_async else WebBrowser,
+                searcher_type=search_engine,
+                topk=6,
+                api_key=os.getenv("WEB_SEARCH_API_KEY"),
+            )
         )
     ]
     agent = (AsyncMindSearchAgent if use_async else MindSearchAgent)(
         llm=llm,
         template=date,
         output_format=InterpreterParser(
-            template=GRAPH_PROMPT_CN if lang == "cn" else GRAPH_PROMPT_EN
+            begin="<|action_start|><|interpreter|>",
+            end="<|action_end|>",
+            template=GRAPH_PROMPT_CN if lang == "cn" else GRAPH_PROMPT_EN,
         ),
         searcher_cfg=dict(
             llm=llm,
             plugins=plugins,
             template=date,
             output_format=PluginParser(
+                begin="<|action_start|><|plugin|>",
+                end="<|action_end|>",
                 template=searcher_system_prompt_cn if lang == "cn" else searcher_system_prompt_en,
                 tool_info=get_plugin_prompt(plugins),
             ),
-            user_input_template=(
-                searcher_input_template_cn if lang == "cn" else searcher_input_template_en
-            ),
-            user_context_template=(
-                searcher_context_template_cn if lang == "cn" else searcher_context_template_en
-            ),
+            user_input_template=(searcher_input_template_cn if lang == "cn" else searcher_input_template_en),
+            user_context_template=(searcher_context_template_cn if lang == "cn" else searcher_context_template_en),
         ),
         summary_prompt=FINAL_RESPONSE_CN if lang == "cn" else FINAL_RESPONSE_EN,
         max_turn=10,