AI Streaming Response với SSE trong NestJS & NextJS

// src/chat/chat.controller.ts
import { Controller, Post, Body, Sse, Res } from '@nestjs/common';
import { Observable, from, map } from 'rxjs';
import OpenAI from 'openai';
import { Response } from 'express';

@Controller('chat')
export class ChatController {
  private openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });

  @Post('stream')
  @Sse()  // Khai báo SSE endpoint
  async streamChat(@Body() body: { message: string }): Promise<Observable<MessageEvent>> {
    const stream = await this.openai.chat.completions.create({
      model: 'gpt-4o-mini',
      messages: [{ role: 'user', content: body.message }],
      stream: true,
    });

    return new Observable(subscriber => {
      (async () => {
        try {
          for await (const chunk of stream) {
            const text = chunk.choices[0]?.delta?.content ?? '';
            if (text) {
              subscriber.next({ data: JSON.stringify({ text }) } as MessageEvent);
            }
          }
          subscriber.next({ data: '[DONE]' } as MessageEvent);
          subscriber.complete();
        } catch (err) {
          subscriber.error(err);
        }
      })();
    });
  }
}

@Post('stream/raw')
async streamRaw(
  @Body() body: { message: string; systemPrompt?: string },
  @Res() res: Response,
) {
  // Cấu hình headers cho SSE
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');
  res.setHeader('Access-Control-Allow-Origin', '*');
  res.flushHeaders();

  const stream = await this.openai.chat.completions.create({
    model: 'gpt-4o-mini',
    messages: [
      ...(body.systemPrompt ? [{ role: 'system' as const, content: body.systemPrompt }] : []),
      { role: 'user', content: body.message },
    ],
    stream: true,
  });

  let fullText = '';

  try {
    for await (const chunk of stream) {
      const text = chunk.choices[0]?.delta?.content ?? '';
      if (text) {
        fullText += text;
        res.write(`data: ${JSON.stringify({ text })}\n\n`);
      }
    }

    // Gửi signal kết thúc kèm full text
    res.write(`data: ${JSON.stringify({ done: true, fullText })}\n\n`);
  } catch (error) {
    res.write(`data: ${JSON.stringify({ error: 'Stream error' })}\n\n`);
  } finally {
    res.end();
  }
}

@Post('rag/stream')
async ragStream(@Body() body: { question: string }, @Res() res: Response) {
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.flushHeaders();

  // 1. Tìm context từ vector DB (không stream bước này)
  const queryVector = await this.openai.createEmbedding(body.question);
  const chunks = await this.qdrant.search(queryVector, 5);
  const context = chunks.map(c => c.text).join('\n\n');

  // Thông báo cho client biết đã tìm xong context
  res.write(`data: ${JSON.stringify({ type: 'context_ready', sources: chunks.map(c => c.source) })}\n\n`);

  // 2. Stream câu trả lời từ LLM
  const stream = await this.openai.client.chat.completions.create({
    model: 'gpt-4o-mini',
    messages: [
      { role: 'system', content: `Trả lời dựa trên ngữ cảnh:\n${context}` },
      { role: 'user', content: body.question },
    ],
    stream: true,
  });

  for await (const chunk of stream) {
    const text = chunk.choices[0]?.delta?.content ?? '';
    if (text) {
      res.write(`data: ${JSON.stringify({ type: 'text', text })}\n\n`);
    }
  }

  res.write(`data: ${JSON.stringify({ type: 'done' })}\n\n`);
  res.end();
}

// app/api/chat/stream/route.ts
import { NextRequest } from 'next/server';

export async function POST(request: NextRequest) {
  const body = await request.json();

  const backendResponse = await fetch(`${process.env.API_URL}/chat/stream`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(body),
  });

  // Proxy stream từ backend sang client
  return new Response(backendResponse.body, {
    headers: {
      'Content-Type': 'text/event-stream',
      'Cache-Control': 'no-cache',
    },
  });
}