🔧 工具系统详解

"""
@tool 装饰器基础示例
功能：创建一个简单的数据库搜索工具
"""
from langchain.tools import tool

@tool
def search_database(query: str, limit: int = 10) -> str:
    """在客户数据库中搜索匹配查询的记录。

    Args:
        query: 要查找的搜索词
        limit: 返回的最大结果数量

    Returns:
        搜索结果的描述
    """
    # 模拟数据库搜索
    results = [
        f"客户 {i+1}: 匹配 '{query}'"
        for i in range(min(limit, 5))
    ]
    return f"找到 {len(results)} 条结果：\n" + "\n".join(results)

# 查看工具属性
print(f"工具名称: {search_database.name}")
print(f"工具描述: {search_database.description}")
print(f"工具参数: {search_database.args}")

"""
自定义工具名称和描述
"""
from langchain.tools import tool

# 自定义工具名称
@tool("web_search")
def search(query: str) -> str:
    """在网络上搜索信息。"""
    return f"'{query}' 的搜索结果..."

print(search.name)  # 输出: web_search


# 自定义名称和描述
@tool(
    "calculator",
    description="执行算术计算。用于任何数学问题。"
)
def calc(expression: str) -> str:
    """计算数学表达式。"""
    try:
        result = eval(expression)  # 仅用于演示
        return f"结果: {result}"
    except Exception as e:
        return f"计算错误: {str(e)}"

print(calc.name)         # 输出: calculator
print(calc.description)  # 输出: 执行算术计算。用于任何数学问题。

"""
使用 Pydantic 定义复杂工具 Schema
功能：创建一个功能完善的天气查询工具
"""
from pydantic import BaseModel, Field
from typing import Literal
from langchain.tools import tool

class WeatherInput(BaseModel):
    """天气查询的输入参数。"""
    location: str = Field(
        description="城市名称或坐标"
    )
    units: Literal["celsius", "fahrenheit"] = Field(
        default="celsius",
        description="温度单位偏好"
    )
    include_forecast: bool = Field(
        default=False,
        description="是否包含 5 天预报"
    )

@tool(args_schema=WeatherInput)
def get_weather(
    location: str,
    units: str = "celsius",
    include_forecast: bool = False
) -> str:
    """获取当前天气和可选的天气预报。

    Args:
        location: 城市名称
        units: 温度单位
        include_forecast: 是否包含预报

    Returns:
        天气信息描述
    """
    # 模拟天气数据
    temp = 22 if units == "celsius" else 72
    result = f"{location} 的当前天气：{temp}° {units[0].upper()}"

    if include_forecast:
        result += "\n未来 5 天：晴朗"

    return result

# 测试工具
print(get_weather.invoke({
    "location": "北京",
    "units": "celsius",
    "include_forecast": True
}))

"""
访问对话状态示例
功能：工具可以读取对话历史
"""
from langchain.tools import tool, ToolRuntime

@tool
def summarize_conversation(runtime: ToolRuntime) -> str:
    """总结到目前为止的对话。

    Returns:
        对话统计摘要
    """
    # 访问对话状态中的消息列表
    messages = runtime.state["messages"]

    # 统计不同类型的消息
    human_msgs = sum(
        1 for m in messages
        if m.__class__.__name__ == "HumanMessage"
    )
    ai_msgs = sum(
        1 for m in messages
        if m.__class__.__name__ == "AIMessage"
    )
    tool_msgs = sum(
        1 for m in messages
        if m.__class__.__name__ == "ToolMessage"
    )

    return (
        f"对话包含 {human_msgs} 条用户消息、"
        f"{ai_msgs} 条 AI 响应和 "
        f"{tool_msgs} 条工具结果"
    )

@tool
def get_user_preference(pref_name: str, runtime: ToolRuntime) -> str:
    """获取用户偏好设置的值。

    Args:
        pref_name: 偏好设置名称

    Returns:
        偏好设置的值
    """
    # 从状态中读取用户偏好
    preferences = runtime.state.get("user_preferences", {})
    return preferences.get(pref_name, "未设置")

"""
自定义上下文访问示例
功能：工具可以访问用户特定的上下文信息
"""
from dataclasses import dataclass
from langchain.tools import tool, ToolRuntime
from langchain.chat_models import init_chat_model
from langchain.agents import create_agent

