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