【AI Agent系列】【LangGraph】1. 进阶实战：给你的 LangGraph 加入条件分支（Conditional edges）

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, BaseMessage
from langgraph.graph import END, MessageGraph
import json
from langchain_core.messages import ToolMessage
from langchain_core.tools import tool
from langchain_core.utils.function_calling import convert_to_openai_tool
from typing import List
@tool
def multiply(first_number: int, second_number: int):
    """Multiplies two numbers together."""
    return first_number * second_number
model = ChatOpenAI(temperature=0)
model_with_tools = model.bind(tools=[convert_to_openai_tool(multiply)])
graph = MessageGraph()
def invoke_model(state: List[BaseMessage]):
    return model_with_tools.invoke(state)
graph.add_node("oracle", invoke_model)
def invoke_tool(state: List[BaseMessage]):
    tool_calls = state[-1].additional_kwargs.get("tool_calls", [])
    multiply_call = None
    for tool_call in tool_calls:
        if tool_call.get("function").get("name") == "multiply":
            multiply_call = tool_call
    if multiply_call is None:
        raise Exception("No adder input found.")
    res = multiply.invoke(
        json.loads(multiply_call.get("function").get("arguments"))
    )
    return ToolMessage(
        tool_call_id=multiply_call.get("id"),
        content=res
    )
graph.add_node("multiply", invoke_tool)
graph.add_edge("multiply", END)
graph.set_entry_point("oracle")
def router(state: List[BaseMessage]):
    tool_calls = state[-1].additional_kwargs.get("tool_calls", [])
    if len(tool_calls):
        return "multiply"
    else:
        return "end"
graph.add_conditional_edges("oracle", router, {
    "multiply": "multiply",
    "end": END,
})
runnable = graph.compile()
response = runnable.invoke(HumanMessage("What is 123 * 456?"))
print(response)

graph.add_conditional_edges("oracle", router, {
    "multiply": "multiply",
    "end": END,
})

def add_conditional_edges(
    self,
    start_key: str,
    condition: Callable[..., str],
    conditional_edge_mapping: Optional[Dict[str, str]] = None,
) -> None:
    if self.compiled:
        logger.warning(
            "Adding an edge to a graph that has already been compiled. This will "
            "not be reflected in the compiled graph."
        )
    if start_key not in self.nodes:
        raise ValueError(f"Need to add_node `{start_key}` first")
    if iscoroutinefunction(condition):
        raise ValueError("Condition cannot be a coroutine function")
    if conditional_edge_mapping and set(
        conditional_edge_mapping.values()
    ).difference([END]).difference(self.nodes):
        raise ValueError(
            f"Missing nodes which are in conditional edge mapping. Mapping "
            f"contains possible destinations: "
            f"{list(conditional_edge_mapping.values())}. Possible nodes are "
            f"{list(self.nodes.keys())}."
        )
    self.branches[start_key].append(Branch(condition, conditional_edge_mapping))

def router(state: List[BaseMessage]):
    tool_calls = state[-1].additional_kwargs.get("tool_calls", [])
    if len(tool_calls):
        return "multiply"
    else:
        return "end"

@tool
def multiply(first_number: int, second_number: int):
    """Multiplies two numbers together."""
    return first_number * second_number
model = ChatOpenAI(temperature=0)
model_with_tools = model.bind(tools=[convert_to_openai_tool(multiply)])
graph = MessageGraph()
def invoke_model(state: List[BaseMessage]):
    return model_with_tools.invoke(state)
graph.add_node("oracle", invoke_model)

def invoke_tool(state: List[BaseMessage]):
    tool_calls = state[-1].additional_kwargs.get("tool_calls", [])
    multiply_call = None
    for tool_call in tool_calls:
        if tool_call.get("function").get("name") == "multiply":
            multiply_call = tool_call
    if multiply_call is None:
        raise Exception("No adder input found.")
    res = multiply.invoke(
        json.loads(multiply_call.get("function").get("arguments"))
    )
    return ToolMessage(
        tool_call_id=multiply_call.get("id"),
        content=res
    )
graph.add_node("multiply", invoke_tool)