NestJS + OpenAI 实现流式输出

9 阅读3分钟

在现代 web 应用中,AI 交互的实时性体验越来越重要。本文将详细介绍如何在 NestJS 中集成 OpenAI,并实现流式输出功能,让用户能够看到 AI 回复的实时过程,而不是等待完整回复。

安装依赖

首先,我们需要安装必要的依赖包:

pnpm i @langchain/openai @langchain/core

依赖版本

{
  "dependencies": {
    "@langchain/core": "^1.1.39",
    "@langchain/openai": "^1.4.4"
  }
}

NestJS 配置

为了安全管理 OpenAI API 密钥等配置信息,我们需要设置 NestJS 的环境配置。这里参考了 NestJS 多环境 YAML 配置方案 的实现方式。

config.yaml 中添加 OpenAI 相关配置:

# config/config.yaml
openai:
  apiKey: 'your-api-key'
  model: 'gpt-3.5-turbo'
  baseURL: 'https://api.openai.com/v1'

服务端实现代码

类型定义

首先,我们定义全局类型别名,方便在整个项目中使用:

// types\global.d.ts

import type { Request, Response } from 'express'

declare global {
  /**
   * Express 请求对象类型别名
   * - 说明:用于描述 HTTP 请求的完整信息,包含请求路径、请求参数、请求头、请求体、Cookie 等核心内容
   * - 用途:通常用于定义 Express 接口的请求参数类型,约束和解析请求数据
   */
  type ExpressRequest = Request

  /**
   * Express 响应对象类型别名
   * - 说明:用于描述 HTTP 响应的配置信息,包含响应状态码、响应头、响应体、重定向等核心功能
   * - 用途:通常用于定义 Express 接口的响应格式,返回指定结构的数据给客户端
   */
  type ExpressResponse = Response
}

AI 服务实现

创建 AI 服务,实现流式调用 OpenAI API 的核心逻辑:

// src\modules\ai\ai.service.ts

import { Injectable } from '@nestjs/common'
import { ConfigService } from '@nestjs/config'
import { ChatPromptTemplate } from '@langchain/core/prompts'
import { ChatOpenAI, ChatOpenAIFields } from '@langchain/openai'

@Injectable()
export class AiService {
  constructor(private readonly configService: ConfigService) {}

  /** 调用 OpenAI 模型(流式) */
  public async streamChat(prompt: string, response: ExpressResponse) {
    try {
      // 1. 构造提示词模板
      const chatPromptTemplate = ChatPromptTemplate.fromMessages([
        ['system', '你是一个专业的前端开发人员'],
        ['human', '{input}'],
      ])

      // 2. 拼接链
      const chain = chatPromptTemplate.pipe(this.llm)

      // 3. 流式调用(关键:传入对象,不是字符串)
      const stream = await chain.stream({ input: prompt })

      // 4. 处理流式响应
      for await (const chunk of stream) {
        const content = chunk.content || ''
        if (!content) continue // 过滤空内容
        response.write(`data: ${JSON.stringify({ content })}\n\n`)
      }

      // 5. 发送完成通知
      response.write(`data: ${JSON.stringify({ status: 'DONE' })}\n\n`)
    } catch (error) {
      // 处理错误
      response.write(`data: ${JSON.stringify({ error: error.message })}\n\n`)
    } finally {
      // 6. 结束响应
      response.end()
    }
  }

  /** 获取 OpenAI 模型实例 */
  private get llm(): ChatOpenAI {
    const options: ChatOpenAIFields = {}
    options.apiKey = this.configService.get('openai.apiKey')
    options.model = this.configService.get('openai.model')
    options.temperature = 0.5
    options.streaming = true // 开启流式输出
    options.configuration = {}
    options.configuration.baseURL = this.configService.get('openai.baseURL')
    return new ChatOpenAI(options)
  }
}

控制器实现

创建控制器,处理客户端的流式聊天请求:

// src\modules\ai\ai.controller.ts

import { AiService } from './ai.service'
import { Body, Controller, Post, Res } from '@nestjs/common'

@Controller('ai')
export class AiController {
  constructor(private readonly aiService: AiService) {}

