zhimingge/lib/ai/stream.ts

"use server";

import { streamText } from "ai";
import { createOpenAI } from "@ai-sdk/openai";
import { createStreamableValue } from "ai/rsc";
import { ERROR_PREFIX } from "@/lib/constant";

const model =
  process.env.OPENAI_MODEL ?? "huihui_ai/gemma-4-abliterated:e4b";
const openai = createOpenAI({
  baseURL: process.env.OPENAI_BASE_URL ?? "https://op.bz121.com/v1",
});

const STREAM_INTERVAL = 60;
const MAX_SIZE = 6;

export async function streamAIResponse(
  system: string,
  user: string,
): Promise<{ data?: ReturnType<typeof createStreamableValue<string>>["value"]; error?: string }> {
  if (!process.env.OPENAI_API_KEY?.trim()) {
    return { error: "未配置 OPENAI_API_KEY，请在 .env.local 或 Docker env_file 中设置" };
  }

  const stream = createStreamableValue<string>();

  try {
    const { fullStream } = streamText({
      temperature: 0.5,
      model: openai(model),
      messages: [
        { role: "system", content: system },
        { role: "user", content: user },
      ],
      maxRetries: 0,
    });

    let buffer = "";
    let done = false;
    const intervalId = setInterval(() => {
      if (done && buffer.length === 0) {
        clearInterval(intervalId);
        stream.done();
        return;
      }
      if (buffer.length <= MAX_SIZE) {
        stream.update(buffer);
        buffer = "";
      } else {
        const chunk = buffer.slice(0, MAX_SIZE);
        buffer = buffer.slice(MAX_SIZE);
        stream.update(chunk);
      }
    }, STREAM_INTERVAL);

    (async () => {
      for await (const part of fullStream) {
        switch (part.type) {
          case "text-delta":
            buffer += part.textDelta;
            break;
          case "error": {
            const err = part.error as { message?: string };
            stream.update(ERROR_PREFIX + (err.message ?? String(part.error)));
            break;
          }
        }
      }
    })()
      .catch(console.error)
      .finally(() => {
        done = true;
      });

    return { data: stream.value };
  } catch (err) {
    stream.done();
    const message = err instanceof Error ? err.message : String(err);
    return { error: message };
  }
}