Vercel AI SDK 流式输出实战

// app/api/chat/route.ts
import { NextRequest } from 'next/server'
import { streamText } from 'ai'
import { openai } from '@ai-sdk/openai'
import { z } from 'zod'
import { buildSystemPrompt } from '@/lib/prompts'
import { rateLimit } from '@/lib/rate-limit'

export const runtime = 'nodejs'
export const maxDuration = 60

const bodySchema = z.object({
  messages: z
    .array(
      z.object({
        role: z.enum(['user', 'assistant']),
        content: z.string().min(1).max(8000),
      }),
    )
    .min(1)
    .max(50),
})

function getClientIp(request: NextRequest): string {
  return (
    request.headers.get('cf-connecting-ip') ||
    request.headers.get('x-forwarded-for')?.split(',')[0]?.trim() ||
    'unknown'
  )
}

export async function POST(request: NextRequest) {
  const ip = getClientIp(request)

  if (!rateLimit(`chat:${ip}`, 20, 60_000)) {
    return new Response(JSON.stringify({ error: '请求过于频繁，请稍后再试' }), {
      status: 429,
      headers: { 'Content-Type': 'application/json' },
    })
  }

  let parsed
  try {
    const json = await request.json()
    parsed = bodySchema.parse(json)
  } catch (error) {
    console.error('[chat parse]', error)
    return new Response(JSON.stringify({ error: '参数格式错误' }), {
      status: 400,
      headers: { 'Content-Type': 'application/json' },
    })
  }

  try {
    const result = streamText({
      model: openai('gpt-4o-mini'),
      system: buildSystemPrompt('chat'),
      messages: parsed.messages,
      temperature: 0.7,
      maxTokens: 1024,
      abortSignal: request.signal,
      onError: error => {
        console.error('[streamText error]', error)
      },
    })

    return result.toDataStreamResponse()
  } catch (error) {
    console.error('[POST /api/chat]', error)
    return new Response(JSON.stringify({ error: '生成失败，请稍后重试' }), {
      status: 500,
      headers: { 'Content-Type': 'application/json' },
    })
  }
}

// lib/prompts/index.ts
export type PromptKey = 'chat' | 'codeReview' | 'summary'

type PromptDefinition = {
  version: string
  template: string
}

const prompts: Record<PromptKey, PromptDefinition> = {
  chat: {
    version: '2026-06-13',
    template: `你是一个严谨的助手。
规则：
1. 不要编造事实。
2. 不确定时明确说明不确定。
3. 回答使用简体中文。`,
  },
  codeReview: {
    version: '2026-06-13',
    template: `你是一名资深工程师，负责代码审查。
重点关注：安全、边界条件、错误处理。
输出使用简体中文。`,
  },
  summary: {
    version: '2026-06-13',
    template: `请将用户提供的内容总结为要点列表，保持客观。`,
  },
}

export function buildSystemPrompt(key: PromptKey, context?: Record<string, string>): string {
  const definition = prompts[key]
  if (!definition) {
    throw new Error(`Unknown prompt key: ${key}`)
  }

  let result = definition.template

  if (context) {
    for (const [name, value] of Object.entries(context)) {
      result = result.replaceAll(`{{${name}}}`, value)
    }
  }

  return result
}

export function getPromptVersion(key: PromptKey): string {
  return prompts[key]?.version ?? 'unknown'
}

// lib/ai-call.ts
import { generateText } from 'ai'
import { openai } from '@ai-sdk/openai'

type CallOptions = {
  system: string
  prompt: string
  retries?: number
  timeoutMs?: number
}

export async function callModelWithRetry(options: CallOptions): Promise<string> {
  const { system, prompt, retries = 2, timeoutMs = 30_000 } = options

  let lastError: unknown = null

  for (let attempt = 0; attempt <= retries; attempt += 1) {
    const controller = new AbortController()
    const timer = setTimeout(() => controller.abort(), timeoutMs)

    try {
      const { text } = await generateText({
        model: openai('gpt-4o-mini'),
        system,
        prompt,
        maxTokens: 1024,
        abortSignal: controller.signal,
      })

      clearTimeout(timer)
      return text
    } catch (error) {
      clearTimeout(timer)
      lastError = error
      console.error(`[callModelWithRetry] attempt ${attempt} failed`, error)

      // 指数退避，避免立即重试打爆上游
      if (attempt < retries) {
        await new Promise(resolve => setTimeout(resolve, 500 * 2 ** attempt))
      }
    }
  }

  throw new Error(`AI call failed after ${retries + 1} attempts: ${String(lastError)}`)
}

// lib/usage.ts
import { prisma } from '@/lib/prisma'

const DAILY_TOKEN_LIMIT = 100_000

export async function canConsume(userId: string): Promise<boolean> {
  if (!userId) return false

  try {
    const today = new Date()
    today.setHours(0, 0, 0, 0)

    const usage = await prisma.aiUsage.findFirst({
      where: { userId, date: today },
    })

    return (usage?.tokens ?? 0) < DAILY_TOKEN_LIMIT
  } catch (error) {
    console.error('[canConsume]', error)
    // 出错时保守拒绝，避免被刷
    return false
  }
}

export async function recordUsage(userId: string, tokens: number): Promise<void> {
  if (!userId || tokens <= 0) return

  try {
    const today = new Date()
    today.setHours(0, 0, 0, 0)

    await prisma.aiUsage.upsert({
      where: { userId_date: { userId, date: today } },
      create: { userId, date: today, tokens },
      update: { tokens: { increment: tokens } },
    })
  } catch (error) {
    console.error('[recordUsage]', error)
  }
}

// 在 streamText 之前
const userId = request.cookies.get('user_id')?.value
if (!userId) {
  return new Response(JSON.stringify({ error: '请先登录' }), { status: 401 })
}

if (!(await canConsume(userId))) {
  return new Response(JSON.stringify({ error: '今日额度已用完' }), { status: 429 })
}

const result = streamText({
  model: openai('gpt-4o-mini'),
  system: buildSystemPrompt('chat'),
  messages: parsed.messages,
  onFinish: async ({ usage }) => {
    await recordUsage(userId, usage.totalTokens ?? 0)
  },
})