Spaces:
Running
Running
| import os | |
| from copy import deepcopy | |
| from datetime import datetime | |
| from lagent.actions import AsyncWebBrowser, WebBrowser | |
| from lagent.agents.stream import get_plugin_prompt | |
| from lagent.prompts import InterpreterParser, PluginParser | |
| from lagent.utils import create_object | |
| from . import models as llm_factory | |
| from .mindsearch_agent import AsyncMindSearchAgent, MindSearchAgent | |
| from .mindsearch_prompt import ( | |
| FINAL_RESPONSE_CN, | |
| FINAL_RESPONSE_EN, | |
| GRAPH_PROMPT_CN, | |
| GRAPH_PROMPT_EN, | |
| searcher_context_template_cn, | |
| searcher_context_template_en, | |
| searcher_input_template_cn, | |
| searcher_input_template_en, | |
| searcher_system_prompt_cn, | |
| searcher_system_prompt_en, | |
| ) | |
| LLM = {} | |
| 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__) | |
| 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"), | |
| ))] | |
| 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), | |
| searcher_cfg=dict( | |
| llm=llm, | |
| plugins=plugins, | |
| template=date, | |
| output_format=PluginParser( | |
| 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), | |
| ), | |
| summary_prompt=FINAL_RESPONSE_CN | |
| if lang == "cn" else FINAL_RESPONSE_EN, | |
| max_turn=10, | |
| ) | |
| return agent | |