AI Agent 工具系统实现

// lib/ai/tools.ts
import { tool } from 'ai'
import { z } from 'zod'
 
// 天气查询工具（示例）
export const weatherTool = tool({
  description: 'Get current weather for a location',
  parameters: z.object({
    city: z.string().describe('City name'),
    unit: z.enum(['celsius', 'fahrenheit']).default('celsius'),
  }),
  execute: async ({ city, unit }) => {
    const res = await fetch(`https://api.weather.com/v1?city=${city}&unit=${unit}`)
    return res.json()
  },
})

// lib/ai/tools/project-tools.ts
 
export const listProjects = tool({
  description: 'List all projects in the current organization, optionally filtered by status',
  parameters: z.object({
    status: z.enum(['active', 'completed', 'archived']).optional()
      .describe('Filter by project status'),
    limit: z.number().max(20).default(10)
      .describe('Maximum number of projects to return'),
  }),
  execute: async ({ status, limit }, { tenantId }) => {
    const query = db.select({
      id: projects.id,
      name: projects.name,
      status: projects.status,
      createdAt: projects.createdAt,
    })
      .from(projects)
      .where(and(
        eq(projects.tenantId, tenantId),
        status ? eq(projects.status, status) : undefined,
      ))
      .orderBy(desc(projects.createdAt))
      .limit(limit)
 
    return query
  },
})
 
export const createTask = tool({
  description: 'Create a new task in a project',
  parameters: z.object({
    projectId: z.string().uuid().describe('The project ID to add the task to'),
    title: z.string().describe('Task title'),
    description: z.string().optional().describe('Task description'),
    priority: z.enum(['low', 'medium', 'high', 'urgent']).default('medium'),
    assigneeEmail: z.string().email().optional()
      .describe('Email of the person to assign the task to'),
  }),
  execute: async ({ projectId, title, description, priority, assigneeEmail }, { tenantId, userId }) => {
    // 验证项目归属
    const [project] = await db.select().from(projects)
      .where(and(eq(projects.id, projectId), eq(projects.tenantId, tenantId)))
    if (!project) return { error: 'Project not found' }
 
    let assigneeId: string | undefined
    if (assigneeEmail) {
      const [user] = await db.select().from(users).where(eq(users.email, assigneeEmail))
      assigneeId = user?.id
    }
 
    const [task] = await db.insert(tasks).values({
      projectId, title, description, priority, assigneeId,
      createdBy: userId,
    }).returning()
 
    return { success: true, taskId: task.id, title: task.title }
  },
})
 
export const queryMetrics = tool({
  description: 'Query business metrics like revenue, user count, or project stats for a given time range',
  parameters: z.object({
    metric: z.enum(['revenue', 'activeUsers', 'newSignups', 'projectCount'])
      .describe('The metric to query'),
    period: z.enum(['today', 'thisWeek', 'thisMonth', 'lastMonth', 'last30days'])
      .describe('Time period for the metric'),
  }),
  execute: async ({ metric, period }, { tenantId }) => {
    const dateRange = getDateRange(period)
 
    switch (metric) {
      case 'revenue':
        return queryRevenue(tenantId, dateRange)
      case 'activeUsers':
        return queryActiveUsers(tenantId, dateRange)
      case 'newSignups':
        return queryNewSignups(tenantId, dateRange)
      case 'projectCount':
        return queryProjectCount(tenantId, dateRange)
    }
  },
})

// lib/ai/tools/registry.ts
 
interface ToolCategory {
  name: string
  description: string
  tools: Record<string, any>
  requiredPermission?: string
}
 
const TOOL_REGISTRY: ToolCategory[] = [
  {
    name: 'projects',
    description: 'Project and task management',
    tools: { listProjects, createTask, getProjectDetails },
    requiredPermission: 'ai:agent:use',
  },
  {
    name: 'analytics',
    description: 'Business metrics and reporting',
    tools: { queryMetrics, generateChart },
    requiredPermission: 'ai:agent:use',
  },
  {
    name: 'communication',
    description: 'Send emails and notifications',
    tools: { sendEmail, sendNotification },
    requiredPermission: 'ai:agent:manage',  // 更高权限
  },
  {
    name: 'search',
    description: 'Web search and knowledge base query',
    tools: { webSearch, queryKnowledgeBase },
    requiredPermission: 'ai:agent:use',
  },
]
 
