使用按需技能构建 SQL 助手

SKILLS: list[Skill] = [
    {
        "name": "sales_analytics",
        "description": "用于销售数据分析的数据库 schema 和业务逻辑，包括 customers、orders 和 revenue。",
        "content": """# 销售分析 Schema

## 表

### customers
- customer_id (PRIMARY KEY)
- name
- email
- signup_date
- status (active/inactive)
- customer_tier (bronze/silver/gold/platinum)

### orders
- order_id (PRIMARY KEY)
- customer_id (FOREIGN KEY -> customers)
- order_date
- status (pending/completed/cancelled/refunded)
- total_amount
- sales_region (north/south/east/west)

### order_items
- item_id (PRIMARY KEY)
- order_id (FOREIGN KEY -> orders)
- product_id
- quantity
- unit_price
- discount_percent

## 业务逻辑

**活跃客户**：status = 'active' AND signup_date <= CURRENT_DATE - INTERVAL '90 days'

**收入计算**：仅统计 status = 'completed' 的订单。使用 orders 表中的 total_amount，它已经包含折扣计算。

**客户生命周期价值 (CLV)**：某个客户所有已完成订单金额的总和。

**高价值订单**：total_amount > 1000 的订单

## 示例查询

-- 获取上一季度收入最高的 10 个客户
SELECT
    c.customer_id,
    c.name,
    c.customer_tier,
    SUM(o.total_amount) as total_revenue
FROM customers c
JOIN orders o ON c.customer_id = o.customer_id
WHERE o.status = 'completed'
  AND o.order_date >= CURRENT_DATE - INTERVAL '3 months'
GROUP BY c.customer_id, c.name, c.customer_tier
ORDER BY total_revenue DESC
LIMIT 10;
""",
    },
    {
        "name": "inventory_management",
        "description": "用于库存跟踪的数据库 schema 和业务逻辑，包括 products、warehouses 和 stock levels。",
        "content": """# 库存管理 Schema

## 表

### products
- product_id (PRIMARY KEY)
- product_name
- sku
- category
- unit_cost
- reorder_point (重新订购前的最低库存水平)
- discontinued (boolean)

### warehouses
- warehouse_id (PRIMARY KEY)
- warehouse_name
- location
- capacity

### inventory
- inventory_id (PRIMARY KEY)
- product_id (FOREIGN KEY -> products)
- warehouse_id (FOREIGN KEY -> warehouses)
- quantity_on_hand
- last_updated

### stock_movements
- movement_id (PRIMARY KEY)
- product_id (FOREIGN KEY -> products)
- warehouse_id (FOREIGN KEY -> warehouses)
- movement_type (inbound/outbound/transfer/adjustment)
- quantity (inbound 为正数，outbound 为负数)
- movement_date
- reference_number

## 业务逻辑

**可用库存**：inventory 表中 quantity_on_hand > 0 的 quantity_on_hand

**需要重新订购的产品**：所有仓库中的 total quantity_on_hand 小于或等于产品 reorder_point 的产品

**仅活跃产品**：除非专门分析已停用商品，否则排除 discontinued = true 的产品

**库存估值**：每个产品的 quantity_on_hand * unit_cost

## 示例查询

-- 查找所有仓库中低于重新订购点的产品
SELECT
    p.product_id,
    p.product_name,
    p.reorder_point,
    SUM(i.quantity_on_hand) as total_stock,
    p.unit_cost,
    (p.reorder_point - SUM(i.quantity_on_hand)) as units_to_reorder
FROM products p
JOIN inventory i ON p.product_id = i.product_id
WHERE p.discontinued = false
GROUP BY p.product_id, p.product_name, p.reorder_point, p.unit_cost
HAVING SUM(i.quantity_on_hand) <= p.reorder_point
ORDER BY units_to_reorder DESC;
""",
    },
]

from langchain.agents.middleware import AgentState

class CustomState(AgentState):
    skills_loaded: NotRequired[list[str]]  # 跟踪已加载的技能  #

from langgraph.types import Command  
from langchain.tools import tool, ToolRuntime
from langchain.messages import ToolMessage  

@tool
def load_skill(skill_name: str, runtime: ToolRuntime) -> Command:
    """将技能的完整内容加载到智能体上下文中。

    当你需要有关如何处理特定类型请求的详细信息时使用此工具。
    它会提供该技能领域的完整指令、策略和指南。

    Args:
        skill_name: 要加载的技能名称
    """
    # 查找并返回请求的技能
    for skill in SKILLS:
        if skill["name"] == skill_name:
            skill_content = f"已加载技能：{skill_name}\n\n{skill['content']}"

            # 更新状态以跟踪已加载技能
            return Command(
                update={
                    "messages": [
                        ToolMessage(
                            content=skill_content,
                            tool_call_id=runtime.tool_call_id,
                        )
                    ],
                    "skills_loaded": [skill_name],
                }
            )

    # 未找到技能
    available = ", ".join(s["name"] for s in SKILLS)
    return Command(
        update={
            "messages": [
                ToolMessage(
                    content=f"未找到技能 '{skill_name}'。可用技能：{available}",
                    tool_call_id=runtime.tool_call_id,
                )
            ]
        }
    )

