builder.add_conditional_edges( "start", # Either expand/rollout or finish should_loop, ) builder.add_conditional_edges( "expand", # Either continue to rollout or finish should_loop, )
（5）编译图
graph = builder.compile()
这就定义好了一个LangGraph，图的执行路径是这样的：
其中 should_loop 函数的定义：
def should_loop(state: TreeState): """Determine whether to continue the tree search.""" root = state["root"] if root.is_solved: return END if root.height > 5: return END return"expand"
prompt = ChatPromptTemplate.from_messages( [ ( "system", "Reflect and grade the assistant response to the user question below.", ), ("user", "{input}"), MessagesPlaceholder(variable_name="candidate"), ] )
class Reflection(BaseModel): reflections: str = Field( description="The critique and reflections on the sufficiency, superfluency," " and general quality of the response" ) score: int = Field( description="Score from 0-10 on the quality of the candidate response.", gte=0, lte=10, ) found_solution: bool = Field( description="Whether the response has fully solved the question or task." )
def expand(state: TreeState, config: RunnableConfig) -> dict: """Starting from the "best" node in the tree, generate N candidates for the next step.""" root = state["root"] best_candidate: Node = root.best_child if root.children else root messages = best_candidate.get_trajectory() # Generate N candidates from the single child candidate new_candidates = expansion_chain.invoke( {"input": state["input"], "messages": messages}, config ) parsed = parser.batch(new_candidates) flattened = [ (i, tool_call) for i, tool_calls inenumerate(parsed) for tool_call in tool_calls ] tool_responses = tool_executor.batch( [ ToolInvocation(tool=tool_call["type"], tool_input=tool_call["args"]) for _, tool_call in flattened ] ) collected_responses = defaultdict(list) for (i, tool_call), resp inzip(flattened, tool_responses): collected_responses[i].append( ToolMessage(content=json.dumps(resp), tool_call_id=tool_call["id"]) ) output_messages = [] for i, candidate inenumerate(new_candidates): output_messages.append([candidate] + collected_responses[i])
# Reflect on each candidate # For tasks with external validation, you'd add that here. reflections = reflection_chain.batch( [{"input": state["input"], "candidate": msges} for msges in output_messages], config, ) # Grow tree child_nodes = [ Node(cand, parent=best_candidate, reflection=reflection) for cand, reflection inzip(output_messages, reflections) ] best_candidate.children.extend(child_nodes) # We have already extended the tree directly, so we just return the state return state
1.3.2.1 选择 best_candidate
选择当前最优的节点。
以下代码从 best_child开始看，首先是获取树的全部节点和子孙节点。取分数最高的节点。
怎么取分数最高的节点？这里调用了 upper_confidence_bound函数。这个函数用来计算UCT分数，UCT 是一种常用于多臂赌博机问题（Multi-Armed Bandit Problem）和蒙特卡洛树搜索（Monte Carlo Tree Search, MCTS）的算法，它有助于在探索（exploration）和利用（exploitation）之间找到一个平衡。
def upper_confidence_bound(self, exploration_weight=1.0): """Return the UCT score. This helps balance exploration vs. exploitation of a branch.""" if self.parent isNone: raise ValueError("Cannot obtain UCT from root node") if self.visits == 0: return self.value # Encourages exploitation of high-value trajectories average_reward = self.value / self.visits # Encourages exploration of less-visited trajectories exploration_term = math.sqrt(math.log(self.parent.visits) / self.visits) return average_reward + exploration_weight * exploration_term
@property def best_child(self): """Select the child with the highest UCT to search next.""" ifnot self.children: returnNone all_nodes = self._get_all_children() returnmax(all_nodes, key=lambda child: child.upper_confidence_bound())
# This generates N candidate values for a single input to sample actions from the environment def generate_candidates(messages: ChatPromptValue, config: RunnableConfig): n = config["configurable"].get("N", 5) bound_kwargs = llm.bind_tools(tools=tools).kwargs chat_result = llm.generate( [messages.to_messages()], n=n, callbacks=config["callbacks"], run_name="GenerateCandidates", **bound_kwargs ) return [gen.message for gen in chat_result.generations[0]]
def backpropagate(self, reward: float): """Update the score of this node and its parents.""" node = self while node: node.visits += 1 node.value = (node.value * (node.visits - 1) + reward) / node.visits node = node.parent
1.4 执行
question = "Write a research report on lithium pollution." for step in graph.stream({"input": question}): step_name, step_state = next(iter(step.items())) print(step_name) print("rolled out: ", step_state["root"].height) print("---")
粤ICP备14082021号
免费POC，
零成本试错
