数据库Agent架构

// src/agents/database/prompt.ts
 
// 只暴露允许 Agent 访问的表和字段
const DATABASE_SCHEMA = `
可用的表结构：
- users: id(integer), email(text), name(text), role(text), created_at(timestamp)
- orders: id(integer), user_id(integer), status(text), total(decimal), created_at(timestamp)
- products: id(integer), name(text), price(decimal), stock(integer), category(text)
 
规则：
1. 只能查询上述表，不能访问其他表
2. 写操作（INSERT/UPDATE/DELETE）需要用户明确授权
3. 查询结果默认限制 100 行
4. 不要生成 DROP、TRUNCATE、ALTER 等 DDL 语句
`

// src/agents/database/prompt.ts
const SQL_EXAMPLES = `
示例 1：
用户：上个月注册了多少新用户
SQL：SELECT COUNT(*) AS new_users FROM users WHERE created_at >= date('now', '-1 month')
 
示例 2：
用户：销售额最高的 5 个商品
SQL：SELECT p.name, SUM(o.total) AS total_sales FROM orders o JOIN products p ON o.product_id = p.id GROUP BY p.name ORDER BY total_sales DESC LIMIT 5
 
示例 3：
用户：把订单 12345 的状态改成已发货
SQL：UPDATE orders SET status = 'shipped' WHERE id = 12345
`

// src/agents/database/formatter.ts
 
// 把数据库原始结果转成 LLM 能理解的文本
export function formatQueryResult(
  columns: string[],
  rows: unknown[][],
): string {
  if (rows.length === 0) return '查询结果为空，没有匹配的数据。'
 
  // 转成 CSV 格式，LLM 对这种格式的理解比较好
  const header = columns.join(' | ')
  const separator = columns.map(() => '---').join(' | ')
  const body = rows
    .map((row) => row.map((v) => String(v ?? 'NULL')).join(' | '))
    .join('\n')
 
  return `查询结果（共 ${rows.length} 行）：\n${header}\n${separator}\n${body}`
}

// src/agents/database/tools.ts
import type { ToolDefinition } from '../../tools/types'
 
export const databaseTools: ToolDefinition[] = [
  {
    name: 'query',
    description: '执行只读 SQL 查询，返回结果。用于查询数据、统计分析。只能执行 SELECT 语句。',
    parameters: {
      type: 'object',
      properties: {
        sql: {
          type: 'string',
          description: 'SELECT 语句，例如：SELECT * FROM users WHERE role = \'admin\' LIMIT 10'
        },
        reason: {
          type: 'string',
          description: '执行这条查询的原因，用于审计日志'
        }
      },
      required: ['sql', 'reason']
    }
  },
  {
    name: 'insert',
    description: '向数据库插入一条记录。用于新增数据。',
    parameters: {
      type: 'object',
      properties: {
        table: { type: 'string', description: '目标表名' },
        data: { type: 'object', description: '要插入的字段和值' },
        reason: { type: 'string', description: '插入原因' }
      },
      required: ['table', 'data', 'reason']
    }
  },
  {
    name: 'update',
    description: '更新数据库中的记录。必须提供 WHERE 条件，不能无条件更新全表。',
    parameters: {
      type: 'object',
      properties: {
        table: { type: 'string', description: '目标表名' },
        data: { type: 'object', description: '要更新的字段和新值' },
        where: { type: 'string', description: 'WHERE 条件，例如：id = 123' },
        reason: { type: 'string', description: '更新原因' }
      },
      required: ['table', 'data', 'where', 'reason']
    }
  },
  {
    name: 'delete',
    description: '删除数据库中的记录。必须提供 WHERE 条件，不能无条件删除全表。',
    parameters: {
      type: 'object',
      properties: {
        table: { type: 'string', description: '目标表名' },
        where: { type: 'string', description: 'WHERE 条件，例如：id = 123' },
        reason: { type: 'string', description: '删除原因' }
      },
      required: ['table', 'where', 'reason']
    }
  }
]

// src/agents/database/executor.ts
import type { ToolContext } from '../../tools/types'
import { validateSQL, isReadOnly } from './security'
 
