第 4 阶段：Tool Calling

本阶段我们将 AI 节点与工具执行节点连接起来，形成一个循环。

• @tool 装饰器: 将 Python 函数快速声明为可供 LLM 识别的工具。
• ToolNode: LangGraph 预置的节点，专门用于执行模型发出的工具调用请求。
• ReAct 循环: agent -> tools -> agent。模型先思考，如果需要工具则跳转到 tools，tools 执行完后将结果传回给 agent，模型再总结。

测试场景：

1. 输入 "北京天气怎么样？"
2. 输入 "iPhone 多少钱？"
3. 输入 "你好" (观察不触发工具的情况)

import json
from typing import Annotated, List, Literal, TypedDict, Union

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.tools import tool

from langchain_test.models import qwen


# 1. 定义工具 (Mock 外部系统)
@tool
def get_weather(city: str):
    """获取指定城市的实时天气情况。"""
    # 模拟外部 API 调用
    if "北京" in city:
        return "北京天气晴朗，气温 15°C。"
    elif "上海" in city:
        return "上海有小雨，气温 18°C。"
    else:
        return f"抱歉，暂时无法获取 {city} 的天气信息。"


@tool
def search_product_price(product_name: str):
    """查询产品的最新价格。"""
    prices = {
        "iphone": "8000 元",
        "macbook": "15000 元",
        "ipad": "4000 元"
    }
    for key, value in prices.items():
        if key in product_name.lower():
            return f"{product_name} 的价格是 {value}。"
    return f"未找到 {product_name} 的价格信息。"


tools = [get_weather, search_product_price]
tool_node = ToolNode(tools)


# 2. 定义状态
class MessagesState(TypedDict):
    messages: Annotated[List[BaseMessage], add_messages]


# 3. 初始化模型并绑定工具
# 注意：qwen 模型需要支持 tool calling
llm = qwen.bind_tools(tools)


# 4. 定义节点逻辑
def call_model(state: MessagesState):
    """调用模型决定是否需要使用工具"""
    response = llm.invoke(state["messages"])
    return {"messages": [response]}


def should_continue(state: MessagesState) -> Literal["tools", END]:
    """路由逻辑：检查最后一条消息是否包含工具调用请求"""
    last_message = state["messages"][-1]
    if last_message.tool_calls:
        return "tools"
    return END


# 5. 构建图 (ReAct 架构)
workflow = StateGraph(MessagesState)

# 添加节点
workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)

# 设置入口
workflow.set_entry_point("agent")

# 添加条件边：agent -> tools 或 agent -> END
workflow.add_conditional_edges(
    "agent",
    should_continue,
)

# 工具执行完后，必须回到 agent 节点，让模型根据工具结果生成最终回答
workflow.add_edge("tools", "agent")

graph = workflow.compile()

# --- FastAPI 接口 ---

app = FastAPI(title="LangGraph Lesson 4 Backend")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class ChatMessage(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    messages: List[ChatMessage]


class ChatResponse(BaseModel):
    reply: str
    # 返回执行过程中的消息列表，方便前端展示工具调用细节
    history: List[dict]


@app.post("/api/chat")
async def chat(request: ChatRequest):
    history = []
    for msg in request.messages:
        if msg.role == "user":
            history.append(HumanMessage(content=msg.content))
        elif msg.role == "assistant":
            history.append(AIMessage(content=msg.content))

    # 运行图
    result = graph.invoke({"messages": history})

    # 格式化历史消息发送给前端
    formatted_history = []
    for m in result["messages"]:
        if isinstance(m, HumanMessage):
            formatted_history.append({"role": "user", "content": m.content})
        elif isinstance(m, AIMessage):
            content = m.content
            if m.tool_calls:
                # 如果有工具调用，记录下来
                tool_info = [f"调用工具: {tc['name']}({tc['args']})" for tc in m.tool_calls]
                content += "\n" + "\n".join(tool_info)
            formatted_history.append({"role": "assistant", "content": content})
        elif isinstance(m, ToolMessage):
            formatted_history.append({"role": "system", "content": f"[工具返回]: {m.content}"})

    last_message = result["messages"][-1]
    return ChatResponse(
        reply=last_message.content,
        history=formatted_history
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)