export function getAvailableTools(userPermissions: string[]): Record<string, any> {
  const tools: Record<string, any> = {}
 
  for (const category of TOOL_REGISTRY) {
    if (!category.requiredPermission || userPermissions.includes(category.requiredPermission)) {
      Object.assign(tools, category.tools)
    }
  }
 
  return tools
}

// app/api/agent/route.ts
import { streamText } from 'ai'
import { openai } from '@ai-sdk/openai'
 
export async function POST(req: Request) {
  const session = await getSession()
  if (!session) return new Response('Unauthorized', { status: 401 })
 
  const tenantId = await getCurrentTenantId()
  const { messages } = await req.json()
 
  const tools = getAvailableTools(session.user.permissions)
 
  const result = streamText({
    model: openai('gpt-4o'),  // Agent 推荐用更强的模型
    system: buildAgentSystemPrompt(),
    messages,
    tools,
    maxSteps: 5,  // 最多 5 轮工具调用
    // 每个工具调用都会注入上下文
    toolChoice: 'auto',  // 让模型自己决定是否调用工具
 
    onStepFinish: async ({ stepType, toolCalls, toolResults, usage }) => {
      // 记录每一步的执行情况
      if (stepType === 'tool-result') {
        await logAgentStep(tenantId!, session.user.id, {
          toolCalls,
          toolResults,
          usage,
        })
      }
    },
 
    onFinish: async ({ text, usage, steps }) => {
      await recordUsage(tenantId!, session.user.id, 'gpt-4o', {
        promptTokens: steps.reduce((sum, s) => sum + (s.usage?.promptTokens || 0), 0),
        completionTokens: steps.reduce((sum, s) => sum + (s.usage?.completionTokens || 0), 0),
      })
    },
  })
 
  return result.toDataStreamResponse()
}

function buildAgentSystemPrompt(): string {
  return `You are an AI assistant with access to tools that can perform actions in the user's workspace.
 
BEHAVIOR:
- Think step by step. Before taking action, briefly explain what you're going to do and why.
- If a task requires multiple steps, plan them out first, then execute one by one.
- After executing tools, summarize the results clearly for the user.
- If a tool returns an error, explain what went wrong and suggest alternatives.
 
SAFETY:
- Never perform destructive actions (delete, modify) without explicit user confirmation.
- If the user's request is ambiguous, ask for clarification before acting.
- Do not call the same tool with the same parameters more than twice (avoid infinite loops).
- Respect rate limits: do not call more than 3 tools in a single response.
 
LIMITATIONS:
- You can only access data within the user's current organization.
- Email sending requires explicit user consent.
- You cannot access external systems not exposed as tools.`
}

// components/agent-message.tsx
import type { Message } from 'ai'
 
export function AgentMessage({ message }: { message: Message }) {
  return (
    <div className="space-y-3">
      {/* 工具调用步骤 */}
      {message.toolInvocations?.map((invocation, i) => (
        <ToolInvocationCard key={i} invocation={invocation} />
      ))}
 
      {/* 最终文本回答 */}
      {message.content && (
        <div className="rounded-lg bg-muted px-4 py-3">
          <MarkdownRenderer content={message.content} />
        </div>
      )}
    </div>
  )
}
 
function ToolInvocationCard({ invocation }: { invocation: any }) {
  const isComplete = invocation.state === 'result'
 
  return (
    <div className="rounded-lg border bg-card p-3 text-sm">
      <div className="flex items-center gap-2">
        {isComplete ? (
          <span className="text-green-500">✓</span>
        ) : (
          <span className="animate-spin">⟳</span>
        )}
        <span className="font-medium">{formatToolName(invocation.toolName)}</span>
      </div>
 
      {/* 工具参数 */}
      <details className="mt-2">
        <summary className="cursor-pointer text-muted-foreground text-xs">
          Parameters
        </summary>
        <pre className="mt-1 rounded bg-muted p-2 text-xs overflow-x-auto">
          {JSON.stringify(invocation.args, null, 2)}
        </pre>
      </details>
 
      {/* 工具结果 */}
      {isComplete && invocation.result && (
        <details className="mt-2">
          <summary className="cursor-pointer text-muted-foreground text-xs">
            Result
          </summary>
          <pre className="mt-1 rounded bg-muted p-2 text-xs overflow-x-auto">
            {JSON.stringify(invocation.result, null, 2)}
          </pre>
        </details>
      )}
    </div>
  )
}
 