type DatabaseAgentContext = ToolContext & {
  db: D1Database
  allowedTables: string[]
}
 
export async function executeDatabaseTool(
  toolName: string,
  args: Record<string, unknown>,
  ctx: DatabaseAgentContext
): Promise<unknown> {
  switch (toolName) {
    case 'query':
      return executeQuery(args.sql as string, ctx)
    case 'insert':
      return executeInsert(args.table as string, args.data as Record<string, unknown>, ctx)
    case 'update':
      return executeUpdate(
        args.table as string,
        args.data as Record<string, unknown>,
        args.where as string,
        ctx
      )
    case 'delete':
      return executeDelete(args.table as string, args.where as string, ctx)
    default:
      return { error: `未知工具：${toolName}` }
  }
}
 
async function executeQuery(sql: string, ctx: DatabaseAgentContext) {
  // 第一步：检查是否为只读语句
  if (!isReadOnly(sql)) {
    return { error: 'query 工具只允许执行 SELECT 语句' }
  }
 
  // 第二步：安全校验
  const validationError = validateSQL(sql, ctx.allowedTables)
  if (validationError) {
    return { error: validationError }
  }
 
  // 第三步：强制加 LIMIT 防止返回过多数据
  const limitedSQL = enforceLimit(sql, 100)
 
  // 第四步：执行查询
  try {
    const result = await ctx.db.prepare(limitedSQL).all()
    return {
      columns: result.columns ?? Object.keys(result.results?.[0] ?? {}),
      rows: result.results,
      rowCount: result.results?.length ?? 0
    }
  } catch (err) {
    return { error: `查询执行失败：${err instanceof Error ? err.message : '未知错误'}` }
  }
}
 
async function executeInsert(
  table: string,
  data: Record<string, unknown>,
  ctx: DatabaseAgentContext
) {
  if (!ctx.allowedTables.includes(table)) {
    return { error: `不允许操作表：${table}` }
  }
 
  const columns = Object.keys(data)
  const placeholders = columns.map(() => '?').join(', ')
  const sql = `INSERT INTO ${table} (${columns.join(', ')}) VALUES (${placeholders})`
 
  try {
    await ctx.db.prepare(sql).bind(...Object.values(data)).run()
    return { success: true, message: `已向 ${table} 插入 1 条记录` }
  } catch (err) {
    return { error: `插入失败：${err instanceof Error ? err.message : '未知错误'}` }
  }
}
 
// enforceLimit：如果 SQL 没有 LIMIT 子句，自动追加
function enforceLimit(sql: string, maxRows: number): string {
  if (/LIMIT\s+\d+/i.test(sql)) return sql
  return `${sql.replace(/;?\s*$/, '')} LIMIT ${maxRows}`
}

// src/agents/database/agent.ts
import { Hono } from 'hono'
import type { AppEnv } from '../../types'
import { databaseTools } from './tools'
import { executeDatabaseTool } from './executor'
import { DATABASE_SCHEMA, SQL_EXAMPLES } from './prompt'
 
const dbAgentApp = new Hono<AppEnv>()
 
dbAgentApp.post('/query', async (c) => {
  const { message, userId } = await c.req.json<{
    message: string
    userId: string
  }>()
 
  const systemPrompt = `你是一个数据库查询助手。根据用户的需求生成并执行 SQL，然后用自然语言回复结果。
 
${DATABASE_SCHEMA}
 
${SQL_EXAMPLES}
 
注意事项：
1. 查询操作使用 query 工具
2. 插入操作使用 insert 工具
3. 更新操作使用 update 工具，必须提供 WHERE 条件
4. 删除操作使用 delete 工具，必须提供 WHERE 条件
5. 不要编造数据，所有事实必须来自工具返回的结果
6. 如果查询结果为空，如实告知用户`
 
  const messages = [
    { role: 'system', content: systemPrompt },
    { role: 'user', content: message }
  ]
 
  // Agent 循环
  const MAX_ROUNDS = 5
  for (let i = 0; i < MAX_ROUNDS; i++) {
    const response = await callLLM(messages, databaseTools)
 
    // 没有工具调用，LLM 已生成最终回复
    if (!response.toolCalls?.length) {
      return c.json({ reply: response.text })
    }
 
    // 执行工具调用
    messages.push({
      role: 'assistant',
      content: response.text,
      tool_calls: response.toolCalls
    })
 
    for (const toolCall of response.toolCalls) {
      const result = await executeDatabaseTool(
        toolCall.name,
        toolCall.arguments,
        {
          userId,
          db: c.env.DB,
          allowedTables: ['users', 'orders', 'products'],
          signal: new AbortController().signal,
          env: c.env as unknown as Record<string, string>
        }
      )
 
      messages.push({
        role: 'tool',
        tool_call_id: toolCall.id,
        content: JSON.stringify(result)
      })
    }
  }
 
  return c.json({ error: '达到最大执行轮次' }, 500)
})
 
