提交

app.py

@@ -0,0 +1,314 @@
+import json
+import tempfile
+
+import requests
+import streamlit as st
+from lagent.schema import AgentStatusCode
+from pyvis.network import Network
+
+
+# Function to create the network graph
+def create_network_graph(nodes, adjacency_list):
+    net = Network(height="500px", width="60%", bgcolor="white", font_color="black")
+    for node_id, node_content in nodes.items():
+        net.add_node(node_id, label=node_id, title=node_content, color="#FF5733", size=25)
+    for node_id, neighbors in adjacency_list.items():
+        for neighbor in neighbors:
+            if neighbor["name"] in nodes:
+                net.add_edge(node_id, neighbor["name"])
+    net.show_buttons(filter_=["physics"])
+    return net
+
+
+# Function to draw the graph and return the HTML file path
+def draw_graph(net):
+    path = tempfile.mktemp(suffix=".html")
+    net.save_graph(path)
+    return path
+
+
+def streaming(raw_response):
+    for chunk in raw_response.iter_lines(chunk_size=8192, decode_unicode=False, delimiter=b"\n"):
+        if chunk:
+            decoded = chunk.decode("utf-8")
+            if decoded == "\r":
+                continue
+            if decoded[:6] == "data: ":
+                decoded = decoded[6:]
+            elif decoded.startswith(": ping - "):
+                continue
+            response = json.loads(decoded)
+            yield (
+                response["current_node"],
+                (
+                    response["response"]["formatted"]["node"][response["current_node"]]["response"]
+                    if response["current_node"]
+                    else response["response"]
+                ),
+                response["response"]["formatted"]["adjacency_list"],
+            )
+
+
+# Initialize Streamlit session state
+if "queries" not in st.session_state:
+    st.session_state["queries"] = []
+    st.session_state["responses"] = []
+    st.session_state["graphs_html"] = []
+    st.session_state["nodes_list"] = []
+    st.session_state["adjacency_list_list"] = []
+    st.session_state["history"] = []
+    st.session_state["already_used_keys"] = list()
+
+# Set up page layout
+st.set_page_config(layout="wide")
+st.title("MindSearch-思索")
+
+
+# Function to update chat
+def update_chat(query):
+    with st.chat_message("user"):
+        st.write(query)
+    if query not in st.session_state["queries"]:
+        # Mock data to simulate backend response
+        # response, history, nodes, adjacency_list
+        st.session_state["queries"].append(query)
+        st.session_state["responses"].append([])
+        history = None
+        # 暂不支持多轮
+        # message = [dict(role='user', content=query)]
+
+        url = "http://localhost:8002/solve"
+        headers = {"Content-Type": "application/json"}
+        data = {"inputs": query}
+        raw_response = requests.post(
+            url, headers=headers, data=json.dumps(data), timeout=20, stream=True
+        )
+
+        _nodes, _node_cnt = {}, 0
+        for resp in streaming(raw_response):
+            node_name, response, adjacency_list = resp
+            for name in set(adjacency_list) | {
+                val["name"] for vals in adjacency_list.values() for val in vals
+            }:
+                if name not in _nodes:
+                    _nodes[name] = query if name == "root" else name
+                elif response["stream_state"] == 0:
+                    _nodes[node_name or "response"] = response["formatted"] and response[
+                        "formatted"
+                    ].get("thought")
+            if len(_nodes) != _node_cnt or response["stream_state"] == 0:
+                net = create_network_graph(_nodes, adjacency_list)
+                graph_html_path = draw_graph(net)
+                with open(graph_html_path, encoding="utf-8") as f:
+                    graph_html = f.read()
+                _node_cnt = len(_nodes)
+            else:
+                graph_html = None
+            if "graph_placeholder" not in st.session_state:
+                st.session_state["graph_placeholder"] = st.empty()
+            if "expander_placeholder" not in st.session_state:
+                st.session_state["expander_placeholder"] = st.empty()
+            if graph_html:
+                with st.session_state["expander_placeholder"].expander(
+                    "Show Graph", expanded=False
+                ):
+                    st.session_state["graph_placeholder"]._html(graph_html, height=500)
+            if "container_placeholder" not in st.session_state:
+                st.session_state["container_placeholder"] = st.empty()
+            with st.session_state["container_placeholder"].container():
+                if "columns_placeholder" not in st.session_state:
+                    st.session_state["columns_placeholder"] = st.empty()
+                col1, col2 = st.session_state["columns_placeholder"].columns([2, 1])
+                with col1:
+                    if "planner_placeholder" not in st.session_state:
+                        st.session_state["planner_placeholder"] = st.empty()
+                    if "session_info_temp" not in st.session_state:
+                        st.session_state["session_info_temp"] = ""
+                    if not node_name:
+                        if response["stream_state"] in [
+                            AgentStatusCode.STREAM_ING,
+                            AgentStatusCode.CODING,
+                            AgentStatusCode.CODE_END,
+                        ]:
+                            content = response["formatted"]["thought"]
+                            if response["formatted"]["tool_type"]:
+                                action = response["formatted"]["action"]
+                                if isinstance(action, dict):
+                                    action = json.dumps(action, ensure_ascii=False, indent=4)
+                                content += "\n" + action
+                            st.session_state["session_info_temp"] = content.replace(
+                                "<|action_start|><|interpreter|>\n", "\n"
+                            )
+                        elif response["stream_state"] == AgentStatusCode.CODE_RETURN:
+                            # assert history[-1]["role"] == "environment"
+                            st.session_state["session_info_temp"] += "\n" + response["content"]
+                        st.session_state["planner_placeholder"].markdown(
+                            st.session_state["session_info_temp"]
+                        )
+                        if response["stream_state"] == AgentStatusCode.CODE_RETURN:
+                            st.session_state["responses"][-1].append(
+                                st.session_state["session_info_temp"]
+                            )
+                            st.session_state["session_info_temp"] = ""
+                    else:
+                        st.session_state["planner_placeholder"].markdown(
+                            st.session_state["responses"][-1][-1]
+                            if not st.session_state["session_info_temp"]
+                            else st.session_state["session_info_temp"]
+                        )
+                with col2:
+                    if "selectbox_placeholder" not in st.session_state:
+                        st.session_state["selectbox_placeholder"] = st.empty()
+                    if "searcher_placeholder" not in st.session_state:
+                        st.session_state["searcher_placeholder"] = st.empty()
+                    if node_name:
+                        selected_node_key = (
+                            f"selected_node_{len(st.session_state['queries'])}_{node_name}"
+                        )
+                        if selected_node_key not in st.session_state:
+                            st.session_state[selected_node_key] = node_name
+                        if selected_node_key not in st.session_state["already_used_keys"]:
+                            selected_node = st.session_state["selectbox_placeholder"].selectbox(
+                                "Select a node:",
+                                list(_nodes.keys()),
+                                key=f"key_{selected_node_key}",
+                                index=list(_nodes.keys()).index(node_name),
+                            )
+                            st.session_state["already_used_keys"].append(selected_node_key)
+                        else:
+                            selected_node = node_name
+                        st.session_state[selected_node_key] = selected_node
+                        node_info_key = f"{selected_node}_info"
+                        if node_info_key not in st.session_state:
+                            st.session_state[node_info_key] = [["thought", ""]]
+                        if response["stream_state"] in [AgentStatusCode.STREAM_ING]:
+                            content = response["formatted"]["thought"]
+                            st.session_state[node_info_key][-1][1] = content.replace(
+                                "<|action_start|><|plugin|>\n", "\n```json\n"
+                            )
+                        elif response["stream_state"] in [
+                            AgentStatusCode.PLUGIN_START,
+                            AgentStatusCode.PLUGIN_END,
+                        ]:
+                            thought = response["formatted"]["thought"]
+                            action = response["formatted"]["action"]
+                            if isinstance(action, dict):
+                                action = json.dumps(action, ensure_ascii=False, indent=4)
+                            content = thought + "\n```json\n" + action
+                            if response["stream_state"] == AgentStatusCode.PLUGIN_RETURN:
+                                content += "\n```"
+                            st.session_state[node_info_key][-1][1] = content
+                        elif (
+                            response["stream_state"] == AgentStatusCode.PLUGIN_RETURN
+                            and st.session_state[node_info_key][-1][1]
+                        ):
+                            try:
+                                content = json.loads(response["content"])
+                            except json.decoder.JSONDecodeError:
+                                content = response["content"]
+                            st.session_state[node_info_key].append(
+                                [
+                                    "observation",
+                                    (
+                                        content
+                                        if isinstance(content, str)
+                                        else f"```json\n{json.dumps(content, ensure_ascii=False, indent=4)}\n```"
+                                    ),
+                                ]
+                            )
+                        st.session_state["searcher_placeholder"].markdown(
+                            st.session_state[node_info_key][-1][1]
+                        )
+                        if (
+                            response["stream_state"] == AgentStatusCode.PLUGIN_RETURN
+                            and st.session_state[node_info_key][-1][1]
+                        ):
+                            st.session_state[node_info_key].append(["thought", ""])
+        if st.session_state["session_info_temp"]:
+            st.session_state["responses"][-1].append(st.session_state["session_info_temp"])
+            st.session_state["session_info_temp"] = ""
+        # st.session_state['responses'][-1] = '\n'.join(st.session_state['responses'][-1])
+        st.session_state["graphs_html"].append(graph_html)
+        st.session_state["nodes_list"].append(_nodes)
+        st.session_state["adjacency_list_list"].append(adjacency_list)
+        st.session_state["history"] = history
+
+
+def display_chat_history():
+    for i, query in enumerate(st.session_state["queries"][-1:]):
+        # with st.chat_message('assistant'):
+        if st.session_state["graphs_html"][i]:
+            with st.session_state["expander_placeholder"].expander("Show Graph", expanded=False):
+                st.session_state["graph_placeholder"]._html(
+                    st.session_state["graphs_html"][i], height=500
+                )
+            with st.session_state["container_placeholder"].container():
+                col1, col2 = st.session_state["columns_placeholder"].columns([2, 1])
+                with col1:
+                    st.session_state["planner_placeholder"].markdown(
+                        st.session_state["responses"][-1][-1]
+                    )
+                with col2:
+                    selected_node_key = st.session_state["already_used_keys"][-1]
+                    st.session_state["selectbox_placeholder"] = st.empty()
+                    selected_node = st.session_state["selectbox_placeholder"].selectbox(
+                        "Select a node:",
+                        list(st.session_state["nodes_list"][i].keys()),
+                        key=f"replay_key_{i}",
+                        index=list(st.session_state["nodes_list"][i].keys()).index(
+                            st.session_state[selected_node_key]
+                        ),
+                    )
+                    st.session_state[selected_node_key] = selected_node
+                    if (
+                        selected_node not in ["root", "response"]
+                        and selected_node in st.session_state["nodes_list"][i]
+                    ):
+                        node_info_key = f"{selected_node}_info"
+                        for item in st.session_state[node_info_key]:
+                            if item[0] in ["thought", "answer"]:
+                                st.session_state["searcher_placeholder"] = st.empty()
+                                st.session_state["searcher_placeholder"].markdown(item[1])
+                            elif item[0] == "observation":
+                                st.session_state["observation_expander"] = st.empty()
+                                with st.session_state["observation_expander"].expander("Results"):
+                                    st.write(item[1])
+                        # st.session_state['searcher_placeholder'].markdown(st.session_state[node_info_key])
+
+
+def clean_history():
+    st.session_state["queries"] = []
+    st.session_state["responses"] = []
+    st.session_state["graphs_html"] = []
+    st.session_state["nodes_list"] = []
+    st.session_state["adjacency_list_list"] = []
+    st.session_state["history"] = []
+    st.session_state["already_used_keys"] = list()
+    for k in st.session_state:
+        if k.endswith("placeholder") or k.endswith("_info"):
+            del st.session_state[k]
+
+
+# Main function to run the Streamlit app
+def main():
+    st.sidebar.title("Model Control")
+    col1, col2 = st.columns([4, 1])
+    with col1:
+        user_input = st.chat_input("Enter your query:")
+    with col2:
+        if st.button("Clear History"):
+            clean_history()
+    if user_input:
+        update_chat(user_input)
+    display_chat_history()
+
+
+import os
+
+if __name__ == "__main__":
+    os.system("pip show starlette")
+    # os.system("pip install -r requirements.txt")
+    os.system("pip install tenacity")
+    os.system("python -m mindsearch.app --lang en --model_format internlm_silicon --search_engine GoogleSearch &")
+
+    main()