# 模拟用户数据库
USER_DATABASE = {
    "user123": {
        "name": "张三",
        "account_type": "VIP",
        "balance": 5000,
        "email": "[email protected]"
    },
    "user456": {
        "name": "李四",
        "account_type": "普通",
        "balance": 1200,
        "email": "[email protected]"
    }
}

# 定义用户上下文结构
@dataclass
class UserContext:
    """用户上下文信息"""
    user_id: str

# 创建使用上下文的工具
@tool
def get_account_info(runtime: ToolRuntime[UserContext]) -> str:
    """获取当前用户的账户信息。

    Returns:
        用户账户详情
    """
    # 从上下文获取用户 ID
    user_id = runtime.context.user_id

    # 查询用户数据
    if user_id in USER_DATABASE:
        user = USER_DATABASE[user_id]
        return (
            f"账户持有人：{user['name']}\n"
            f"账户类型：{user['account_type']}\n"
            f"余额：¥{user['balance']}"
        )
    return "用户未找到"

@tool
def check_vip_status(runtime: ToolRuntime[UserContext]) -> str:
    """检查当前用户的 VIP 状态。

    Returns:
        VIP 状态信息
    """
    user_id = runtime.context.user_id

    if user_id in USER_DATABASE:
        user = USER_DATABASE[user_id]
        is_vip = user["account_type"] == "VIP"
        return f"VIP 状态：{'是' if is_vip else '否'}"
    return "用户未找到"

# 创建 Agent
model = init_chat_model("gpt-4o")
agent = create_agent(
    model,
    tools=[get_account_info, check_vip_status],
    context_schema=UserContext,
    system_prompt="你是一个金融助手。"
)

# 调用 Agent（传入用户上下文）
result = agent.invoke(
    {"messages": [{"role": "user", "content": "我的当前余额是多少？"}]},
    context=UserContext(user_id="user123")  # 注入用户上下文
)

print(result["messages"][-1].content)

"""
状态更新示例
功能：工具可以修改 Agent 的运行时状态
"""
from langchain.tools import tool, ToolRuntime
from langgraph.types import Command
from langchain.messages import RemoveMessage
from langgraph.graph.message import REMOVE_ALL_MESSAGES

@tool
def clear_conversation() -> Command:
    """清除对话历史。

    Returns:
        状态更新命令
    """
    return Command(
        update={
            "messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)],
        }
    )

@tool
def update_user_name(new_name: str, runtime: ToolRuntime) -> Command:
    """更新用户姓名。

    Args:
        new_name: 新的用户姓名

    Returns:
        状态更新命令
    """
    return Command(
        update={"user_name": new_name}
    )

@tool
def set_language_preference(language: str) -> Command:
    """设置语言偏好。

    Args:
        language: 语言代码（如 'zh', 'en'）

    Returns:
        状态更新命令
    """
    return Command(
        update={"language": language}
    )

"""
Memory/Store 访问示例
功能：持久化用户数据
"""
from typing import Any
from langchain.tools import tool, ToolRuntime
from langgraph.store.memory import InMemoryStore
from langchain.agents import create_agent
from langchain.chat_models import init_chat_model

@tool
def get_user_info(user_id: str, runtime: ToolRuntime) -> str:
    """查找用户信息。

    Args:
        user_id: 用户 ID

    Returns:
        用户信息
    """
    store = runtime.store
    user_info = store.get(("users",), user_id)
    return str(user_info.value) if user_info else "未知用户"

@tool
def save_user_info(
    user_id: str,
    user_info: dict[str, Any],
    runtime: ToolRuntime
) -> str:
    """保存用户信息。

    Args:
        user_id: 用户 ID
        user_info: 用户信息字典

    Returns:
        成功消息
    """
    store = runtime.store
    store.put(("users",), user_id, user_info)
    return "成功保存用户信息。"

# 创建存储
store = InMemoryStore()

# 创建 Agent
model = init_chat_model("gpt-4o")
agent = create_agent(
    model,
    tools=[get_user_info, save_user_info],
    store=store
)

