构建具备按需技能的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'

    **收入计算**：
    仅计算状态为 '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

    **需要重新订购的产品**：
    所有仓库中 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 { v4 as uuidv4 } from "uuid";
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'

    **收入计算**：
    仅计算状态为 '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

    **需要重新订购的产品**：
    所有仓库中 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 = uuidv4();
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}`);
}