Upload 11 files

mindsearch/agent/__init__.py

@@ -0,0 +1,60 @@
+import os
+from datetime import datetime
+
+from lagent.actions import ActionExecutor, BingBrowser
+
+import mindsearch.agent.models as llm_factory
+from mindsearch.agent.mindsearch_agent import (MindSearchAgent,
+                                               MindSearchProtocol)
+from mindsearch.agent.mindsearch_prompt import (
+    FINAL_RESPONSE_CN, FINAL_RESPONSE_EN, GRAPH_PROMPT_CN, GRAPH_PROMPT_EN,
+    fewshot_example_cn, fewshot_example_en, graph_fewshot_example_cn,
+    graph_fewshot_example_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='DuckDuckGoSearch'):
+    llm = LLM.get(model_format, None)
+    if llm is None:
+        llm_cfg = getattr(llm_factory, model_format)
+        if llm_cfg is None:
+            raise NotImplementedError
+        llm_cfg = llm_cfg.copy()
+        llm = llm_cfg.pop('type')(**llm_cfg)
+        LLM[model_format] = llm
+
+    interpreter_prompt = GRAPH_PROMPT_CN if lang == 'cn' else GRAPH_PROMPT_EN
+    plugin_prompt = searcher_system_prompt_cn if lang == 'cn' else searcher_system_prompt_en
+    if not model_format.lower().startswith('internlm'):
+        interpreter_prompt += graph_fewshot_example_cn if lang == 'cn' else graph_fewshot_example_en
+        plugin_prompt += fewshot_example_cn if lang == 'cn' else fewshot_example_en
+
+    agent = MindSearchAgent(
+        llm=llm,
+        protocol=MindSearchProtocol(meta_prompt=datetime.now().strftime(
+            'The current date is %Y-%m-%d.'),
+                                    interpreter_prompt=interpreter_prompt,
+                                    response_prompt=FINAL_RESPONSE_CN
+                                    if lang == 'cn' else FINAL_RESPONSE_EN),
+        searcher_cfg=dict(
+            llm=llm,
+            plugin_executor=ActionExecutor(
+                BingBrowser(searcher_type=search_engine,
+                            topk=6,
+                            api_key=os.environ.get('BING_API_KEY',
+                                                   'YOUR BING API'))),
+            protocol=MindSearchProtocol(
+                meta_prompt=datetime.now().strftime(
+                    'The current date is %Y-%m-%d.'),
+                plugin_prompt=plugin_prompt,
+            ),
+            template=dict(input=searcher_input_template_cn
+                          if lang == 'cn' else searcher_input_template_en,
+                          context=searcher_context_template_cn
+                          if lang == 'cn' else searcher_context_template_en)),
+        max_turn=10)
+    return agent

mindsearch/agent/mindsearch_agent.py