  @Post('chat/stream')
  public chatStream(@Body('prompt') prompt: string, @Res({ passthrough: true }) response: ExpressResponse) {
    // 设置 SSE 响应头
    response.setHeader('Content-Type', 'text/event-stream') // 设置响应头为事件流
    response.setHeader('Cache-Control', 'no-cache') // 禁用缓存
    response.setHeader('Connection', 'keep-alive') // 保持连接
    return this.aiService.streamChat(prompt, response)
  }
}

客户端实现代码

在前端,我们使用 Vue 3 的 Composition API 实现流式接收和展示 AI 回复:

<script setup lang="ts">
import { ref } from 'vue'
import axios from 'axios'

const prompt = ref('说一下岳阳楼的故事')
const content = ref('')
const loading = ref(false)

async function chat() {
  if (!prompt.value) {
    alert('请先输入问题')
    return
  }

  content.value = ''
  loading.value = true

  try {
    const url = `http://localhost:3000/api/ai/chat/stream`
    const data = { prompt: prompt.value }
    const response = await axios.post(url, data, { responseType: 'stream', adapter: 'fetch' })

    const reader = response.data.getReader()
    const decoder = new TextDecoder()

    while (true) {
      const { done, value } = await reader.read()
      if (done) break
      const chunk = decoder.decode(value, { stream: true })
      const lines = chunk.split('\n').filter((i: string) => i.startsWith('data: '))
      for (const line of lines) {
        try {
          const data = JSON.parse(line.replace('data: ', ''))
          if (data.status === 'DONE') return
          if (data.error) throw new Error(data.error)
          if (data.content) {
            content.value += data.content
            console.log('🚀 ~ data.content :', data.content)
          }
        } catch (error) {
          console.error('解析数据失败:', error)
        }
      }
    }
  } catch (error) {
    console.error('请求失败:', error)
    content.value = `请求失败: ${error.message}`
  } finally {
    loading.value = false
  }
}
</script>

<template>
  <div class="chat-container">
    <div class="input-area">
      <input
        v-model="prompt"
        type="text"
        placeholder="请输入问题..."
        :disabled="loading"
      />
      <button @click="chat" :disabled="loading">
        {{ loading ? '发送中...' : '发送' }}
      </button>
    </div>
    <div class="response-area" v-if="content">
      <h3>AI 回复:</h3>
      <div class="content">{{ content }}</div>
    </div>
  </div>
</template>

<style scoped>
.chat-container {
  max-width: 800px;
  margin: 0 auto;
  padding: 20px;
}

.input-area {
  display: flex;
  gap: 10px;
  margin-bottom: 20px;
}

input {
  flex: 1;
  padding: 10px;
  border: 1px solid #ddd;
  border-radius: 4px;
}

button {
  padding: 0 20px;
  background-color: #42b983;
  color: white;
  border: none;
  border-radius: 4px;
  cursor: pointer;
}

button:disabled {
  background-color: #ccc;
  cursor: not-allowed;
}

.response-area {
  border: 1px solid #ddd;
  border-radius: 4px;
  padding: 15px;
  min-height: 200px;
  white-space: pre-wrap;
}
</style>

实现原理

  1. 服务端实现

    • 使用 text/event-stream 响应头启用 Server-Sent Events (SSE)
    • 通过 LangChain 的 stream 方法获取 OpenAI 的流式响应
    • 逐块处理响应数据并通过 SSE 发送给客户端
    • 发送完成信号标记流式传输结束
  2. 客户端实现

    • 使用 responseType: 'stream' 接收流式响应
    • 使用 getReader() 读取响应流
    • 逐块解码并解析数据
    • 实时更新 UI 展示 AI 回复

注意事项

  1. API 密钥安全:确保 OpenAI API 密钥不会被提交到代码仓库,使用环境配置管理
  2. 错误处理:添加适当的错误处理,确保流式传输过程中的错误能够被捕获和处理
  3. 性能优化:对于较长的回复,可以考虑添加节流处理,避免 UI 更新过于频繁
  4. CORS 配置:如果前端和后端分离部署,需要配置适当的 CORS 策略