function formatToolName(name: string): string {
  return name.replace(/([A-Z])/g, ' $1').replace(/^./, s => s.toUpperCase()).trim()
}

// lib/ai/tools/web-search.ts
import { tool } from 'ai'
import { z } from 'zod'
 
export const webSearch = tool({
  description: 'Search the web for current information. Use this when the user asks about recent events, real-time data, or topics that may have changed since your training cutoff.',
  parameters: z.object({
    query: z.string().describe('The search query'),
    maxResults: z.number().max(5).default(3)
      .describe('Maximum number of results to return'),
  }),
  execute: async ({ query, maxResults }) => {
    // 使用 Tavily API（专为 AI Agent 设计的搜索 API）
    const res = await fetch('https://api.tavily.com/search', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        api_key: process.env.TAVILY_API_KEY,
        query,
        max_results: maxResults,
        include_answer: true,
      }),
    })
 
    const data = await res.json()
 
    return {
      answer: data.answer,
      results: data.results?.map((r: any) => ({
        title: r.title,
        url: r.url,
        snippet: r.content?.slice(0, 300),
      })),
    }
  },
})

export const queryKnowledgeBase = tool({
  description: 'Search the organization knowledge base for internal documents and policies. Use this for questions about company processes, product documentation, or team guidelines.',
  parameters: z.object({
    query: z.string().describe('The search query for the knowledge base'),
    knowledgeBaseId: z.string().uuid().optional()
      .describe('Specific knowledge base ID. If omitted, searches all knowledge bases.'),
  }),
  execute: async ({ query, knowledgeBaseId }, { tenantId }) => {
    if (knowledgeBaseId) {
      const results = await retrieveRelevantChunks(query, knowledgeBaseId, { topK: 3 })
      return formatResults(results)
    }
 
    // 搜索所有知识库
    const kbs = await db.select().from(knowledgeBases)
      .where(eq(knowledgeBases.tenantId, tenantId))
    
    const allResults = []
    for (const kb of kbs) {
      const results = await retrieveRelevantChunks(query, kb.id, { topK: 2 })
      allResults.push(...results.map(r => ({ ...r, knowledgeBaseName: kb.name })))
    }
 
    return formatResults(allResults.sort((a, b) => b.similarity - a.similarity).slice(0, 5))
  },
})
 
function formatResults(results: any[]) {
  return results.map(r => ({
    content: r.content,
    source: r.documentName,
    relevance: Math.round(r.similarity * 100) + '%',
  }))
}

// 需要确认的工具：不提供 execute，客户端处理确认流程
export const deleteProject = tool({
  description: 'Delete a project permanently. This action cannot be undone.',
  parameters: z.object({
    projectId: z.string().uuid().describe('The project ID to delete'),
    confirm: z.literal(true).describe('Must be true to confirm deletion'),
  }),
  // 注意：没有 execute 函数！
})

// components/tool-confirmation.tsx
const { messages, addToolResult } = useChat({
  api: '/api/agent',
  maxSteps: 5,
})
 
// 检测需要确认的工具调用
const pendingConfirmation = messages
  .flatMap(m => m.toolInvocations || [])
  .find(t => t.state === 'call' && !t.result)
 
if (pendingConfirmation) {
  return (
    <div className="rounded-lg border border-yellow-200 bg-yellow-50 p-4">
      <p className="font-medium">⚠️ Confirmation Required</p>
      <p className="mt-1 text-sm text-muted-foreground">
        The AI wants to: <strong>{formatToolName(pendingConfirmation.toolName)}</strong>
      </p>
      <pre className="mt-2 text-xs">{JSON.stringify(pendingConfirmation.args, null, 2)}</pre>
      <div className="mt-3 flex gap-2">
        <button
          onClick={() => {
            // 执行操作并返回结果
            executeConfirmedTool(pendingConfirmation).then(result => {
              addToolResult({ toolCallId: pendingConfirmation.toolCallId, result })
            })
          }}
          className="rounded bg-red-500 px-3 py-1.5 text-sm text-white"
        >
          Confirm
        </button>
        <button
          onClick={() => {
            addToolResult({
              toolCallId: pendingConfirmation.toolCallId,
              result: { error: 'User cancelled the operation' },
            })
          }}
          className="rounded border px-3 py-1.5 text-sm"
        >
          Cancel
        </button>
      </div>
    </div>
  )
}

