144 lines
4.0 KiB
Python
144 lines
4.0 KiB
Python
"""LangGraph 工具调用示例：天气查询 + 计算器。
|
||
|
||
图结构：
|
||
|
||
START -> model
|
||
| 有 tool_calls
|
||
v
|
||
tools
|
||
|
|
||
+---------> model
|
||
|
||
model -- 无 tool_calls --> END
|
||
|
||
天气数据是本地模拟数据，重点是学习 LangGraph 的工具调用循环。
|
||
"""
|
||
|
||
import os
|
||
from typing import Literal
|
||
|
||
from dotenv import load_dotenv
|
||
from langchain_core.tools import tool
|
||
from langchain_openai import ChatOpenAI
|
||
from langgraph.graph import END, START, MessagesState, StateGraph
|
||
from langgraph.prebuilt import ToolNode
|
||
|
||
load_dotenv()
|
||
|
||
|
||
# 1. 定义工具。类型注解和 docstring 会变成提供给模型的工具 schema。
|
||
@tool
|
||
def get_weather(city: str) -> str:
|
||
"""查询指定城市当前的天气。城市名称使用中文，例如北京、上海。"""
|
||
weather_data = {
|
||
"北京": "晴，25°C，北风 2 级",
|
||
"上海": "多云，28°C，东南风 3 级",
|
||
"广州": "雷阵雨，30°C，南风 2 级",
|
||
"深圳": "阴，29°C，东风 2 级",
|
||
}
|
||
return weather_data.get(city, f"暂时没有 {city} 的天气数据")
|
||
|
||
|
||
@tool
|
||
def calculator(a: float, b: float, operation: Literal["add", "subtract", "multiply", "divide"]) -> float:
|
||
"""对两个数字进行加、减、乘、除运算。
|
||
|
||
Args:
|
||
a: 第一个数字。
|
||
b: 第二个数字。
|
||
operation: 运算类型：add、subtract、multiply 或 divide。
|
||
"""
|
||
if operation == "add":
|
||
return a + b
|
||
if operation == "subtract":
|
||
return a - b
|
||
if operation == "multiply":
|
||
return a * b
|
||
if b == 0:
|
||
raise ValueError("除数不能为 0")
|
||
return a / b
|
||
|
||
|
||
# 2. 将工具绑定给模型。绑定后，模型可以在 AIMessage.tool_calls 中请求工具。
|
||
tools = [get_weather, calculator]
|
||
model = ChatOpenAI(
|
||
model=os.getenv("OPENAI_MODEL", "openai/gpt-5.4"),
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
base_url=os.getenv("OPENAI_BASE_URL", "https://api.qnaigc.com/v1"),
|
||
temperature=0,
|
||
max_retries=2,
|
||
)
|
||
model_with_tools = model.bind_tools(tools)
|
||
|
||
|
||
# 3. 模型节点：读取完整消息历史，返回一条新的 AIMessage。
|
||
def call_model(state: MessagesState) -> dict:
|
||
response = model_with_tools.invoke(state["messages"])
|
||
return {"messages": [response]}
|
||
|
||
|
||
# 4. 路由函数：检查模型是否发起了工具调用。
|
||
def route_after_model(state: MessagesState) -> Literal["tools", "__end__"]:
|
||
last_message = state["messages"][-1]
|
||
if last_message.tool_calls:
|
||
return "tools"
|
||
return END
|
||
|
||
|
||
# 5. ToolNode 自动完成：找到工具、传入参数、执行工具、生成 ToolMessage。
|
||
tool_node = ToolNode(tools)
|
||
|
||
builder = StateGraph(MessagesState)
|
||
builder.add_node("model", call_model)
|
||
builder.add_node("tools", tool_node)
|
||
|
||
builder.add_edge(START, "model")
|
||
builder.add_conditional_edges(
|
||
"model",
|
||
route_after_model,
|
||
{
|
||
"tools": "tools",
|
||
END: END,
|
||
},
|
||
)
|
||
builder.add_edge("tools", "model")
|
||
|
||
graph = builder.compile()
|
||
|
||
|
||
def ask(question: str) -> str:
|
||
"""调用助手，并打印每个节点产生的更新。"""
|
||
final_answer = ""
|
||
|
||
for event in graph.stream(
|
||
{"messages": [{"role": "user", "content": question}]},
|
||
stream_mode="updates",
|
||
):
|
||
for node_name, update in event.items():
|
||
print(f"\n[{node_name}]")
|
||
for message in update.get("messages", []):
|
||
message.pretty_print()
|
||
if node_name == "model" and not message.tool_calls:
|
||
final_answer = str(message.content)
|
||
|
||
return final_answer
|
||
|
||
|
||
def main() -> None:
|
||
print("=== 天气 + 计算器工具助手 ===")
|
||
print("输入 quit 退出。")
|
||
|
||
while True:
|
||
question = input("\n用户：").strip()
|
||
if question.lower() in {"quit", "exit", "q"}:
|
||
break
|
||
if not question:
|
||
continue
|
||
|
||
answer = ask(question)
|
||
print(f"\n最终回答：{answer}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|