@tool
def write_sql_query(
    query: str,
    vertical: str,
    runtime: ToolRuntime,
) -> str:
    """为特定业务条线编写并验证 SQL 查询。

    此工具帮助格式化和验证 SQL 查询。你必须先加载相应技能，
    以了解数据库 schema。

    Args:
        query: 要编写的 SQL 查询
        vertical: 业务条线（sales_analytics 或 inventory_management）
    """
    # 检查是否已经加载所需技能
    skills_loaded = runtime.state.get("skills_loaded", [])

    if vertical not in skills_loaded:
        return (
            f"错误：你必须先加载 '{vertical}' 技能，"
            f"才能在编写查询前了解数据库 schema。"
            f"使用 load_skill('{vertical}') 加载 schema。"
        )

    # 验证并格式化查询
    return (
        f"{vertical} 的 SQL 查询：\n\n"
        f"```sql\n{query}\n```\n\n"
        f"✓ 已根据 {vertical} schema 验证查询\n"
        f"可以对数据库执行。"
    )

class SkillMiddleware(AgentMiddleware[CustomState]):
    """将技能描述注入系统提示词的中间件。"""

    state_schema = CustomState  
    tools = [load_skill, write_sql_query]

    # ... 中间件实现的其余部分保持不变

agent = create_agent(
    model,
    system_prompt=(
        "你是一个 SQL 查询助手，帮助用户"
        "针对业务数据库编写查询。"
    ),
    middleware=[SkillMiddleware()],
    checkpointer=InMemorySaver(),
)

from langchain_core.utils.uuid import uuid7
from typing import TypedDict, NotRequired
from langchain.tools import tool
from langchain.agents import create_agent
from langchain.agents.middleware import ModelRequest, ModelResponse, AgentMiddleware
from langchain.messages import SystemMessage
from langgraph.checkpoint.memory import InMemorySaver
from typing import Callable

# 定义技能结构
class Skill(TypedDict):
    """可以渐进式披露给智能体的技能。"""
    name: str
    description: str
    content: str

# 定义包含 schema 和业务逻辑的技能
SKILLS: list[Skill] = [
    {
        "name": "sales_analytics",
        "description": "用于销售数据分析的数据库 schema 和业务逻辑，包括 customers、orders 和 revenue。",
        "content": """# 销售分析 Schema

## 表

### customers
- customer_id (PRIMARY KEY)
- name
- email
- signup_date
- status (active/inactive)
- customer_tier (bronze/silver/gold/platinum)

### orders
- order_id (PRIMARY KEY)
- customer_id (FOREIGN KEY -> customers)
- order_date
- status (pending/completed/cancelled/refunded)
- total_amount
- sales_region (north/south/east/west)

### order_items
- item_id (PRIMARY KEY)
- order_id (FOREIGN KEY -> orders)
- product_id
- quantity
- unit_price
- discount_percent

## 业务逻辑

**活跃客户**：status = 'active' AND signup_date <= CURRENT_DATE - INTERVAL '90 days'

**收入计算**：仅统计 status = 'completed' 的订单。使用 orders 表中的 total_amount，它已经包含折扣计算。

**客户生命周期价值 (CLV)**：某个客户所有已完成订单金额的总和。

**高价值订单**：total_amount > 1000 的订单

## 示例查询

-- 获取上一季度收入最高的 10 个客户
SELECT
    c.customer_id,
    c.name,
    c.customer_tier,
    SUM(o.total_amount) as total_revenue
FROM customers c
JOIN orders o ON c.customer_id = o.customer_id
WHERE o.status = 'completed'
  AND o.order_date >= CURRENT_DATE - INTERVAL '3 months'
GROUP BY c.customer_id, c.name, c.customer_tier
ORDER BY total_revenue DESC
LIMIT 10;
""",
    },
    {
        "name": "inventory_management",
        "description": "用于库存跟踪的数据库 schema 和业务逻辑，包括 products、warehouses 和 stock levels。",
        "content": """# 库存管理 Schema

## 表

### products
- product_id (PRIMARY KEY)
- product_name
- sku
- category
- unit_cost
- reorder_point (重新订购前的最低库存水平)
- discontinued (boolean)

### warehouses
- warehouse_id (PRIMARY KEY)
- warehouse_name
- location
- capacity

### inventory
- inventory_id (PRIMARY KEY)
- product_id (FOREIGN KEY -> products)
- warehouse_id (FOREIGN KEY -> warehouses)
- quantity_on_hand
- last_updated

### stock_movements
- movement_id (PRIMARY KEY)
- product_id (FOREIGN KEY -> products)
- warehouse_id (FOREIGN KEY -> warehouses)
- movement_type (inbound/outbound/transfer/adjustment)
- quantity (inbound 为正数，outbound 为负数)
- movement_date
- reference_number

## 业务逻辑

**可用库存**：inventory 表中 quantity_on_hand > 0 的 quantity_on_hand

**需要重新订购的产品**：所有仓库中的 total quantity_on_hand 小于或等于产品 reorder_point 的产品

**仅活跃产品**：除非专门分析已停用商品，否则排除 discontinued = true 的产品

**库存估值**：每个产品的 quantity_on_hand * unit_cost

## 示例查询

-- 查找所有仓库中低于重新订购点的产品
SELECT
    p.product_id,
    p.product_name,
    p.reorder_point,
    SUM(i.quantity_on_hand) as total_stock,
    p.unit_cost,
    (p.reorder_point - SUM(i.quantity_on_hand)) as units_to_reorder
FROM products p
JOIN inventory i ON p.product_id = i.product_id
WHERE p.discontinued = false
GROUP BY p.product_id, p.product_name, p.reorder_point, p.unit_cost
HAVING SUM(i.quantity_on_hand) <= p.reorder_point
ORDER BY units_to_reorder DESC;
""",
    },
]