mindsearch/agent/__init__.py

@@ -0,0 +1,82 @@
+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

mindsearch/agent/graph.py

@@ -0,0 +1,307 @@
+import asyncio
+import queue
+import random
+import re
+import uuid
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from copy import deepcopy
+from threading import Thread
+from typing import Dict, List
+
+from lagent.actions import BaseAction
+from lagent.schema import AgentMessage, AgentStatusCode
+
+from .streaming import AsyncStreamingAgentForInternLM, StreamingAgentForInternLM
+
+
+class SearcherAgent(StreamingAgentForInternLM):
+    def __init__(
+        self,
+        user_input_template: str = "{question}",
+        user_context_template: str = None,
+        **kwargs,
+    ):
+        self.user_input_template = user_input_template
+        self.user_context_template = user_context_template
+        super().__init__(**kwargs)
+
+    def forward(
+        self,
+        question: str,
+        topic: str,
+        history: List[dict] = None,
+        session_id=0,
+        **kwargs,
+    ):
+        message = [self.user_input_template.format(question=question, topic=topic)]
+        if history and self.user_context_template:
+            message = [self.user_context_template.format_map(item) for item in history] + message
+        message = "\n".join(message)
+        return super().forward(message, session_id=session_id, **kwargs)
+
+
+class AsyncSearcherAgent(AsyncStreamingAgentForInternLM):
+    def __init__(
+        self,
+        user_input_template: str = "{question}",
+        user_context_template: str = None,
+        **kwargs,
+    ):
+        self.user_input_template = user_input_template
+        self.user_context_template = user_context_template
+        super().__init__(**kwargs)
+
+    async def forward(
+        self,
+        question: str,
+        topic: str,
+        history: List[dict] = None,
+        session_id=0,
+        **kwargs,
+    ):
+        message = [self.user_input_template.format(question=question, topic=topic)]
+        if history and self.user_context_template:
+            message = [self.user_context_template.format_map(item) for item in history] + message
+        message = "\n".join(message)
+        async for message in super().forward(message, session_id=session_id, **kwargs):
+            yield message
+
+
+class WebSearchGraph:
+    is_async = False
+    SEARCHER_CONFIG = {}
+    _SEARCHER_LOOP = []
+    _SEARCHER_THREAD = []
+
+    def __init__(self):
+        self.nodes: Dict[str, Dict[str, str]] = {}
+        self.adjacency_list: Dict[str, List[dict]] = defaultdict(list)
+        self.future_to_query = dict()
+        self.searcher_resp_queue = queue.Queue()
+        self.executor = ThreadPoolExecutor(max_workers=10)
+        self.n_active_tasks = 0
+
+    def add_root_node(
+        self,
+        node_content: str,
+        node_name: str = "root",
+    ):
+        """添加起始节点
+
+        Args:
+            node_content (str): 节点内容
+            node_name (str, optional): 节点名称. Defaults to 'root'.
+
+        """
+        self.nodes[node_name] = dict(content=node_content, type="root")
+        self.adjacency_list[node_name] = []
+
+    def add_node(
+        self,
+        node_name: str,
+        node_content: str,
+    ):
+        """添加搜索子问题节点
+
+        Args:
+            node_name (str): 节点名称
+            node_content (str): 子问题内容
+
+        Returns:
+            str: 返回搜索结果
+        """
+        self.nodes[node_name] = dict(content=node_content, type="searcher")
+        self.adjacency_list[node_name] = []
+
+        parent_nodes = []
+        for start_node, adj in self.adjacency_list.items():
+            for neighbor in adj:
+                if (
+                    node_name == neighbor
+                    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
+        ]
+
+        if self.is_async:
+
+            async def _async_search_node_stream():
+                cfg = {
+                    **self.SEARCHER_CONFIG,
+                    "plugins": deepcopy(self.SEARCHER_CONFIG.get("plugins")),
+                }
+                agent, session_id = AsyncSearcherAgent(**cfg), random.randint(0, 999999)
+                searcher_message = AgentMessage(sender="SearcherAgent", content="")
+                try:
+                    async for searcher_message in agent(
+                        question=node_content,
+                        topic=self.nodes["root"]["content"],
+                        history=parent_response,
+                        session_id=session_id,
+                    ):
+                        self.nodes[node_name]["response"] = searcher_message.model_dump()
+                        self.nodes[node_name]["memory"] = agent.state_dict(session_id=session_id)
+                        self.nodes[node_name]["session_id"] = session_id
+                        self.searcher_resp_queue.put((node_name, self.nodes[node_name], []))
+                    self.searcher_resp_queue.put((None, None, None))
+                except Exception as exc:
+                    self.searcher_resp_queue.put((exc, None, None))
+
+            self.future_to_query[
+                asyncio.run_coroutine_threadsafe(
+                    _async_search_node_stream(), random.choice(self._SEARCHER_LOOP)
+                )
+            ] = f"{node_name}-{node_content}"
+            # self.future_to_query[
+            #     self.executor.submit(asyncio.run, _async_search_node_stream())
+            # ] = f"{node_name}-{node_content}"
+        else:
+
+            def _search_node_stream():
+                cfg = {
+                    **self.SEARCHER_CONFIG,
+                    "plugins": deepcopy(self.SEARCHER_CONFIG.get("plugins")),
+                }
+                agent, session_id = SearcherAgent(**cfg), random.randint(0, 999999)
+                searcher_message = AgentMessage(sender="SearcherAgent", content="")
+                try:
+                    for searcher_message in agent(
+                        question=node_content,
+                        topic=self.nodes["root"]["content"],
+                        history=parent_response,
+                        session_id=session_id,
+                    ):
+                        self.nodes[node_name]["response"] = searcher_message.model_dump()
+                        self.nodes[node_name]["memory"] = agent.state_dict(session_id=session_id)
+                        self.nodes[node_name]["session_id"] = session_id
+                        self.searcher_resp_queue.put((node_name, self.nodes[node_name], []))
+                    self.searcher_resp_queue.put((None, None, None))
+                except Exception as exc:
+                    self.searcher_resp_queue.put((exc, None, None))
+
+            self.future_to_query[
+                self.executor.submit(_search_node_stream)
+            ] = f"{node_name}-{node_content}"
+
+        self.n_active_tasks += 1
+
+    def add_response_node(self, node_name="response"):
+        """添加回复节点
+
+        Args:
+            thought (str): 思考过程
+            node_name (str, optional): 节点名称. Defaults to '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: str, end_node: str):
+        """添加边
+
+        Args:
+            start_node (str): 起始节点名称
+            end_node (str): 结束节点名称
+        """
+        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: str) -> str:
+        return self.nodes[node_name].copy()
+
+    @classmethod
+    def start_loop(cls, n: int = 32):
+        if not cls.is_async:
+            raise RuntimeError("Event loop cannot be launched as `is_async` is disabled")
+
+        assert len(cls._SEARCHER_LOOP) == len(cls._SEARCHER_THREAD)
+        for i, (loop, thread) in enumerate(
+            zip(cls._SEARCHER_LOOP.copy(), cls._SEARCHER_THREAD.copy())
+        ):
+            if not (loop.is_running() and thread.is_alive()):
+                cls._SEARCHER_LOOP.pop(i)
+                cls._SEARCHER_THREAD.pop(i)
+
+        while len(cls._SEARCHER_THREAD) < n:
+
+            def _start_loop():
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+                cls._SEARCHER_LOOP.append(loop)
+                loop.run_forever()
+
+            thread = Thread(target=_start_loop, daemon=True)
+            thread.start()
+            cls._SEARCHER_THREAD.append(thread)
+
+
+class ExecutionAction(BaseAction):
+    """Tool used by MindSearch planner to execute graph node query."""
+
+    def run(self, command, local_dict, global_dict, stream_graph=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
+
+        command = extract_code(command)
+        exec(command, global_dict, local_dict)
+
+        # 匹配所有 graph.node 中的内容
+        node_list = re.findall(r"graph.node\((.*?)\)", command)
+        graph: WebSearchGraph = local_dict["graph"]
+        while graph.n_active_tasks:
+            while not graph.searcher_resp_queue.empty():
+                node_name, _, _ = graph.searcher_resp_queue.get(timeout=60)
+                if isinstance(node_name, Exception):
+                    raise node_name
+                if node_name is None:
+                    graph.n_active_tasks -= 1
+                    continue
+                if stream_graph:
+                    for neighbors in graph.adjacency_list.values():
+                        for neighbor in neighbors:
+                            # state  1进行中，2未开始，3已结束
+                            if not (
+                                neighbor["name"] in graph.nodes
+                                and "response" in graph.nodes[neighbor["name"]]
+                            ):
+                                neighbor["state"] = 2
+                            elif (
+                                graph.nodes[neighbor["name"]]["response"]["stream_state"]
+                                == AgentStatusCode.END
+                            ):
+                                neighbor["state"] = 3
+                            else:
+                                neighbor["state"] = 1
+                    if all(
+                        "response" in node
+                        for name, node in graph.nodes.items()
+                        if name not in ["root", "response"]
+                    ):
+                        yield AgentMessage(
+                            sender=self.name,
+                            content=dict(current_node=node_name),
+                            formatted=dict(
+                                node=deepcopy(graph.nodes),
+                                adjacency_list=deepcopy(graph.adjacency_list),
+                            ),
+                            stream_state=AgentStatusCode.STREAM_ING,
+                        )
+        res = [graph.nodes[node.strip().strip('"').strip("'")] for node in node_list]
+        return res, graph.nodes, graph.adjacency_list

mindsearch/agent/mindsearch_agent.py

@@ -0,0 +1,210 @@
+import json
+import logging
+import re
+from copy import deepcopy
+from typing import Dict, Tuple
+
+from lagent.schema import AgentMessage, AgentStatusCode, ModelStatusCode
+from lagent.utils import GeneratorWithReturn
+
+from .graph import ExecutionAction, WebSearchGraph
+from .streaming import AsyncStreamingAgentForInternLM, StreamingAgentForInternLM
+
+
+def _update_ref(ref: str, ref2url: Dict[str, str], ptr: int) -> str:
+    numbers = list({int(n) for n in re.findall(r"\[\[(\d+)\]\]", ref)})
+    numbers = {n: idx + 1 for idx, n in enumerate(numbers)}
+    updated_ref = re.sub(
+        r"\[\[(\d+)\]\]",
+        lambda match: f"[[{numbers[int(match.group(1))] + ptr}]]",
+        ref,
+    )
+    updated_ref2url = {}
+    if numbers:
+        try:
+            assert all(elem in ref2url for elem in numbers)
+        except Exception as exc:
+            logging.info(f"Illegal reference id: {str(exc)}")
+        if ref2url:
+            updated_ref2url = {
+                numbers[idx] + ptr: ref2url[idx] for idx in numbers if idx in ref2url
+            }
+    return updated_ref, updated_ref2url, len(numbers) + 1
+
+
+def _generate_references_from_graph(graph: Dict[str, dict]) -> Tuple[str, Dict[int, dict]]:
+    ptr, references, references_url = 0, [], {}
+    for name, data_item in graph.items():
+        if name in ["root", "response"]:
+            continue
+        # only search once at each node, thus the result offset is 2
+        assert data_item["memory"]["agent.memory"][2]["sender"].endswith("ActionExecutor")
+        ref2url = {
+            int(k): v
+            for k, v in json.loads(data_item["memory"]["agent.memory"][2]["content"]).items()
+        }
+        updata_ref, ref2url, added_ptr = _update_ref(
+            data_item["response"]["content"], ref2url, ptr
+        )
+        ptr += added_ptr
+        references.append(f'## {data_item["content"]}\n\n{updata_ref}')
+        references_url.update(ref2url)
+    return "\n\n".join(references), references_url
+
+
+class MindSearchAgent(StreamingAgentForInternLM):
+    def __init__(
+        self,
+        searcher_cfg: dict,
+        summary_prompt: str,
+        finish_condition=lambda m: "add_response_node" in m.content,
+        max_turn: int = 10,
+        **kwargs,
+    ):
+        WebSearchGraph.SEARCHER_CONFIG = searcher_cfg
+        super().__init__(finish_condition=finish_condition, max_turn=max_turn, **kwargs)
+        self.summary_prompt = summary_prompt
+        self.action = ExecutionAction()
+
+    def forward(self, message: AgentMessage, session_id=0, **kwargs):
+        if isinstance(message, str):
+            message = AgentMessage(sender="user", content=message)
+        _graph_state = dict(node={}, adjacency_list={}, ref2url={})
+        local_dict, global_dict = {}, globals()
+        for _ in range(self.max_turn):
+            last_agent_state = AgentStatusCode.SESSION_READY
+            for message in self.agent(message, session_id=session_id, **kwargs):
+                if isinstance(message.formatted, dict) and message.formatted.get("tool_type"):
+                    if message.stream_state == ModelStatusCode.END:
+                        message.stream_state = last_agent_state + int(
+                            last_agent_state
+                            in [
+                                AgentStatusCode.CODING,
+                                AgentStatusCode.PLUGIN_START,
+                            ]
+                        )
+                    else:
+                        message.stream_state = (
+                            AgentStatusCode.PLUGIN_START
+                            if message.formatted["tool_type"] == "plugin"
+                            else AgentStatusCode.CODING
+                        )
+                else:
+                    message.stream_state = AgentStatusCode.STREAM_ING
+                message.formatted.update(deepcopy(_graph_state))
+                yield message
+                last_agent_state = message.stream_state
+            if not message.formatted["tool_type"]:
+                message.stream_state = AgentStatusCode.END
+                yield message
+                return
+
+            gen = GeneratorWithReturn(
+                self.action.run(message.content, local_dict, global_dict, True)
+            )
+            for graph_exec in gen:
+                graph_exec.formatted["ref2url"] = deepcopy(_graph_state["ref2url"])
+                yield graph_exec
+
+            reference, references_url = _generate_references_from_graph(gen.ret[1])
+            _graph_state.update(node=gen.ret[1], adjacency_list=gen.ret[2], ref2url=references_url)
+            if self.finish_condition(message):
+                message = AgentMessage(
+                    sender="ActionExecutor",
+                    content=self.summary_prompt,
+                    formatted=deepcopy(_graph_state),
+                    stream_state=message.stream_state + 1,  # plugin or code return
+                )
+                yield message
+                # summarize the references to generate the final answer
+                for message in self.agent(message, session_id=session_id, **kwargs):
+                    message.formatted.update(deepcopy(_graph_state))
+                    yield message
+                return
+            message = AgentMessage(
+                sender="ActionExecutor",
+                content=reference,
+                formatted=deepcopy(_graph_state),
+                stream_state=message.stream_state + 1,  # plugin or code return
+            )
+            yield message
+
+
+class AsyncMindSearchAgent(AsyncStreamingAgentForInternLM):
+    def __init__(
+        self,
+        searcher_cfg: dict,
+        summary_prompt: str,
+        finish_condition=lambda m: "add_response_node" in m.content,
+        max_turn: int = 10,
+        **kwargs,
+    ):
+        WebSearchGraph.SEARCHER_CONFIG = searcher_cfg
+        WebSearchGraph.is_async = True
+        WebSearchGraph.start_loop()
+        super().__init__(finish_condition=finish_condition, max_turn=max_turn, **kwargs)
+        self.summary_prompt = summary_prompt
+        self.action = ExecutionAction()
+
+    async def forward(self, message: AgentMessage, session_id=0, **kwargs):
+        if isinstance(message, str):
+            message = AgentMessage(sender="user", content=message)
+        _graph_state = dict(node={}, adjacency_list={}, ref2url={})
+        local_dict, global_dict = {}, globals()
+        for _ in range(self.max_turn):
+            last_agent_state = AgentStatusCode.SESSION_READY
+            async for message in self.agent(message, session_id=session_id, **kwargs):
+                if isinstance(message.formatted, dict) and message.formatted.get("tool_type"):
+                    if message.stream_state == ModelStatusCode.END:
+                        message.stream_state = last_agent_state + int(
+                            last_agent_state
+                            in [
+                                AgentStatusCode.CODING,
+                                AgentStatusCode.PLUGIN_START,
+                            ]
+                        )
+                    else:
+                        message.stream_state = (
+                            AgentStatusCode.PLUGIN_START
+                            if message.formatted["tool_type"] == "plugin"
+                            else AgentStatusCode.CODING
+                        )
+                else:
+                    message.stream_state = AgentStatusCode.STREAM_ING
+                message.formatted.update(deepcopy(_graph_state))
+                yield message
+                last_agent_state = message.stream_state
+            if not message.formatted["tool_type"]:
+                message.stream_state = AgentStatusCode.END
+                yield message
+                return
+
+            gen = GeneratorWithReturn(
+                self.action.run(message.content, local_dict, global_dict, True)
+            )
+            for graph_exec in gen:
+                graph_exec.formatted["ref2url"] = deepcopy(_graph_state["ref2url"])
+                yield graph_exec
+
+            reference, references_url = _generate_references_from_graph(gen.ret[1])
+            _graph_state.update(node=gen.ret[1], adjacency_list=gen.ret[2], ref2url=references_url)
+            if self.finish_condition(message):
+                message = AgentMessage(
+                    sender="ActionExecutor",
+                    content=self.summary_prompt,
+                    formatted=deepcopy(_graph_state),
+                    stream_state=message.stream_state + 1,  # plugin or code return
+                )
+                yield message
+                # summarize the references to generate the final answer
+                async for message in self.agent(message, session_id=session_id, **kwargs):
+                    message.formatted.update(deepcopy(_graph_state))
+                    yield message
+                return
+            message = AgentMessage(
+                sender="ActionExecutor",
+                content=reference,
+                formatted=deepcopy(_graph_state),
+                stream_state=message.stream_state + 1,  # plugin or code return
+            )
+            yield message

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,115 @@
+import os
+
+from dotenv import load_dotenv
+from lagent.llms import (
+    GPTAPI,
+    INTERNLM2_META,
+    HFTransformerCasualLM,
+    LMDeployClient,
+    LMDeployServer,
+)
+
+internlm_server = dict(
+    type=LMDeployServer,
+    path="internlm/internlm2_5-7b-chat",
+    model_name="internlm2_5-7b-chat",
+    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"),
+    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"),
+    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"),
+    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|>"],
+)
+
+
+internlm_api = dict(
+    type=GPTAPI,
+    model_type="internlm2.5-latest",
+    key=os.environ.get("InternLM_API_KEY", "YOUR InternLM API KEY https://internlm.intern-ai.org.cn/api/document"),
+    api_base="https://internlm-chat.intern-ai.org.cn/puyu/api/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/agent/streaming.py

@@ -0,0 +1,203 @@
+import copy
+from typing import List, Union
+
+from lagent.agents import Agent, AgentForInternLM, AsyncAgent, AsyncAgentForInternLM
+from lagent.schema import AgentMessage, AgentStatusCode, ModelStatusCode
+
+
+class StreamingAgentMixin:
+    """Make agent calling output a streaming response."""
+
+    def __call__(self, *message: Union[AgentMessage, List[AgentMessage]], session_id=0, **kwargs):
+        for hook in self._hooks.values():
+            message = copy.deepcopy(message)
+            result = hook.before_agent(self, message, session_id)
+            if result:
+                message = result
+        self.update_memory(message, session_id=session_id)
+        response_message = AgentMessage(sender=self.name, content="")
+        for response_message in self.forward(*message, session_id=session_id, **kwargs):
+            if not isinstance(response_message, AgentMessage):
+                model_state, response = response_message
+                response_message = AgentMessage(
+                    sender=self.name,
+                    content=response,
+                    stream_state=model_state,
+                )
+            yield response_message.model_copy()
+        self.update_memory(response_message, session_id=session_id)
+        for hook in self._hooks.values():
+            response_message = response_message.model_copy(deep=True)
+            result = hook.after_agent(self, response_message, session_id)
+            if result:
+                response_message = result
+        yield response_message
+
+
+class AsyncStreamingAgentMixin:
+    """Make asynchronous agent calling output a streaming response."""
+
+    async def __call__(
+        self, *message: Union[AgentMessage, List[AgentMessage]], session_id=0, **kwargs
+    ):
+        for hook in self._hooks.values():
+            message = copy.deepcopy(message)
+            result = hook.before_agent(self, message, session_id)
+            if result:
+                message = result
+        self.update_memory(message, session_id=session_id)
+        response_message = AgentMessage(sender=self.name, content="")
+        async for response_message in self.forward(*message, session_id=session_id, **kwargs):
+            if not isinstance(response_message, AgentMessage):
+                model_state, response = response_message
+                response_message = AgentMessage(
+                    sender=self.name,
+                    content=response,
+                    stream_state=model_state,
+                )
+            yield response_message.model_copy()
+        self.update_memory(response_message, session_id=session_id)
+        for hook in self._hooks.values():
+            response_message = response_message.model_copy(deep=True)
+            result = hook.after_agent(self, response_message, session_id)
+            if result:
+                response_message = result
+        yield response_message
+
+
+class StreamingAgent(StreamingAgentMixin, Agent):
+    """Base streaming agent class"""
+
+    def forward(self, *message: AgentMessage, session_id=0, **kwargs):
+        formatted_messages = self.aggregator.aggregate(
+            self.memory.get(session_id),
+            self.name,
+            self.output_format,
+            self.template,
+        )
+        for model_state, response, _ in self.llm.stream_chat(
+            formatted_messages, session_id=session_id, **kwargs
+        ):
+            yield AgentMessage(
+                sender=self.name,
+                content=response,
+                formatted=self.output_format.parse_response(response),
+                stream_state=model_state,
+            ) if self.output_format else (model_state, response)
+
+
+class AsyncStreamingAgent(AsyncStreamingAgentMixin, AsyncAgent):
+    """Base asynchronous streaming agent class"""
+
+    async def forward(self, *message: AgentMessage, session_id=0, **kwargs):
+        formatted_messages = self.aggregator.aggregate(
+            self.memory.get(session_id),
+            self.name,
+            self.output_format,
+            self.template,
+        )
+        async for model_state, response, _ in self.llm.stream_chat(
+            formatted_messages, session_id=session_id, **kwargs
+        ):
+            yield AgentMessage(
+                sender=self.name,
+                content=response,
+                formatted=self.output_format.parse_response(response),
+                stream_state=model_state,
+            ) if self.output_format else (model_state, response)
+
+
+class StreamingAgentForInternLM(StreamingAgentMixin, AgentForInternLM):
+    """Streaming implementation of `lagent.agents.AgentForInternLM`"""
+
+    _INTERNAL_AGENT_CLS = StreamingAgent
+
+    def forward(self, message: AgentMessage, session_id=0, **kwargs):
+        if isinstance(message, str):
+            message = AgentMessage(sender="user", content=message)
+        for _ in range(self.max_turn):
+            last_agent_state = AgentStatusCode.SESSION_READY
+            for message in self.agent(message, session_id=session_id, **kwargs):
+                if isinstance(message.formatted, dict) and message.formatted.get("tool_type"):
+                    if message.stream_state == ModelStatusCode.END:
+                        message.stream_state = last_agent_state + int(
+                            last_agent_state
+                            in [
+                                AgentStatusCode.CODING,
+                                AgentStatusCode.PLUGIN_START,
+                            ]
+                        )
+                    else:
+                        message.stream_state = (
+                            AgentStatusCode.PLUGIN_START
+                            if message.formatted["tool_type"] == "plugin"
+                            else AgentStatusCode.CODING
+                        )
+                else:
+                    message.stream_state = AgentStatusCode.STREAM_ING
+                yield message
+                last_agent_state = message.stream_state
+            if self.finish_condition(message):
+                message.stream_state = AgentStatusCode.END
+                yield message
+                return
+            if message.formatted["tool_type"]:
+                tool_type = message.formatted["tool_type"]
+                executor = getattr(self, f"{tool_type}_executor", None)
+                if not executor:
+                    raise RuntimeError(f"No available {tool_type} executor")
+                tool_return = executor(message, session_id=session_id)
+                tool_return.stream_state = message.stream_state + 1
+                message = tool_return
+                yield message
+            else:
+                message.stream_state = AgentStatusCode.STREAM_ING
+                yield message
+
+
+class AsyncStreamingAgentForInternLM(AsyncStreamingAgentMixin, AsyncAgentForInternLM):
+    """Streaming implementation of `lagent.agents.AsyncAgentForInternLM`"""
+
+    _INTERNAL_AGENT_CLS = AsyncStreamingAgent
+
+    async def forward(self, message: AgentMessage, session_id=0, **kwargs):
+        if isinstance(message, str):
+            message = AgentMessage(sender="user", content=message)
+        for _ in range(self.max_turn):
+            last_agent_state = AgentStatusCode.SESSION_READY
+            async for message in self.agent(message, session_id=session_id, **kwargs):
+                if isinstance(message.formatted, dict) and message.formatted.get("tool_type"):
+                    if message.stream_state == ModelStatusCode.END:
+                        message.stream_state = last_agent_state + int(
+                            last_agent_state
+                            in [
+                                AgentStatusCode.CODING,
+                                AgentStatusCode.PLUGIN_START,
+                            ]
+                        )
+                    else:
+                        message.stream_state = (
+                            AgentStatusCode.PLUGIN_START
+                            if message.formatted["tool_type"] == "plugin"
+                            else AgentStatusCode.CODING
+                        )
+                else:
+                    message.stream_state = AgentStatusCode.STREAM_ING
+                yield message
+                last_agent_state = message.stream_state
+            if self.finish_condition(message):
+                message.stream_state = AgentStatusCode.END
+                yield message
+                return
+            if message.formatted["tool_type"]:
+                tool_type = message.formatted["tool_type"]
+                executor = getattr(self, f"{tool_type}_executor", None)
+                if not executor:
+                    raise RuntimeError(f"No available {tool_type} executor")
+                tool_return = await executor(message, session_id=session_id)
+                tool_return.stream_state = message.stream_state + 1
+                message = tool_return
+                yield message
+            else:
+                message.stream_state = AgentStatusCode.STREAM_ING
+                yield message

mindsearch/app.py

@@ -0,0 +1,176 @@
+import asyncio
+import json
+import logging
+import random
+from typing import Dict, List, Union
+
+import janus
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.requests import Request
+from pydantic import BaseModel, Field
+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("--host", default="0.0.0.0", type=str, help="Service host")
+    parser.add_argument("--port", default=8002, type=int, help="Service port")
+    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="BingSearch", type=str, help="Search engine")
+    parser.add_argument("--asy", default=False, action="store_true", help="Agent mode")
+    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]]
+    session_id: int = Field(default_factory=lambda: random.randint(0, 999999))
+    agent_cfg: Dict = dict()
+
+
+def _postprocess_agent_message(message: dict) -> dict:
+    content, fmt = message["content"], message["formatted"]
+    current_node = content["current_node"] if isinstance(content, dict) else None
+    if current_node:
+        message["content"] = None
+        for key in ["ref2url"]:
+            fmt.pop(key, None)
+        graph = fmt["node"]
+        for key in graph.copy():
+            if key != current_node:
+                graph.pop(key)
+        if current_node not in ["root", "response"]:
+            node = graph[current_node]
+            for key in ["memory", "session_id"]:
+                node.pop(key, None)
+            node_fmt = node["response"]["formatted"]
+            if isinstance(node_fmt, dict) and "thought" in node_fmt and "action" in node_fmt:
+                node["response"]["content"] = None
+                node_fmt["thought"] = (
+                    node_fmt["thought"] and node_fmt["thought"].split("<|action_start|>")[0]
+                )
+                if isinstance(node_fmt["action"], str):
+                    node_fmt["action"] = node_fmt["action"].split("<|action_end|>")[0]
+    else:
+        if isinstance(fmt, dict) and "thought" in fmt and "action" in fmt:
+            message["content"] = None
+            fmt["thought"] = fmt["thought"] and fmt["thought"].split("<|action_start|>")[0]
+            if isinstance(fmt["action"], str):
+                fmt["action"] = fmt["action"].split("<|action_end|>")[0]
+        for key in ["node"]:
+            fmt.pop(key, None)
+    return dict(current_node=current_node, response=message)
+
+
+async def run(request: GenerationParams, _request: Request):
+    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(inputs, session_id=session_id):
+                        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
+                stop_event.set()  # Inform sync_generator_wrapper to stop
+
+            async for message in async_generator_wrapper():
+                response_json = json.dumps(
+                    _postprocess_agent_message(message.model_dump()),
+                    ensure_ascii=False,
+                )
+                yield {"data": response_json}
+                if await _request.is_disconnected():
+                    break
+        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}
+        finally:
+            await stop_event.wait()  # Waiting for async_generator_wrapper to stop
+            queue.close()
+            await queue.wait_closed()
+            agent.agent.memory.memory_map.pop(session_id, None)
+
+    inputs = request.inputs
+    session_id = request.session_id
+    agent = init_agent(
+        lang=args.lang,
+        model_format=args.model_format,
+        search_engine=args.search_engine,
+    )
+    return EventSourceResponse(generate(), ping=300)
+
+
+async def run_async(request: GenerationParams, _request: Request):
+    async def generate():
+        try:
+            async for message in agent(inputs, session_id=session_id):
+                response_json = json.dumps(
+                    _postprocess_agent_message(message.model_dump()),
+                    ensure_ascii=False,
+                )
+                yield {"data": response_json}
+                if await _request.is_disconnected():
+                    break
+        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}
+        finally:
+            agent.agent.memory.memory_map.pop(session_id, None)
+
+    inputs = request.inputs
+    session_id = request.session_id
+    agent = init_agent(
+        lang=args.lang,
+        model_format=args.model_format,
+        search_engine=args.search_engine,
+        use_async=True,
+    )
+    return EventSourceResponse(generate(), ping=300)
+
+
+app.add_api_route("/solve", run_async if args.asy else run, methods=["POST"])
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

mindsearch/terminal.py

@@ -0,0 +1,66 @@
+import os
+import sys
+from datetime import datetime
+
+from lagent.actions import WebBrowser
+from lagent.agents.stream import get_plugin_prompt
+from lagent.llms import INTERNLM2_META, LMDeployServer
+from lagent.prompts import InterpreterParser, PluginParser
+
+from mindsearch.agent.mindsearch_agent import MindSearchAgent
+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"
+date = datetime.now().strftime("The current date is %Y-%m-%d.")
+llm = LMDeployServer(
+    path="internlm/internlm2_5-7b-chat",
+    model_name="internlm2",
+    meta_template=INTERNLM2_META,
+    top_p=0.8,
+    top_k=1,
+    temperature=1.0,
+    max_new_tokens=8192,
+    repetition_penalty=1.02,
+    stop_words=["<|im_end|>", "<|action_end|>"],
+)
+plugins = [WebBrowser(searcher_type="BingSearch", topk=6)]
+agent = 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,
+)
+
+for agent_return in agent("上海今天适合穿什么衣服"):
+    pass
+
+print(agent_return.sender)
+print(agent_return.content)
+print(agent_return.formatted["ref2url"])

requirements.txt

@@ -0,0 +1,18 @@
+duckduckgo_search==5.3.1b1
+einops
+fastapi
+gradio==5.7.1
+janus
+lagent==0.5.0rc2
+matplotlib
+pydantic==2.6.4
+python-dotenv
+pyvis
+schemdraw
+sse-starlette
+termcolor
+transformers==4.41.0
+uvicorn
+tenacity
+streamlit
+git+https://github.com/vansin/ms_gradio_agentchatbot.git