使用按需技能构建 SQL 助手

SKILLS: list[Skill] = [
    {
        "name": "sales_analytics",
        "description": "销售数据分析的数据库模式和业务逻辑，包括客户、订单和收入。",
        "content": """# 销售分析模式

## 表

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

### orders
- order_id (主键)
- customer_id (外键 -> customers)
- order_date
- status (pending/completed/cancelled/refunded)
- total_amount
- sales_region (north/south/east/west)

### order_items
- item_id (主键)
- order_id (外键 -> 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": "库存跟踪的数据库模式和业务逻辑，包括产品、仓库和库存水平。",
        "content": """# 库存管理模式

## 表

### products
- product_id (主键)
- product_name
- sku
- category
- unit_cost
- reorder_point (重新订购前的最低库存水平)
- discontinued (布尔值)

### warehouses
- warehouse_id (主键)
- warehouse_name
- location
- capacity

### inventory
- inventory_id (主键)
- product_id (外键 -> products)
- warehouse_id (外键 -> warehouses)
- quantity_on_hand
- last_updated

### stock_movements
- movement_id (主键)
- product_id (外键 -> products)
- warehouse_id (外键 -> warehouses)
- movement_type (inbound/outbound/transfer/adjustment)
- quantity (正数表示入库，负数表示出库)
- movement_date
- reference_number

## 业务逻辑

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

**需要重新订购的产品**：所有仓库中总 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 查询。您必须先加载相应的技能以了解数据库模式。

    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"以了解数据库模式，然后才能编写查询。"
            f"使用 load_skill('{vertical}') 加载模式。"
        )

    # 验证并格式化查询
    return (
        f"SQL 查询，适用于 {vertical}：\n\n"
        f"```sql\n{query}\n```\n\n"
        f"✓ 查询已根据 {vertical} 模式验证\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

# 定义包含模式和业务逻辑的技能
SKILLS: list[Skill] = [
    {
        "name": "sales_analytics",
        "description": "销售数据分析的数据库模式和业务逻辑，包括客户、订单和收入。",
        "content": """# 销售分析模式

## 表

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

### orders
- order_id (主键)
- customer_id (外键 -> customers)
- order_date
- status (pending/completed/cancelled/refunded)
- total_amount
- sales_region (north/south/east/west)

### order_items
- item_id (主键)
- order_id (外键 -> 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": "库存跟踪的数据库模式和业务逻辑，包括产品、仓库和库存水平。",
        "content": """# 库存管理模式

## 表

### products
- product_id (主键)
- product_name
- sku
- category
- unit_cost
- reorder_point (重新订购前的最低库存水平)
- discontinued (布尔值)

### warehouses
- warehouse_id (主键)
- warehouse_name
- location
- capacity

### inventory
- inventory_id (主键)
- product_id (外键 -> products)
- warehouse_id (外键 -> warehouses)
- quantity_on_hand
- last_updated

### stock_movements
- movement_id (主键)
- product_id (外键 -> products)
- warehouse_id (外键 -> warehouses)
- movement_type (inbound/outbound/transfer/adjustment)
- quantity (正数表示入库，负数表示出库)
- movement_date
- reference_number

## 业务逻辑

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

**需要重新订购的产品**：所有仓库中总 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}")