构建具有按需技能的 SQL 助手

import { context } from "langchain";

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

    ## 表

    ### 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: context`
    # 库存管理模式

    ## 表

    ### 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;`,
  },
];

import { StateSchema } from "@langchain/langgraph";
import { z } from "zod";

const CustomState = new StateSchema({
  skillsLoaded: z.array(z.string()).optional(),  // 跟踪已加载的技能  
});

import { tool, ToolMessage, type ToolRuntime } from "langchain";
import { Command } from "@langchain/langgraph";
import { z } from "zod";

const loadSkill = tool(  
  async ({ skillName }, runtime: ToolRuntime<typeof CustomState.State>) => {
    // 查找并返回请求的技能
    const skill = SKILLS.find((s) => s.name === skillName);

    if (skill) {
      const skillContent = `已加载技能：${skillName}\n\n${skill.content}`;

      // 更新状态以跟踪已加载的技能
      return new Command({
        update: {
          messages: [  
            new ToolMessage({
              content: skillContent,
              tool_call_id: runtime.toolCallId,
            }),
          ],
          skillsLoaded: [skillName],
        },
      });
    }

    // 技能未找到
    const available = SKILLS.map((s) => s.name).join(", ");
    return new Command({
      update: {
        messages: [
          new ToolMessage({
            content: `技能 '${skillName}' 未找到。可用技能：${available}`,
            tool_call_id: runtime.toolCallId,
          }),
        ],
      },
    });
  },
  {
    name: "load_skill",
    description: `将技能的完整内容加载到代理的上下文中。`,
    schema: z.object({
      skillName: z.string().describe("要加载的技能的名称"),
    }),
  }
);

const writeSqlQuery = tool(  
  async ({ query, vertical }, runtime: ToolRuntime<typeof CustomState.State>) => {
    // 检查是否已加载所需的技能
    const skillsLoaded = runtime.state.skillsLoaded ?? [];

    if (!skillsLoaded.includes(vertical)) {
      return (  
        `错误：您必须首先加载 '${vertical}' 技能` +
        `以理解数据库模式，然后才能编写查询。` +
        `请使用 load_skill('${vertical}') 加载模式。`
      );
    }

    // 验证并格式化查询
    return (
      `${vertical} 的 SQL 查询：\n\n` +
      `\`\`\`sql\n${query}\n\`\`\`\n\n` +
      `✓ 查询已针对 ${vertical} 模式验证\n` +
      `准备好对数据库执行。`
    );
  },
  {
    name: "write_sql_query",
    description: `为特定业务部门编写并验证 SQL 查询。

此工具有助于格式化和验证 SQL 查询。您必须首先加载相应的技能以理解数据库模式。`,
    schema: z.object({
      query: z.string().describe("要编写的 SQL 查询"),
      vertical: z.string().describe("业务部门 (sales_analytics 或 inventory_management)"),
    }),
  }
);

const skillMiddleware = createMiddleware({
  name: "skillMiddleware",
  stateSchema: CustomState,
  tools: [loadSkill, writeSqlQuery],
  // ... 中间件实现的其余部分保持不变
});

const agent = createAgent({
  model,
  systemPrompt:
    "您是一名 SQL 查询助手，帮助用户" +
    "针对业务数据库编写查询。",
  middleware: [skillMiddleware],
  checkpointer: new MemorySaver(),
});

import {
  tool,
  createAgent,
  createMiddleware,
  ToolMessage,
  context,
  type ToolRuntime,
} from "langchain";
import { MemorySaver, Command } from "@langchain/langgraph";
import { ChatOpenAI } from "@langchain/openai";
import { z } from "zod";

// 可以向代理渐进式披露的技能
const SkillSchema = z.object({
  name: z.string(), // 技能的唯一标识符
  description: z.string(), // 1-2 句话的描述，显示在系统提示中
  content: z.string(), // 完整的技能内容，包含详细说明
});

type Skill = z.infer<typeof SkillSchema>;

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

    ## 表

    ### 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: context`
    # 库存管理模式

    ## 表

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

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


    ### 库存表
    - 库存标识（主键）
    - 产品标识（外键 -> 产品表）
    - 仓库标识（外键 -> 仓库表）
    - 当前库存数量
    - 最后更新时间

    ### 库存变动记录
    - 变动标识（主键）
    - 产品标识（外键 -> 产品表）
    - 仓库标识（外键 -> 仓库表）
    - 变动类型（入库/出库/调拨/调整）
    - 数量（入库为正数，出库为负数）
    - 变动日期
    - 参考编号

    ## 业务逻辑

    **可用库存**：
    来自库存表中库存数量大于零的 `quantity_on_hand` 记录

    **需要补货的产品**：
    所有仓库的总库存数量（`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;`,
    },
    ];

    // const loadSkill = tool(
    //   async ({ skillName }) => {
    //     // 查找并返回请求的技能
    //     const skill = SKILLS.find((s) => s.name === skillName);
    //     if (skill) {
    //       return `已加载技能: ${skillName}\n\n${skill.content}`;
    //     }

    //     // 技能未找到
    //     const available = SKILLS.map((s) => s.name).join(", ");
    //     return `技能 '${skillName}' 未找到。可用技能: ${available}`;
    //   },
    //   {
    //     name: "load_skill",
    //     description: `将一个技能的完整内容加载到代理的上下文中。
    //
    // 当您需要关于如何处理特定类型请求的详细信息时使用此工具。它将为您提供该技能领域全面的说明、策略和指南。`,
    //     schema: z.object({
    //       skillName: z.string().describe("要加载的技能名称"),
    //     }),
    //   }
    // );

    // 根据 SKILLS 列表构建技能提示
    const skillsPrompt = SKILLS.map(
    (skill) => `- **${skill.name}**: ${skill.description}`
    ).join("\n");

    const skillMiddleware = createMiddleware({
    name: "skillMiddleware",
    tools: [loadSkill],
    wrapModelCall: async (request, handler) => {
        // 构建技能附录
        const skillsAddendum =
        `\n\n## 可用技能\n\n${skillsPrompt}\n\n` +
        "当您需要处理特定类型请求的详细信息时，请使用 load_skill 工具。";

        // 追加到系统提示
        const newSystemPrompt = request.systemPrompt + skillsAddendum;

        return handler({
        ...request,
        systemPrompt: newSystemPrompt,
        });
    },
    });

    const model = new ChatOpenAI({
    model: "gpt-5.4-mini",
    temperature: 0,
    });

    // 创建具有技能支持的代理
    const agent = createAgent({
    model,
    systemPrompt:
        "您是一个 SQL 查询助手，帮助用户针对业务数据库编写查询。",
    middleware: [skillMiddleware],
    checkpointer: new MemorySaver(),
    });

    // 此对话线程的配置
    const threadId = crypto.randomUUID();
    const config = { configurable: { thread_id: threadId } };

    // 请求一个 SQL 查询
    const result = await agent.invoke(
    {
        messages: [
        {
            role: "user",
            content:
            "编写一个 SQL 查询，查找上个月消费超过 1000 美元的所有客户",
        },
        ],
    },
    config
    );

    // 打印对话
    for (const message of result.messages) {
    console.log(`${message.type}: ${message.content}`);
    }