LangChain+OpenAI SSE 스트리밍 끊김·중복토큰 해결

import type { NextRequest } from "next/server";

export const runtime = "nodejs";

function sseFormat(event: string, data: unknown) {
  // MDX 빌드 에러 방지를 위해 본문에서 부등호를 직접 쓰지 않습니다.
  // 코드는 코드블록이므로 안전합니다.
  return `event: ${event}\ndata: ${JSON.stringify(data)}\n\n`;
}

export async function POST(req: NextRequest) {
  const encoder = new TextEncoder();

  const stream = new ReadableStream<Uint8Array>({
    start(controller) {
      const write = (event: string, data: unknown) => {
        controller.enqueue(encoder.encode(sseFormat(event, data)));
      };

      // heartbeat: 15초마다 ping
      const heartbeat = setInterval(() => {
        write("ping", { t: Date.now() });
      }, 15000);

      // 클라이언트가 끊으면 정리
      const abort = () => {
        clearInterval(heartbeat);
        try {
          controller.close();
        } catch {
          // ignore
        }
      };

      // NextRequest는 signal을 제공합니다.
      req.signal.addEventListener("abort", abort);

      // 실제 LLM 스트리밍은 별도 async로 실행
      (async () => {
        try {
          write("ready", { ok: true });

          // TODO: 여기에 LangChain 스트리밍을 연결
          // write("token", { text: "..." })

          write("done", { ok: true });
          clearInterval(heartbeat);
          controller.close();
        } catch (e: any) {
          write("error", { message: e?.message ?? "unknown" });
          clearInterval(heartbeat);
          controller.close();
        }
      })();
    },
  });

  return new Response(stream, {
    headers: {
      "Content-Type": "text/event-stream; charset=utf-8",
      "Cache-Control": "no-cache, no-transform",
      Connection: "keep-alive",
      "X-Accel-Buffering": "no",
    },
  });
}

import { ChatOpenAI } from "@langchain/openai";
import { CallbackManager } from "@langchain/core/callbacks/manager";

type TokenSink = (delta: string) => void;

export async function runStreamingLLM(input: string, onToken: TokenSink) {
  const callbackManager = CallbackManager.fromHandlers({
    handleLLMNewToken(token) {
      onToken(token);
    },
  });

  const model = new ChatOpenAI({
    model: "gpt-4o-mini",
    temperature: 0.2,
    streaming: true,
    callbackManager,
  });

  const res = await model.invoke(input);
  return res;
}

function diffNewSuffix(sentText: string, nextFullText: string) {
  if (nextFullText.startsWith(sentText)) {
    return nextFullText.slice(sentText.length);
  }

  // prefix가 안 맞는 경우: 재시도/리셋/모델 변동 등
  // 여기서는 가장 보수적으로 공통 prefix를 찾습니다.
  const max = Math.min(sentText.length, nextFullText.length);
  let i = 0;
  while (i < max && sentText[i] === nextFullText[i]) i += 1;

  return nextFullText.slice(i);
}

type WriteFn = (event: string, data: unknown) => void;

async function streamWithRetry(params: {
  input: string;
  write: WriteFn;
  maxAttempts: number;
}) {
  const { input, write, maxAttempts } = params;

  let attempt = 0;
  let sentText = "";
  let fullText = "";

  while (attempt < maxAttempts) {
    attempt += 1;
    write("meta", { attempt });

    try {
      // 재시도 시에도 동일 input을 쓰되, dedup로 suffix만 내보냄
      await runStreamingLLM(input, (delta) => {
        fullText += delta;

        const suffix = diffNewSuffix(sentText, fullText);
        if (!suffix) return;

        sentText += suffix;
        write("token", { text: suffix });
      });

      write("done", { ok: true, attempts: attempt });
      return;
    } catch (e: any) {
      write("warn", {
        attempt,
        message: e?.message ?? "stream error",
      });

      // 지수 백오프 + 지터
      const backoffMs = Math.min(2000 * 2 ** (attempt - 1), 15000);
      const jitterMs = Math.floor(Math.random() * 250);
      await new Promise((r) => setTimeout(r, backoffMs + jitterMs));

      // 주의: 여기서 fullText를 유지할지 리셋할지 정책이 필요
      // - 유지하면 prefix 비교가 더 잘 맞을 때가 많지만,
      // - 내부적으로 모델이 다른 경로로 생성하면 불일치가 커질 수 있음
      // 실무에서는 "sentText"는 유지, "fullText"는 리셋 후 재구성하는 편이 안전합니다.
      fullText = sentText;
    }
  }

  write("error", { message: "max retry exceeded" });
}

type FlushFn = (chunk: string) => void;

function createChunkBuffer(flush: FlushFn, windowMs: number) {
  let buf = "";
  let timer: any = null;

  const schedule = () => {
    if (timer) return;
    timer = setTimeout(() => {
      timer = null;
      if (!buf) return;
      const out = buf;
      buf = "";
      flush(out);
    }, windowMs);
  };

  return {
    push(delta: string) {
      buf += delta;
      schedule();
    },
    flushNow() {
      if (timer) {
        clearTimeout(timer);
        timer = null;
      }
      if (!buf) return;
      const out = buf;
      buf = "";
      flush(out);
    },
  };
}

const buffer = createChunkBuffer((chunk) => {
  write("token", { text: chunk });
}, 50);

await runStreamingLLM(input, (delta) => {
  // dedup을 통과한 suffix만 buffer에 넣는 구조를 권장
  buffer.push(delta);
});

buffer.flushNow();
write("done", { ok: true });