# 第一次会话：保存用户信息
agent.invoke({
    "messages": [{
        "role": "user",
        "content": "保存以下用户：userid: abc123, name: 王五, age: 25, email: [email protected]"
    }]
})

# 第二次会话：检索持久化数据
result = agent.invoke({
    "messages": [{
        "role": "user",
        "content": "获取用户 ID 为 'abc123' 的用户信息"
    }]
})

print(result["messages"][-1].content)

"""
Stream Writer 实时反馈示例
功能：在工具执行过程中发送进度更新
"""
from langchain.tools import tool, ToolRuntime
import time

@tool
def get_weather(city: str, runtime: ToolRuntime) -> str:
    """获取指定城市的天气。

    Args:
        city: 城市名称

    Returns:
        天气信息
    """
    writer = runtime.stream_writer

    # 发送进度更新
    writer(f"正在查找 {city} 的数据...")
    time.sleep(1)  # 模拟 API 调用

    writer(f"已获取 {city} 的数据")
    time.sleep(0.5)

    writer(f"正在解析天气信息...")
    time.sleep(0.5)

    # 返回最终结果
    return f"{city} 今天总是晴朗的！"

@tool
def search_documents(query: str, runtime: ToolRuntime) -> str:
    """搜索文档库。

    Args:
        query: 搜索查询

    Returns:
        搜索结果
    """
    writer = runtime.stream_writer

    writer(f"开始搜索：'{query}'")
    time.sleep(0.5)

    writer(f"正在索引中查找...")
    time.sleep(1)

    writer(f"找到 3 个相关文档")
    time.sleep(0.5)

    writer(f"正在提取内容...")

    return f"关于 '{query}' 的搜索结果：[文档摘要]"

# ❌ 错误：不能使用保留参数名
@tool
def bad_tool(query: str, config: dict) -> str:
    """这会导致错误！"""
    return "..."

# ✅ 正确：使用 runtime 访问配置
@tool
def good_tool(query: str, runtime: ToolRuntime) -> str:
    """使用 runtime 访问所需信息"""
    # 可以通过 runtime 访问配置和状态
    return "..."

# ❌ 不好：文档字符串不清晰
@tool
def search(q: str) -> str:
    """搜索"""
    return "..."

# ✅ 好：清晰的文档字符串
@tool
def search_database(query: str, limit: int = 10) -> str:
    """在客户数据库中搜索记录。

    使用此工具查找客户信息、订单历史或产品详情。
    支持模糊匹配和多关键词搜索。

    Args:
        query: 搜索关键词（客户名称、订单号、产品名称等）
        limit: 返回的最大结果数量（默认 10）

    Returns:
        搜索结果的格式化字符串
    """
    return "..."

# ❌ 不好：抛出异常
@tool
def divide(a: float, b: float) -> float:
    """除法运算"""
    return a / b  # 如果 b=0 会抛出异常

# ✅ 好：返回清晰的错误消息
@tool
def divide(a: float, b: float) -> str:
    """执行除法运算。

    Args:
        a: 被除数
        b: 除数

    Returns:
        计算结果或错误消息
    """
    if b == 0:
        return "错误：除数不能为零。请提供非零的除数。"

    try:
        result = a / b
        return f"计算结果：{a} ÷ {b} = {result}"
    except Exception as e:
        return f"计算失败：{str(e)}"

from pydantic import BaseModel, Field, validator
from typing import Literal

class EmailInput(BaseModel):
    """发送邮件的输入参数。"""
    to: str = Field(description="收件人邮箱地址")
    subject: str = Field(description="邮件主题")
    body: str = Field(description="邮件正文")
    priority: Literal["low", "normal", "high"] = Field(
        default="normal",
        description="邮件优先级"
    )

    @validator("to")
    def validate_email(cls, v):
        """验证邮箱格式"""
        if "@" not in v:
            raise ValueError("无效的邮箱地址")
        return v

@tool(args_schema=EmailInput)
def send_email(to: str, subject: str, body: str, priority: str = "normal") -> str:
    """发送邮件。"""
    return f"已发送邮件到 {to}"

"""
订单管理工具集完整示例
功能：一组完整的订单管理工具
"""
from dataclasses import dataclass
from typing import Literal
from pydantic import BaseModel, Field
from langchain.tools import tool, ToolRuntime
from langchain.agents import create_agent
from langchain.chat_models import init_chat_model