@@ -0,0 +1,422 @@
+import json
+import logging
+import queue
+import random
+import re
+import threading
+import uuid
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from copy import deepcopy
+from dataclasses import asdict
+from typing import Dict, List, Optional
+
+from lagent.actions import ActionExecutor
+from lagent.agents import BaseAgent, Internlm2Agent
+from lagent.agents.internlm2_agent import Internlm2Protocol
+from lagent.schema import AgentReturn, AgentStatusCode, ModelStatusCode
+from termcolor import colored
+
+# 初始化日志记录
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class SearcherAgent(Internlm2Agent):
+
+    def __init__(self, template='{query}', **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.template = template
+
+    def stream_chat(self,
+                    question: str,
+                    root_question: str = None,
+                    parent_response: List[dict] = None,
+                    **kwargs) -> AgentReturn:
+        message = self.template['input'].format(question=question,
+                                                topic=root_question)
+        if parent_response:
+            if 'context' in self.template:
+                parent_response = [
+                    self.template['context'].format(**item)
+                    for item in parent_response
+                ]
+                message = '\n'.join(parent_response + [message])
+        print(colored(f'current query: {message}', 'green'))
+        for agent_return in super().stream_chat(message,
+                                                session_id=random.randint(
+                                                    0, 999999),
+                                                **kwargs):
+            agent_return.type = 'searcher'
+            agent_return.content = question
+            yield deepcopy(agent_return)
+
+
+class MindSearchProtocol(Internlm2Protocol):
+
+    def __init__(
+        self,
+        meta_prompt: str = None,
+        interpreter_prompt: str = None,
+        plugin_prompt: str = None,
+        few_shot: Optional[List] = None,
+        response_prompt: str = None,
+        language: Dict = dict(
+            begin='',
+            end='',
+            belong='assistant',
+        ),
+        tool: Dict = dict(
+            begin='{start_token}{name}\n',
+            start_token='<|action_start|>',
+            name_map=dict(plugin='<|plugin|>', interpreter='<|interpreter|>'),
+            belong='assistant',
+            end='<|action_end|>\n',
+        ),
+        execute: Dict = dict(role='execute',
+                             begin='',
+                             end='',
+                             fallback_role='environment'),
+    ) -> None:
+        self.response_prompt = response_prompt
+        super().__init__(meta_prompt=meta_prompt,
+                         interpreter_prompt=interpreter_prompt,
+                         plugin_prompt=plugin_prompt,
+                         few_shot=few_shot,
+                         language=language,
+                         tool=tool,
+                         execute=execute)
+
+    def format(self,
+               inner_step: List[Dict],
+               plugin_executor: ActionExecutor = None,
+               **kwargs) -> list:
+        formatted = []
+        if self.meta_prompt:
+            formatted.append(dict(role='system', content=self.meta_prompt))
+        if self.plugin_prompt:
+            plugin_prompt = self.plugin_prompt.format(tool_info=json.dumps(
+                plugin_executor.get_actions_info(), ensure_ascii=False))
+            formatted.append(
+                dict(role='system', content=plugin_prompt, name='plugin'))
+        if self.interpreter_prompt:
+            formatted.append(
+                dict(role='system',
+                     content=self.interpreter_prompt,
+                     name='interpreter'))
+        if self.few_shot:
+            for few_shot in self.few_shot:
+                formatted += self.format_sub_role(few_shot)
+        formatted += self.format_sub_role(inner_step)
+        return formatted
+
+
+class WebSearchGraph:
+    end_signal = 'end'
+    searcher_cfg = dict()
+
+    def __init__(self):
+        self.nodes = {}
+        self.adjacency_list = defaultdict(list)
+        self.executor = ThreadPoolExecutor(max_workers=10)
+        self.future_to_query = dict()
+        self.searcher_resp_queue = queue.Queue()
+
+    def add_root_node(self, node_content, node_name='root'):
+        self.nodes[node_name] = dict(content=node_content, type='root')
+        self.adjacency_list[node_name] = []
+        self.searcher_resp_queue.put((node_name, self.nodes[node_name], []))
+
+    def add_node(self, node_name, node_content):
+        self.nodes[node_name] = dict(content=node_content, type='searcher')
+        self.adjacency_list[node_name] = []
+
+        def model_stream_thread():
+            agent = SearcherAgent(**self.searcher_cfg)
+            try:
+                parent_nodes = []
+                for start_node, adj in self.adjacency_list.items():
+                    for neighbor in adj:
+                        if node_name == neighbor[
+                                'name'] and start_node in self.nodes and 'response' in self.nodes[
+                                    start_node]:
+                            parent_nodes.append(self.nodes[start_node])
+                parent_response = [
+                    dict(question=node['content'], answer=node['response'])
+                    for node in parent_nodes
+                ]
+                for answer in agent.stream_chat(
+                        node_content,
+                        self.nodes['root']['content'],
+                        parent_response=parent_response):
+                    self.searcher_resp_queue.put(
+                        deepcopy((node_name,
+                                  dict(response=answer.response,
+                                       detail=answer), [])))
+                self.nodes[node_name]['response'] = answer.response
+                self.nodes[node_name]['detail'] = answer
+            except Exception as e:
+                logger.exception(f'Error in model_stream_thread: {e}')
+
+        self.future_to_query[self.executor.submit(
+            model_stream_thread)] = f'{node_name}-{node_content}'
+
+    def add_response_node(self, node_name='response'):
+        self.nodes[node_name] = dict(type='end')
+        self.searcher_resp_queue.put((node_name, self.nodes[node_name], []))
+
+    def add_edge(self, start_node, end_node):
+        self.adjacency_list[start_node].append(
+            dict(id=str(uuid.uuid4()), name=end_node, state=2))
+        self.searcher_resp_queue.put((start_node, self.nodes[start_node],
+                                      self.adjacency_list[start_node]))
+
+    def reset(self):
+        self.nodes = {}
+        self.adjacency_list = defaultdict(list)
+
+    def node(self, node_name):
+        return self.nodes[node_name].copy()
+
+
+class MindSearchAgent(BaseAgent):
+
+    def __init__(self,
+                 llm,
+                 searcher_cfg,
+                 protocol=MindSearchProtocol(),
+                 max_turn=10):
+        self.local_dict = {}
+        self.ptr = 0
+        self.llm = llm
+        self.max_turn = max_turn
+        WebSearchGraph.searcher_cfg = searcher_cfg
+        super().__init__(llm=llm, action_executor=None, protocol=protocol)
+
+    def stream_chat(self, message, **kwargs):
+        if isinstance(message, str):
+            message = [{'role': 'user', 'content': message}]
+        elif isinstance(message, dict):
+            message = [message]
+        as_dict = kwargs.pop('as_dict', False)
+        return_early = kwargs.pop('return_early', False)
+        self.local_dict.clear()
+        self.ptr = 0
+        inner_history = message[:]
+        agent_return = AgentReturn()
+        agent_return.type = 'planner'
+        agent_return.nodes = {}
+        agent_return.adjacency_list = {}
+        agent_return.inner_steps = deepcopy(inner_history)
+        for _ in range(self.max_turn):
+            prompt = self._protocol.format(inner_step=inner_history)
+            code = None
+            for model_state, response, _ in self.llm.stream_chat(
+                    prompt, session_id=random.randint(0, 999999), **kwargs):
+                if model_state.value < 0:
+                    agent_return.state = getattr(AgentStatusCode,
+                                                 model_state.name)
+                    yield deepcopy(agent_return)
+                    return
+                response = response.replace('<|plugin|>', '<|interpreter|>')
+                _, language, action = self._protocol.parse(response)
+                if not language and not action:
+                    continue
+                code = action['parameters']['command'] if action else ''
+                agent_return.state = self._determine_agent_state(
+                    model_state, code, agent_return)
+                agent_return.response = language if not code else code
+
+                # if agent_return.state == AgentStatusCode.STREAM_ING:
+                yield deepcopy(agent_return)
+
+            inner_history.append({'role': 'language', 'content': language})
+            print(colored(response, 'blue'))
+
+            if code:
+                yield from self._process_code(agent_return, inner_history,
+                                              code, as_dict, return_early)
+            else:
+                agent_return.state = AgentStatusCode.END
+                yield deepcopy(agent_return)
+                return
+
+        agent_return.state = AgentStatusCode.END
+        yield deepcopy(agent_return)
+
+    def _determine_agent_state(self, model_state, code, agent_return):
+        if code:
+            return (AgentStatusCode.PLUGIN_START if model_state
+                    == ModelStatusCode.END else AgentStatusCode.PLUGIN_START)
+        return (AgentStatusCode.ANSWER_ING
+                if agent_return.nodes and 'response' in agent_return.nodes else
+                AgentStatusCode.STREAM_ING)
+
+    def _process_code(self,
+                      agent_return,
+                      inner_history,
+                      code,
+                      as_dict=False,
+                      return_early=False):
+        for node_name, node, adj in self.execute_code(
+                code, return_early=return_early):
+            if as_dict and 'detail' in node:
+                node['detail'] = asdict(node['detail'])
+            if not adj:
+                agent_return.nodes[node_name] = node
+            else:
+                agent_return.adjacency_list[node_name] = adj
+            # state  1进行中，2未开始，3已结束
+            for start_node, neighbors in agent_return.adjacency_list.items():
+                for neighbor in neighbors:
+                    if neighbor['name'] not in agent_return.nodes:
+                        state = 2
+                    elif 'detail' not in agent_return.nodes[neighbor['name']]:
+                        state = 2
+                    elif agent_return.nodes[neighbor['name']][
+                            'detail'].state == AgentStatusCode.END:
+                        state = 3
+                    else:
+                        state = 1
+                    neighbor['state'] = state
+            if not adj:
+                yield deepcopy((agent_return, node_name))
+        reference, references_url = self._generate_reference(
+            agent_return, code, as_dict)
+        inner_history.append({
+            'role': 'tool',
+            'content': code,
+            'name': 'plugin'
+        })
+        inner_history.append({
+            'role': 'environment',
+            'content': reference,
+            'name': 'plugin'
+        })
+        agent_return.inner_steps = deepcopy(inner_history)
+        agent_return.state = AgentStatusCode.PLUGIN_RETURN
+        agent_return.references.update(references_url)
+        yield deepcopy(agent_return)
+
+    def _generate_reference(self, agent_return, code, as_dict):
+        node_list = [
+            node.strip().strip('\"') for node in re.findall(
+                r'graph\.node\("((?:[^"\\]|\\.)*?)"\)', code)
+        ]
+        if 'add_response_node' in code:
+            return self._protocol.response_prompt, dict()
+        references = []
+        references_url = dict()
+        for node_name in node_list:
+            ref_results = None
+            ref2url = None
+            if as_dict:
+                actions = agent_return.nodes[node_name]['detail']['actions']
+            else:
+                actions = agent_return.nodes[node_name]['detail'].actions
+            if actions:
+                ref_results = actions[0]['result'][0][
+                    'content'] if as_dict else actions[0].result[0]['content']
+            if ref_results:
+                ref_results = json.loads(ref_results)
+                ref2url = {
+                    idx: item['url']
+                    for idx, item in ref_results.items()
+                }
+
+            ref = f"## {node_name}\n\n{agent_return.nodes[node_name]['response']}\n"
+            updated_ref = re.sub(
+                r'\[\[(\d+)\]\]',
+                lambda match: f'[[{int(match.group(1)) + self.ptr}]]', ref)
+            numbers = [int(n) for n in re.findall(r'\[\[(\d+)\]\]', ref)]
+            if numbers:
+                try:
+                    assert all(str(elem) in ref2url for elem in numbers)
+                except Exception as exc:
+                    logger.info(f'Illegal reference id: {str(exc)}')
+                if ref2url:
+                    references_url.update({
+                        str(idx + self.ptr): ref2url[str(idx)]
+                        for idx in set(numbers) if str(idx) in ref2url
+                    })
+                self.ptr += max(numbers) + 1
+            references.append(updated_ref)
+        return '\n'.join(references), references_url
+
+    def execute_code(self, command: str, return_early=False):
+
+        def extract_code(text: str) -> str:
+            text = re.sub(r'from ([\w.]+) import WebSearchGraph', '', text)
+            triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
+            single_match = re.search(r'`([^`]*)`', text, re.DOTALL)
+            if triple_match:
+                return triple_match.group(1)
+            elif single_match:
+                return single_match.group(1)
+            return text
+
+        def run_command(cmd):
+            try:
+                exec(cmd, globals(), self.local_dict)
+                plan_graph = self.local_dict.get('graph')
+                assert plan_graph is not None
+                for future in as_completed(plan_graph.future_to_query):
+                    future.result()
+                plan_graph.future_to_query.clear()
+                plan_graph.searcher_resp_queue.put(plan_graph.end_signal)
+            except Exception as e:
+                logger.exception(f'Error executing code: {e}')
+                raise
+
+        command = extract_code(command)
+        producer_thread = threading.Thread(target=run_command,
+                                           args=(command, ))
+        producer_thread.start()
+
+        responses = defaultdict(list)
+        ordered_nodes = []
+        active_node = None
+
+        while True:
+            try:
+                item = self.local_dict.get('graph').searcher_resp_queue.get(
+                    timeout=60)
+                if item is WebSearchGraph.end_signal:
+                    for node_name in ordered_nodes:
+                        # resp = None
+                        for resp in responses[node_name]:
+                            yield deepcopy(resp)
+                        # if resp:
+                        #     assert resp[1][
+                        #         'detail'].state == AgentStatusCode.END
+                    break
+                node_name, node, adj = item
+                if node_name in ['root', 'response']:
+                    yield deepcopy((node_name, node, adj))
+                else:
+                    if node_name not in ordered_nodes:
+                        ordered_nodes.append(node_name)
+                    responses[node_name].append((node_name, node, adj))
+                    if not active_node and ordered_nodes:
+                        active_node = ordered_nodes[0]
+                    while active_node and responses[active_node]:
+                        if return_early:
+                            if 'detail' in responses[active_node][-1][
+                                    1] and responses[active_node][-1][1][
+                                        'detail'].state == AgentStatusCode.END:
+                                item = responses[active_node][-1]
+                            else:
+                                item = responses[active_node].pop(0)
+                        else:
+                            item = responses[active_node].pop(0)
+                        if 'detail' in item[1] and item[1][
+                                'detail'].state == AgentStatusCode.END:
+                            ordered_nodes.pop(0)
+                            responses[active_node].clear()
+                            active_node = None
+                        yield deepcopy(item)
+            except queue.Empty:
+                if not producer_thread.is_alive():
+                    break
+        producer_thread.join()
+        return

mindsearch/agent/mindsearch_prompt.py

@@ -0,0 +1,326 @@
+# flake8: noqa
+
+searcher_system_prompt_cn = """## 人物简介
+你是一个可以调用网络搜索工具的智能助手。请根据"当前问题"，调用搜索工具收集信息并回复问题。你能够调用如下工具:
+{tool_info}
+## 回复格式
+
+调用工具时，请按照以下格式:
+```
+你的思考过程...<|action_start|><|plugin|>{{"name": "tool_name", "parameters": {{"param1": "value1"}}}}<|action_end|>
+```
+
+## 要求
+
+- 回答中每个关键点需标注引用的搜索结果来源，以确保信息的可信度。给出索引的形式为`[[int]]`，如果有多个索引，则用多个[[]]表示，如`[[id_1]][[id_2]]`。
+- 基于"当前问题"的搜索结果，撰写详细完备的回复，优先回答"当前问题"。
+
+"""
+
+searcher_system_prompt_en = """## Character Introduction
+You are an intelligent assistant that can call web search tools. Please collect information and reply to the question based on the current problem. You can use the following tools:
+{tool_info}
+## Reply Format
+
+When calling the tool, please follow the format below:
+```
+Your thought process...<|action_start|><|plugin|>{{"name": "tool_name", "parameters": {{"param1": "value1"}}}}<|action_end|>
+```
+
+## Requirements
+
+- Each key point in the response should be marked with the source of the search results to ensure the credibility of the information. The citation format is `[[int]]`. If there are multiple citations, use multiple [[]] to provide the index, such as `[[id_1]][[id_2]]`.
+- Based on the search results of the "current problem", write a detailed and complete reply to answer the "current problem".
+"""
+
+fewshot_example_cn = """
+## 样例
+
+### search
+当我希望搜索"王者荣耀现在是什么赛季"时，我会按照以下格式进行操作:
+现在是2024年，因此我应该搜索王者荣耀赛季关键词<|action_start|><|plugin|>{{"name": "FastWebBrowser.search", "parameters": {{"query": ["王者荣耀 赛季", "2024年王者荣耀赛季"]}}}}<|action_end|>
+
+### select
+为了找到王者荣耀s36赛季最强射手，我需要寻找提及王者荣耀s36射手的网页。初步浏览网页后，发现网页0提到王者荣耀s36赛季的信息，但没有具体提及射手的相关信息。网页3提到“s36最强射手出现？”，有可能包含最强射手信息。网页13提到“四大T0英雄崛起，射手荣耀降临”，可能包含最强射手的信息。因此，我选择了网页3和网页13进行进一步阅读。<|action_start|><|plugin|>{{"name": "FastWebBrowser.select", "parameters": {{"index": [3, 13]}}}}<|action_end|>
+"""
+
+fewshot_example_en = """
+## Example
+
+### search
+When I want to search for "What season is Honor of Kings now", I will operate in the following format:
+Now it is 2024, so I should search for the keyword of the Honor of Kings<|action_start|><|plugin|>{{"name": "FastWebBrowser.search", "parameters": {{"query": ["Honor of Kings Season", "season for Honor of Kings in 2024"]}}}}<|action_end|>
+
+### select
+In order to find the strongest shooters in Honor of Kings in season s36, I needed to look for web pages that mentioned shooters in Honor of Kings in season s36. After an initial browse of the web pages, I found that web page 0 mentions information about Honor of Kings in s36 season, but there is no specific mention of information about the shooter. Webpage 3 mentions that “the strongest shooter in s36 has appeared?”, which may contain information about the strongest shooter. Webpage 13 mentions “Four T0 heroes rise, archer's glory”, which may contain information about the strongest archer. Therefore, I chose webpages 3 and 13 for further reading.<|action_start|><|plugin|>{{"name": "FastWebBrowser.select", "parameters": {{"index": [3, 13]}}}}<|action_end|>
+"""
+
+searcher_input_template_en = """## Final Problem
+{topic}
+## Current Problem
+{question}
+"""
+
+searcher_input_template_cn = """## 主问题
+{topic}
+## 当前问题
+{question}
+"""
+
+searcher_context_template_en = """## Historical Problem
+{question}
+Answer: {answer}
+"""
+
+searcher_context_template_cn = """## 历史问题
+{question}
+回答：{answer}
+"""
+
+search_template_cn = '## {query}\n\n{result}\n'
+search_template_en = '## {query}\n\n{result}\n'
+
+GRAPH_PROMPT_CN = """## 人物简介
+你是一个可以利用 Jupyter 环境 Python 编程的程序员。你可以利用提供的 API 来构建 Web 搜索图，最终生成代码并执行。
+
+## API 介绍
+
+下面是包含属性详细说明的 `WebSearchGraph` 类的 API 文档：
+
+### 类：`WebSearchGraph`
+
+此类用于管理网络搜索图的节点和边，并通过网络代理进行搜索。
+
+#### 初始化方法
+
+初始化 `WebSearchGraph` 实例。
+
+**属性：**
+
+- `nodes` (Dict[str, Dict[str, str]]): 存储图中所有节点的字典。每个节点由其名称索引，并包含内容、类型以及其他相关信息。
+- `adjacency_list` (Dict[str, List[str]]): 存储图中所有节点之间连接关系的邻接表。每个节点由其名称索引，并包含一个相邻节点名称的列表。
+
+
+#### 方法：`add_root_node`
+
+添加原始问题作为根节点。
+**参数：**
+
+- `node_content` (str): 用户提出的问题。
+- `node_name` (str, 可选): 节点名称，默认为 'root'。
+
+
+#### 方法：`add_node`
+
+添加搜索子问题节点并返回搜索结果。
+**参数：
+
+- `node_name` (str): 节点名称。
+- `node_content` (str): 子问题内容。
+
+**返回：**
+
+- `str`: 返回搜索结果。
+
+
+#### 方法：`add_response_node`
+
+当前获取的信息已经满足问题需求，添加回复节点。
+
+**参数：**
+
+- `node_name` (str, 可选): 节点名称，默认为 'response'。
+
+
+#### 方法：`add_edge`
+
+添加边。
+
+**参数：**
+
+- `start_node` (str): 起始节点名称。
+- `end_node` (str): 结束节点名称。
+
+
+#### 方法：`reset`
+
+重置节点和边。
+
+
+#### 方法：`node`
+
+获取节点信息。
+
+```python
+def node(self, node_name: str) -> str
+```
+
+**参数：**
+
+- `node_name` (str): 节点名称。
+
+**返回：**
+
+- `str`: 返回包含节点信息的字典，包含节点的内容、类型、思考过程（如果有）和前驱节点列表。
+
+## 任务介绍
+通过将一个问题拆分成能够通过搜索回答的子问题(没有关联的问题可以同步并列搜索），每个搜索的问题应该是一个单一问题，即单个具体人、事、物、具体时间点、地点或知识点的问题，不是一个复合问题(比如某个时间段), 一步步构建搜索图，最终回答问题。
+
+## 注意事项
+
+1. 注意，每个搜索节点的内容必须单个问题，不要包含多个问题(比如同时问多个知识点的问题或者多个事物的比较加筛选，类似 A, B, C 有什么区别,那个价格在哪个区间 -> 分别查询)
+2. 不要杜撰搜索结果，要等待代码返回结果
+3. 同样的问题不要重复提问，可以在已有问题的基础上继续提问
+4. 添加 response 节点的时候，要单独添加，不要和其他节点一起添加，不能同时添加 response 节点和其他节点
+5. 一次输出中，不要包含多个代码块，每次只能有一个代码块
+6. 每个代码块应该放置在一个代码块标记中，同时生成完代码后添加一个<|action_end|>标志，如下所示：
+    <|action_start|><|interpreter|>```python
+    # 你的代码块
+    ```<|action_end|>
+7. 最后一次回复应该是添加node_name为'response'的 response 节点，必须添加 response 节点，不要添加其他节点
+"""
+
+GRAPH_PROMPT_EN = """## Character Profile
+You are a programmer capable of Python programming in a Jupyter environment. You can utilize the provided API to construct a Web Search Graph, ultimately generating and executing code.
+
+## API Description
+
+Below is the API documentation for the WebSearchGraph class, including detailed attribute descriptions:
+
+### Class: WebSearchGraph
+
+This class manages nodes and edges of a web search graph and conducts searches via a web proxy.
+
+#### Initialization Method
+
+Initializes an instance of WebSearchGraph.
+
+**Attributes:**
+
+- nodes (Dict[str, Dict[str, str]]): A dictionary storing all nodes in the graph. Each node is indexed by its name and contains content, type, and other related information.
+- adjacency_list (Dict[str, List[str]]): An adjacency list storing the connections between all nodes in the graph. Each node is indexed by its name and contains a list of adjacent node names.
+
+#### Method: add_root_node
+
+Adds the initial question as the root node.
+**Parameters:**
+
+- node_content (str): The user's question.
+- node_name (str, optional): The node name, default is 'root'.
+
+#### Method: add_node
+
+Adds a sub-question node and returns search results.
+**Parameters:**
+
+- node_name (str): The node name.
+- node_content (str): The sub-question content.
+
+**Returns:**
+
+- str: Returns the search results.
+
+#### Method: add_response_node
+
+Adds a response node when the current information satisfies the question's requirements.
+
+**Parameters:**
+
+- node_name (str, optional): The node name, default is 'response'.
+
+#### Method: add_edge
+
+Adds an edge.
+
+**Parameters:**
+
+- start_node (str): The starting node name.
+- end_node (str): The ending node name.
+
+#### Method: reset
+
+Resets nodes and edges.
+
+#### Method: node
+
+Get node information.
+
+python
+def node(self, node_name: str) -> str
+
+**Parameters:**
+
+- node_name (str): The node name.
+
+**Returns:**
+
+- str: Returns a dictionary containing the node's information, including content, type, thought process (if any), and list of predecessor nodes.
+
+## Task Description
+By breaking down a question into sub-questions that can be answered through searches (unrelated questions can be searched concurrently), each search query should be a single question focusing on a specific person, event, object, specific time point, location, or knowledge point. It should not be a compound question (e.g., a time period). Step by step, build the search graph to finally answer the question.
+
+## Considerations
+
+1. Each search node's content must be a single question; do not include multiple questions (e.g., do not ask multiple knowledge points or compare and filter multiple things simultaneously, like asking for differences between A, B, and C, or price ranges -> query each separately).
+2. Do not fabricate search results; wait for the code to return results.
+3. Do not repeat the same question; continue asking based on existing questions.
+4. When adding a response node, add it separately; do not add a response node and other nodes simultaneously.
+5. In a single output, do not include multiple code blocks; only one code block per output.
+6. Each code block should be placed within a code block marker, and after generating the code, add an <|action_end|> tag as shown below:
+    <|action_start|><|interpreter|>
+    ```python
+    # Your code block (Note that the 'Get new added node information' logic must be added at the end of the code block, such as 'graph.node('...')')
+    ```<|action_end|>
+7. The final response should add a response node with node_name 'response', and no other nodes should be added.
+"""
+
+graph_fewshot_example_cn = """
+## 返回格式示例
+<|action_start|><|interpreter|>```python
+graph = WebSearchGraph()
+graph.add_root_node(node_content="哪家大模型API最便宜?", node_name="root") # 添加原始问题作为根节点
+graph.add_node(
+        node_name="大模型API提供商", # 节点名称最好有意义
+        node_content="目前有哪些主要的大模型API提供商？")
+graph.add_node(
+        node_name="sub_name_2", # 节点名称最好有意义
+        node_content="content of sub_name_2")
+...
+graph.add_edge(start_node="root", end_node="sub_name_1")
+...
+graph.node("大模型API提供商"), graph.node("sub_name_2"), ...
+```<|action_end|>
+"""
+
+graph_fewshot_example_en = """
+## Response Format
+<|action_start|><|interpreter|>```python
+graph = WebSearchGraph()
+graph.add_root_node(node_content="Which large model API is the cheapest?", node_name="root") # Add the original question as the root node
+graph.add_node(
+        node_name="Large Model API Providers", # The node name should be meaningful
+        node_content="Who are the main large model API providers currently?")
+graph.add_node(
+        node_name="sub_name_2", # The node name should be meaningful
+        node_content="content of sub_name_2")
+...
+graph.add_edge(start_node="root", end_node="sub_name_1")
+...
+# Get node info
+graph.node("Large Model API Providers"), graph.node("sub_name_2"), ...
+```<|action_end|>
+"""
+
+FINAL_RESPONSE_CN = """基于提供的问答对，撰写一篇详细完备的最终回答。
+- 回答内容需要逻辑清晰，层次分明，确保读者易于理解。
+- 回答中每个关键点需标注引用的搜索结果来源(保持跟问答对中的索引一致)，以确保信息的可信度。给出索引的形式为`[[int]]`，如果有多个索引，则用多个[[]]表示，如`[[id_1]][[id_2]]`。
+- 回答部分需要全面且完备，不要出现"基于上述内容"等模糊表达，最终呈现的回答不包括提供给你的问答对。
+- 语言风格需要专业、严谨，避免口语化表达。
+- 保持统一的语法和词汇使用，确保整体文档的一致性和连贯性。"""
+
+FINAL_RESPONSE_EN = """Based on the provided Q&A pairs, write a detailed and comprehensive final response.
+- The response content should be logically clear and well-structured to ensure reader understanding.
+- Each key point in the response should be marked with the source of the search results (consistent with the indices in the Q&A pairs) to ensure information credibility. The index is in the form of `[[int]]`, and if there are multiple indices, use multiple `[[]]`, such as `[[id_1]][[id_2]]`.
+- The response should be comprehensive and complete, without vague expressions like "based on the above content". The final response should not include the Q&A pairs provided to you.
+- The language style should be professional and rigorous, avoiding colloquial expressions.
+- Maintain consistent grammar and vocabulary usage to ensure overall document consistency and coherence."""

mindsearch/agent/models.py

@@ -0,0 +1,77 @@
+import os
+
+from lagent.llms import (GPTAPI, INTERNLM2_META, HFTransformerCasualLM,
+                         LMDeployClient, LMDeployServer)
+
+internlm_server = dict(type=LMDeployServer,
+                       path='internlm/internlm2_5-7b-chat',
+                       model_name='internlm2',
+                       meta_template=INTERNLM2_META,
+                       top_p=0.8,
+                       top_k=1,
+                       temperature=0,
+                       max_new_tokens=8192,
+                       repetition_penalty=1.02,
+                       stop_words=['<|im_end|>'])
+
+internlm_client = dict(type=LMDeployClient,
+                       model_name='internlm2_5-7b-chat',
+                       url='http://127.0.0.1:23333',
+                       meta_template=INTERNLM2_META,
+                       top_p=0.8,
+                       top_k=1,
+                       temperature=0,
+                       max_new_tokens=8192,
+                       repetition_penalty=1.02,
+                       stop_words=['<|im_end|>'])
+
+internlm_hf = dict(type=HFTransformerCasualLM,
+                   path='internlm/internlm2_5-7b-chat',
+                   meta_template=INTERNLM2_META,
+                   top_p=0.8,
+                   top_k=None,
+                   temperature=1e-6,
+                   max_new_tokens=8192,
+                   repetition_penalty=1.02,
+                   stop_words=['<|im_end|>'])
+# openai_api_base needs to fill in the complete chat api address, such as: https://api.openai.com/v1/chat/completions
+gpt4 = dict(type=GPTAPI,
+            model_type='gpt-4-turbo',
+            key=os.environ.get('OPENAI_API_KEY', 'YOUR OPENAI API KEY'),
+            openai_api_base=os.environ.get('OPENAI_API_BASE', 'https://api.openai.com/v1/chat/completions'),
+            )
+
+url = 'https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation'
+qwen = dict(type=GPTAPI,
+            model_type='qwen-max-longcontext',
+            key=os.environ.get('QWEN_API_KEY', 'YOUR QWEN API KEY'),
+            openai_api_base=url,
+            meta_template=[
+                dict(role='system', api_role='system'),
+                dict(role='user', api_role='user'),
+                dict(role='assistant', api_role='assistant'),
+                dict(role='environment', api_role='system')
+            ],
+            top_p=0.8,
+            top_k=1,
+            temperature=0,
+            max_new_tokens=4096,
+            repetition_penalty=1.02,
+            stop_words=['<|im_end|>'])
+
+internlm_silicon = dict(type=GPTAPI,
+                        model_type='internlm/internlm2_5-7b-chat',
+                        key=os.environ.get('SILICON_API_KEY', 'YOUR SILICON API KEY'),
+                        openai_api_base='https://api.siliconflow.cn/v1/chat/completions',
+                        meta_template=[
+                            dict(role='system', api_role='system'),
+                            dict(role='user', api_role='user'),
+                            dict(role='assistant', api_role='assistant'),
+                            dict(role='environment', api_role='system')
+                        ],
+                        top_p=0.8,
+                        top_k=1,
+                        temperature=0,
+                        max_new_tokens=8192,
+                        repetition_penalty=1.02,
+                        stop_words=['<|im_end|>'])

mindsearch/app.py

@@ -0,0 +1,136 @@
+import asyncio
+import json
+import logging
+from copy import deepcopy
+from dataclasses import asdict
+from typing import Dict, List, Union
+
+import janus
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from lagent.schema import AgentStatusCode
+from pydantic import BaseModel
+from sse_starlette.sse import EventSourceResponse
+
+from mindsearch.agent import init_agent
+
+
+def parse_arguments():
+    import argparse
+    parser = argparse.ArgumentParser(description='MindSearch API')
+    parser.add_argument('--lang', default='cn', type=str, help='Language')
+    parser.add_argument('--model_format',
+                        default='internlm_server',
+                        type=str,
+                        help='Model format')
+    parser.add_argument('--search_engine',
+                       default='DuckDuckGoSearch',
+                       type=str,
+                       help='Search engine')
+    return parser.parse_args()
+
+
+args = parse_arguments()
+app = FastAPI(docs_url='/')
+
+app.add_middleware(CORSMiddleware,
+                   allow_origins=['*'],
+                   allow_credentials=True,
+                   allow_methods=['*'],
+                   allow_headers=['*'])
+
+
+class GenerationParams(BaseModel):
+    inputs: Union[str, List[Dict]]
+    agent_cfg: Dict = dict()
+
+
+@app.post('/solve')
+async def run(request: GenerationParams):
+
+    def convert_adjacency_to_tree(adjacency_input, root_name):
+
+        def build_tree(node_name):
+            node = {'name': node_name, 'children': []}
+            if node_name in adjacency_input:
+                for child in adjacency_input[node_name]:
+                    child_node = build_tree(child['name'])
+                    child_node['state'] = child['state']
+                    child_node['id'] = child['id']
+                    node['children'].append(child_node)
+            return node
+
+        return build_tree(root_name)
+
+    async def generate():
+        try:
+            queue = janus.Queue()
+            stop_event = asyncio.Event()
+
+            # Wrapping a sync generator as an async generator using run_in_executor
+            def sync_generator_wrapper():
+                try:
+                    for response in agent.stream_chat(inputs):
+                        queue.sync_q.put(response)
+                except Exception as e:
+                    logging.exception(
+                        f'Exception in sync_generator_wrapper: {e}')
+                finally:
+                    # Notify async_generator_wrapper that the data generation is complete.
+                    queue.sync_q.put(None)
+
+            async def async_generator_wrapper():
+                loop = asyncio.get_event_loop()
+                loop.run_in_executor(None, sync_generator_wrapper)
+                while True:
+                    response = await queue.async_q.get()
+                    if response is None:  # Ensure that all elements are consumed
+                        break
+                    yield response
+                    if not isinstance(
+                            response,
+                            tuple) and response.state == AgentStatusCode.END:
+                        break
+                stop_event.set()  # Inform sync_generator_wrapper to stop
+
+            async for response in async_generator_wrapper():
+                if isinstance(response, tuple):
+                    agent_return, node_name = response
+                else:
+                    agent_return = response
+                    node_name = None
+                origin_adj = deepcopy(agent_return.adjacency_list)
+                adjacency_list = convert_adjacency_to_tree(
+                    agent_return.adjacency_list, 'root')
+                assert adjacency_list[
+                    'name'] == 'root' and 'children' in adjacency_list
+                agent_return.adjacency_list = adjacency_list['children']
+                agent_return = asdict(agent_return)
+                agent_return['adj'] = origin_adj
+                response_json = json.dumps(dict(response=agent_return,
+                                                current_node=node_name),
+                                           ensure_ascii=False)
+                yield {'data': response_json}
+                # yield f'data: {response_json}\n\n'
+        except Exception as exc:
+            msg = 'An error occurred while generating the response.'
+            logging.exception(msg)
+            response_json = json.dumps(
+                dict(error=dict(msg=msg, details=str(exc))),
+                ensure_ascii=False)
+            yield {'data': response_json}
+            # yield f'data: {response_json}\n\n'
+        finally:
+            await stop_event.wait(
+            )  # Waiting for async_generator_wrapper to stop
+            queue.close()
+            await queue.wait_closed()
+
+    inputs = request.inputs
+    agent = init_agent(lang=args.lang, model_format=args.model_format,search_engine=args.search_engine)
+    return EventSourceResponse(generate())
+
+
+if __name__ == '__main__':
+    import uvicorn
+    uvicorn.run(app, host='0.0.0.0', port=8002, log_level='info')

mindsearch/terminal.py

@@ -0,0 +1,50 @@
+from datetime import datetime
+
+from lagent.actions import ActionExecutor, BingBrowser
+from lagent.llms import INTERNLM2_META, LMDeployServer
+
+from mindsearch.agent.mindsearch_agent import (MindSearchAgent,
+                                               MindSearchProtocol)
+from mindsearch.agent.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)
+
+lang = 'cn'
+llm = LMDeployServer(path='internlm/internlm2_5-7b-chat',
+                     model_name='internlm2',
+                     meta_template=INTERNLM2_META,
+                     top_p=0.8,
+                     top_k=1,
+                     temperature=0,
+                     max_new_tokens=8192,
+                     repetition_penalty=1.02,
+                     stop_words=['<|im_end|>'])
+
+agent = MindSearchAgent(
+    llm=llm,
+    protocol=MindSearchProtocol(
+        meta_prompt=datetime.now().strftime('The current date is %Y-%m-%d.'),
+        interpreter_prompt=GRAPH_PROMPT_CN
+        if lang == 'cn' else GRAPH_PROMPT_EN,
+        response_prompt=FINAL_RESPONSE_CN
+        if lang == 'cn' else FINAL_RESPONSE_EN),
+    searcher_cfg=dict(
+        llm=llm,
+        plugin_executor=ActionExecutor(
+            BingBrowser(searcher_type='DuckDuckGoSearch', topk=6)),
+        protocol=MindSearchProtocol(
+            meta_prompt=datetime.now().strftime(
+                'The current date is %Y-%m-%d.'),
+            plugin_prompt=searcher_system_prompt_cn
+            if lang == 'cn' else searcher_system_prompt_en,
+        ),
+        template=dict(input=searcher_input_template_cn
+                      if lang == 'cn' else searcher_input_template_en,
+                      context=searcher_context_template_cn
+                      if lang == 'cn' else searcher_context_template_en)),
+    max_turn=10)
+
+for agent_return in agent.stream_chat('上海今天适合穿什么衣服'):
+    pass