# 创建技能加载工具
@tool
def load_skill(skill_name: str) -> str:
    """将技能的完整内容加载到智能体上下文中。

    当你需要有关如何处理特定类型请求的详细信息时使用此工具。
    它会提供该技能领域的完整指令、策略和指南。

    Args:
        skill_name: 要加载的技能名称（例如 "sales_analytics"、"inventory_management"）
    """
    # 查找并返回请求的技能
    for skill in SKILLS:
        if skill["name"] == skill_name:
            return f"已加载技能：{skill_name}\n\n{skill['content']}"

    # 未找到技能
    available = ", ".join(s["name"] for s in SKILLS)
    return f"未找到技能 '{skill_name}'。可用技能：{available}"

# 创建技能中间件
class SkillMiddleware(AgentMiddleware):
    """将技能描述注入系统提示词的中间件。"""

    # 将 load_skill 工具注册为类变量
    tools = [load_skill]

    def __init__(self):
        """初始化并根据 SKILLS 生成技能提示词。"""
        # 根据 SKILLS 列表构建技能提示词
        skills_list = []
        for skill in SKILLS:
            skills_list.append(
                f"- **{skill['name']}**: {skill['description']}"
            )
        self.skills_prompt = "\n".join(skills_list)

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        """同步：将技能描述注入系统提示词。"""
        # 构建技能补充内容
        skills_addendum = (
            f"\n\n## 可用技能\n\n{self.skills_prompt}\n\n"
            "当你需要有关如何处理特定类型请求的详细信息时，"
            "使用 load_skill 工具。"
        )

        # 追加到系统消息内容块
        new_content = list(request.system_message.content_blocks) + [
            {"type": "text", "text": skills_addendum}
        ]
        new_system_message = SystemMessage(content=new_content)
        modified_request = request.override(system_message=new_system_message)
        return handler(modified_request)

# 初始化你的聊天模型（替换为你的模型）
# 示例：from langchain_anthropic import ChatAnthropic
# model = ChatAnthropic(model="claude-3-5-sonnet-20241022")
from langchain_openai import ChatOpenAI
model = ChatOpenAI(model="gpt-4")

# 创建支持技能的智能体
agent = create_agent(
    model,
    system_prompt=(
        "你是一个 SQL 查询助手，帮助用户"
        "针对业务数据库编写查询。"
    ),
    middleware=[SkillMiddleware()],
    checkpointer=InMemorySaver(),
)

# 示例用法
if __name__ == "__main__":
    # 此对话线程的配置
    thread_id = str(uuid7())
    config = {"configurable": {"thread_id": thread_id}}

    # 请求一个 SQL 查询
    result = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": (
                        "编写一个 SQL 查询，找出上个月"
                        "下过超过 $1000 订单的所有客户"
                    ),
                }
            ]
        },
        config
    )

    # 打印对话
    for message in result["messages"]:
        if hasattr(message, 'pretty_print'):
            message.pretty_print()
        else:
            print(f"{message.type}: {message.content}")