zhimingge/lib/ai/stream.ts

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

function getOpenAiBaseUrl(): string {
  return (
    process.env.OPENAI_BASE_URL ??
    process.env.OPENAI_API_BASE ??
    "https://op.bz121.com/v1"
  );
}

const model =
  process.env.OPENAI_MODEL ?? "huihui_ai/gemma-4-abliterated:e4b";

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

  const openai = createOpenAI({
    apiKey,
    baseURL: getOpenAiBaseUrl(),
  });

  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 <= 6) {
        stream.update(buffer);
        buffer = "";
      } else {
        const chunk = buffer.slice(0, 6);
        buffer = buffer.slice(6);
        stream.update(chunk);
      }
    }, 60);

    (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((err) => {
        const message = err instanceof Error ? err.message : String(err);
        stream.update(ERROR_PREFIX + message);
      })
      .finally(() => {
        done = true;
      });

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