# 模拟订单数据库
ORDERS = {
    "ORD001": {"product": "笔记本电脑", "status": "已发货", "amount": 5999},
    "ORD002": {"product": "鼠标", "status": "处理中", "amount": 99},
}

# 用户上下文
@dataclass
class UserContext:
    user_id: str
    role: str  # 'customer' 或 'admin'

# 查询订单工具
@tool
def query_order(order_id: str) -> str:
    """查询订单状态。

    Args:
        order_id: 订单编号

    Returns:
        订单详情或错误消息
    """
    if order_id in ORDERS:
        order = ORDERS[order_id]
        return (
            f"订单 {order_id}：\n"
            f"  产品：{order['product']}\n"
            f"  状态：{order['status']}\n"
            f"  金额：¥{order['amount']}"
        )
    return f"未找到订单 {order_id}"

# 取消订单工具（需要权限检查）
@tool
def cancel_order(order_id: str, runtime: ToolRuntime[UserContext]) -> str:
    """取消订单。

    Args:
        order_id: 要取消的订单编号

    Returns:
        操作结果
    """
    # 权限检查
    user_role = runtime.context.role
    if user_role not in ["customer", "admin"]:
        return "权限不足：无法取消订单"

    if order_id not in ORDERS:
        return f"未找到订单 {order_id}"

    order = ORDERS[order_id]
    if order["status"] == "已发货":
        return f"订单 {order_id} 已发货，无法取消。请联系客服处理退货。"

    # 取消订单
    order["status"] = "已取消"
    return f"订单 {order_id} 已成功取消。"

# 更新订单状态工具（仅管理员）
class OrderStatusInput(BaseModel):
    """更新订单状态的输入参数。"""
    order_id: str = Field(description="订单编号")
    new_status: Literal["处理中", "已发货", "已完成", "已取消"] = Field(
        description="新的订单状态"
    )

@tool(args_schema=OrderStatusInput)
def update_order_status(
    order_id: str,
    new_status: str,
    runtime: ToolRuntime[UserContext]
) -> str:
    """更新订单状态（仅管理员）。

    Args:
        order_id: 订单编号
        new_status: 新状态

    Returns:
        操作结果
    """
    # 仅管理员可用
    if runtime.context.role != "admin":
        return "权限不足：此操作仅限管理员"

    if order_id not in ORDERS:
        return f"未找到订单 {order_id}"

    ORDERS[order_id]["status"] = new_status
    return f"订单 {order_id} 状态已更新为：{new_status}"

# 创建 Agent
model = init_chat_model("gpt-4o")

customer_agent = create_agent(
    model,
    tools=[query_order, cancel_order],
    context_schema=UserContext,
    system_prompt="你是客服助手，帮助用户管理订单。"
)

admin_agent = create_agent(
    model,
    tools=[query_order, cancel_order, update_order_status],
    context_schema=UserContext,
    system_prompt="你是管理员助手，可以管理所有订单。"
)

# 客户场景
print("=== 客户场景 ===")
result1 = customer_agent.invoke(
    {"messages": [{"role": "user", "content": "查询订单 ORD001 的状态"}]},
    context=UserContext(user_id="customer1", role="customer")
)
print(result1["messages"][-1].content)

# 管理员场景
print("\n=== 管理员场景 ===")
result2 = admin_agent.invoke(
    {"messages": [{"role": "user", "content": "将订单 ORD002 状态更新为已发货"}]},
    context=UserContext(user_id="admin1", role="admin")
)
print(result2["messages"][-1].content)

# 直接调用工具测试
result = search_database.invoke({"query": "测试", "limit": 5})
print(result)

# 使用 pytest 编写单元测试
def test_search_database():
    result = search_database.invoke({"query": "客户", "limit": 3})
    assert "找到" in result
    assert "客户" in result

import asyncio
from langchain.tools import tool

@tool
async def slow_operation(query: str) -> str:
    """执行耗时操作"""
    try:
        # 设置超时
        result = await asyncio.wait_for(
            some_async_function(query),
            timeout=10.0  # 10 秒超时
        )
        return result
    except asyncio.TimeoutError:
        return "操作超时，请稍后重试"