export default dbAgentApp

// src/agents/database/security.ts
 
// 禁止出现在 SQL 中的关键词
const FORBIDDEN_KEYWORDS = [
  'DROP', 'TRUNCATE', 'ALTER', 'CREATE',
  'GRANT', 'REVOKE', 'EXEC', 'EXECUTE',
  'xp_', 'sp_'
]
 
export function validateSQL(
  sql: string,
  allowedTables: string[]
): string | null {
  const upperSQL = sql.toUpperCase()
 
  // 检查禁止关键词
  for (const keyword of FORBIDDEN_KEYWORDS) {
    if (upperSQL.includes(keyword)) {
      return `SQL 包含禁止操作：${keyword}`
    }
  }
 
  // 检查是否引用了未授权的表
  const tablePattern = /(?:FROM|JOIN|INTO|UPDATE|TABLE)\s+(\w+)/gi
  let match
  while ((match = tablePattern.exec(sql)) !== null) {
    const tableName = match[1].toLowerCase()
    if (!allowedTables.includes(tableName)) {
      return `不允许访问表：${match[1]}`
    }
  }
 
  // 检查是否包含子查询中的写操作
  if (/\b(INSERT|UPDATE|DELETE)\b/i.test(sql)) {
    return 'SELECT 语句中不允许包含写操作'
  }
 
  return null
}
 
export function isReadOnly(sql: string): boolean {
  const trimmed = sql.trim().toUpperCase()
  return trimmed.startsWith('SELECT') || trimmed.startsWith('WITH')
}

// src/agents/database/permissions.ts
import { databaseTools } from './tools'
import type { ToolDefinition } from '../../tools/types'
 
const READ_ONLY_TOOLS = ['query']
const WRITE_TOOLS = ['insert', 'update', 'delete']
 
export function getToolsForRole(userRole: string): ToolDefinition[] {
  if (userRole === 'admin') {
    return databaseTools
  }
  // 普通用户只能查询
  return databaseTools.filter((t) => READ_ONLY_TOOLS.includes(t.name))
}

// src/agents/database/executor.ts
// 在 executeQuery 的 catch 块中
catch (err) {
  const errorMessage = err instanceof Error ? err.message : '未知错误'
  return {
    error: `SQL 执行失败：${errorMessage}`,
    hint: '请检查 SQL 语法、表名和字段名是否正确'
  }
}

// 返回给 LLM 的错误信息
{ error: '该操作不被允许：引用了未授权的表。请只使用允许范围内的表进行查询。' }
 
// 而不是
{ error: '安全检查失败：表 user_secrets 不在 allowedTables 白名单 ["users", "orders", "products"] 中' }

// src/agents/database/executor.ts
async function executeWithTimeout<T>(
  fn: () => Promise<T>,
  timeoutMs = 5000
): Promise<T> {
  return new Promise((resolve, reject) => {
    const timer = setTimeout(
      () => reject(new Error(`查询超时（${timeoutMs}ms），请缩小查询范围或添加筛选条件`)),
      timeoutMs
    )
    fn()
      .then((result) => { clearTimeout(timer); resolve(result) })
      .catch((err) => { clearTimeout(timer); reject(err) })
  })
}

// 在 executeQuery 中
if (result.results?.length === 0) {
  return {
    message: '查询成功，但没有匹配的数据。请确认查询条件是否正确，或者告知用户没有找到相关数据。',
    rowCount: 0
  }
}