// lib/ai/agent-limits.ts
 
interface AgentLimits {
  maxStepsPerRequest: number    // 单次请求最多步数
  maxToolCallsPerDay: number    // 每天最多工具调用次数
  maxCostPerRequest: number     // 单次请求最大成本（美元）
  allowedTools: string[]        // 允许使用的工具列表
}
 
const PLAN_LIMITS: Record<string, AgentLimits> = {
  free: {
    maxStepsPerRequest: 2,
    maxToolCallsPerDay: 10,
    maxCostPerRequest: 0.05,
    allowedTools: ['listProjects', 'queryMetrics', 'queryKnowledgeBase'],
  },
  pro: {
    maxStepsPerRequest: 5,
    maxToolCallsPerDay: 100,
    maxCostPerRequest: 0.50,
    allowedTools: ['listProjects', 'createTask', 'queryMetrics', 'webSearch', 'queryKnowledgeBase'],
  },
  business: {
    maxStepsPerRequest: 10,
    maxToolCallsPerDay: 500,
    maxCostPerRequest: 5.00,
    allowedTools: ['*'],  // 所有工具
  },
}
 
export async function checkAgentLimits(tenantId: string, plan: string) {
  const limits = PLAN_LIMITS[plan] || PLAN_LIMITS.free
 
  // 检查每日调用次数
  const todayStart = new Date()
  todayStart.setHours(0, 0, 0, 0)
 
  const [usage] = await db.select({ count: sql<number>`count(*)` })
    .from(aiUsage)
    .where(and(
      eq(aiUsage.tenantId, tenantId),
      eq(aiUsage.operationType, 'tool'),
      gte(aiUsage.createdAt, todayStart),
    ))
 
  if (usage.count >= limits.maxToolCallsPerDay) {
    throw new Error(`Daily tool call limit reached (${limits.maxToolCallsPerDay})`)
  }
 
  return limits
}

// lib/ai/tools/context.ts
 
interface ToolContext {
  tenantId: string
  userId: string
  plan: string
  permissions: string[]
}
 
export function createToolsWithContext(
  tools: Record<string, any>,
  context: ToolContext
): Record<string, any> {
  const wrappedTools: Record<string, any> = {}
 
  for (const [name, t] of Object.entries(tools)) {
    wrappedTools[name] = {
      ...t,
      execute: async (args: any) => {
        // 注入租户上下文
        try {
          const result = await t.execute(args, context)
          // 审计日志
          await logToolExecution(context.tenantId, context.userId, name, args, result)
          return result
        } catch (error) {
          return { error: error instanceof Error ? error.message : 'Tool execution failed' }
        }
      },
    }
  }
 
  return wrappedTools
}

// lib/ai/structured.ts
import { generateObject } from 'ai'
import { openai } from '@ai-sdk/openai'
import { z } from 'zod'
 
// 从自然语言生成结构化任务
export async function parseTaskFromText(text: string) {
  const { object } = await generateObject({
    model: openai('gpt-4o-mini'),
    schema: z.object({
      title: z.string().describe('Task title'),
      description: z.string().describe('Detailed task description'),
      priority: z.enum(['low', 'medium', 'high', 'urgent']),
      estimatedHours: z.number().optional().describe('Estimated hours to complete'),
      tags: z.array(z.string()).describe('Relevant tags'),
    }),
    prompt: `Extract a structured task from this message:\n\n${text}`,
  })
 
  return object  // 类型安全：{ title: string, description: string, ... }
}

import { streamObject } from 'ai'
 
export async function streamReport(data: any) {
  const result = streamObject({
    model: openai('gpt-4o'),
    schema: z.object({
      title: z.string(),
      summary: z.string(),
      sections: z.array(z.object({
        heading: z.string(),
        content: z.string(),
        insights: z.array(z.string()),
      })),
      recommendations: z.array(z.string()),
    }),
    prompt: `Generate a business report based on this data:\n${JSON.stringify(data)}`,
  })
 
  // 客户端可以实时看到报告的各部分逐步生成
  return result.toTextStreamResponse()
}