import mathfrom collections import dequefrom typing import Optionalfrom langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
class Node:def __init__(self,messages: list[BaseMessage],reflection: Reflection,parent: Optional[Node] = None,):self.messages = messagesself.parent = parentself.children = []self.value = 0self.visits = 0self.reflection = reflectionself.depth = parent.depth + 1 if parent is not None else 1self._is_solved = reflection.found_solution if reflection else Falseif self._is_solved:self._mark_tree_as_solved()self.backpropagate(reflection.normalized_score)
def __repr__(self) -> str:return (f"<Node value={self.value}, visits={self.visits},"f" solution={self.messages} reflection={self.reflection}/>")
@propertydef is_solved(self):"""If any solutions exist, we can end the search."""return self._is_solved
@propertydef is_terminal(self):return not self.children
@propertydef best_child(self):"""Select the child with the highest UCT to search next."""if not self.children:return Noneall_nodes = self._get_all_children()return max(all_nodes, key=lambda child: child.upper_confidence_bound())
@propertydef best_child_score(self):"""Return the child with the highest value."""if not self.children:return Nonereturn max(self.children, key=lambda child: int(child.is_solved) * child.value)
@propertydef height(self) -> int:"""Check for how far we've rolled out the tree."""if self.children:return 1 + max([child.height for child in self.children])return 1
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 is None:raise ValueError("Cannot obtain UCT from root node")if self.visits == 0:return self.value# Encourages exploitation of high-value trajectoriesaverage_reward = self.value / self.visits# Encourages exploration of less-visited trajectoriesexploration_term = math.sqrt(math.log(self.parent.visits) / self.visits)return average_reward + exploration_weight * exploration_term
def backpropagate(self, reward: float):"""Update the score of this node and its parents."""node = selfwhile node:node.visits += 1node.value = (node.value * (node.visits - 1) + reward) / node.visitsnode = node.parent
def get_messages(self, include_reflections: bool = True):if include_reflections:return self.messages + [self.reflection.as_message()]return self.messages
def get_trajectory(self, include_reflections: bool = True) -> list[BaseMessage]:"""Get messages representing this search branch."""messages = []node = selfwhile node:messages.extend(node.get_messages(include_reflections=include_reflections)[::-1])node = node.parent# Reverse the final back-tracked trajectory to return in the correct orderreturn messages[::-1]# root solution, reflection, child 1, ...
def _get_all_children(self):all_nodes = []nodes = deque()nodes.append(self)while nodes:node = nodes.popleft()all_nodes.extend(node.children)for n in node.children:nodes.append(n)return all_nodes
def get_best_solution(self):"""Return the best solution from within the current sub-tree."""all_nodes = [self] + self._get_all_children()best_node = max(all_nodes,# We filter out all non-terminal, non-solution trajectorieskey=lambda node: int(node.is_terminal and node.is_solved) * node.value,)return best_node
def _mark_tree_as_solved(self):parent = self.parentwhile parent:parent._is_solved = Trueparent = parent.parent

from langchain_community.tools.tavily_search import TavilySearchResultsfrom langchain_community.utilities.tavily_search import TavilySearchAPIWrapperfrom langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
search = TavilySearchAPIWrapper()tavily_tool = TavilySearchResults(api_wrapper=search, max_results=5)tools = [tavily_tool]tool_executor = ToolExecutor(tools=tools)

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 as_message(self):return HumanMessage(content=f"Reasoning: {self.reflections}\nScore: {self.score}")
@propertydef normalized_score(self) -> float:return self.score / 10.0
prompt = ChatPromptTemplate.from_messages([("system","Reflect and grade the assistant response to the user question below.",),("user", "{input}"),MessagesPlaceholder(variable_name="candidate"),])
reflection_llm_chain = (prompt| llm.bind_tools(tools=[Reflection], tool_choice="Reflection").with_config(run_name="Reflection")| PydanticToolsParser(tools=[Reflection]))
@as_runnabledef reflection_chain(inputs) -> Reflection:tool_choices = reflection_llm_chain.invoke(inputs)reflection = tool_choices[0]if not isinstance(inputs["candidate"][-1], AIMessage):reflection.found_solution = Falsereturn reflection

from langchain_core.prompt_values import ChatPromptValuefrom langchain_core.runnables import RunnableConfig
prompt_template = ChatPromptTemplate.from_messages([("system","You are an AI assistant.",),("user", "{input}"),MessagesPlaceholder(variable_name="messages", optional=True),])
initial_answer_chain = prompt_template | llm.bind_tools(tools=tools).with_config(run_name="GenerateInitialCandidate")
parser = JsonOutputToolsParser(return_id=True)initial_response = initial_answer_chain.invoke({"input": "Write a research report on lithium pollution."})

import json
# Define the node we will add to the graphdef generate_initial_response(state: TreeState) -> dict:"""Generate the initial candidate response."""res = initial_answer_chain.invoke({"input": state["input"]})parsed = parser.invoke(res)tool_responses = tool_executor.batch([ToolInvocation(tool=r["type"], tool_input=r["args"]) for r in parsed])output_messages = [res] + [ToolMessage(content=json.dumps(resp), tool_call_id=tool_call["id"])for resp, tool_call in zip(tool_responses, parsed)]reflection = reflection_chain.invoke({"input": state["input"], "candidate": output_messages})root = Node(output_messages, reflection=reflection)return {**state,"root": root,}

def generate_candidates(messages: ChatPromptValue, config: RunnableConfig):n = config["configurable"].get("N", 5)bound_kwargs = llm.bind_tools(tools=tools).kwargschat_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]]
expansion_chain = prompt_template | generate_candidatesres = expansion_chain.invoke({"input": "Write a research report on lithium pollution."})

from collections import defaultdict
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 rootmessages = best_candidate.get_trajectory()# Generate N candidates from the single child candidatenew_candidates = expansion_chain.invoke({"input": state["input"], "messages": messages}, config)parsed = parser.batch(new_candidates)flattened = [(i, tool_call)for i, tool_calls in enumerate(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 in zip(flattened, tool_responses):collected_responses[i].append(ToolMessage(content=json.dumps(resp), tool_call_id=tool_call["id"]))output_messages = []for i, candidate in enumerate(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 treechild_nodes = [Node(cand, parent=best_candidate, reflection=reflection)for cand, reflection in zip(output_messages, reflections)]best_candidate.children.extend(child_nodes)# We have already extended the tree directly, so we just return the statereturn state

from typing import Literal
from langgraph.graph import END, StateGraph, START
def should_loop(state: TreeState) -> Literal["expand", "__end__"]:"""Determine whether to continue the tree search."""root = state["root"]if root.is_solved:return ENDif root.height > 5:return ENDreturn "expand"

builder = StateGraph(TreeState)builder.add_node("start", generate_initial_response)builder.add_node("expand", expand)builder.add_edge(START, "start")
builder.add_conditional_edges("start",# Either expand/rollout or finishshould_loop,)builder.add_conditional_edges("expand",# Either continue to rollout or finishshould_loop,)
graph = builder.compile()