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";

export type AIStream = ReturnType<typeof createStreamableValue<string>>;

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 function assertOpenAiKey(): string {
  const apiKey = process.env.OPENAI_API_KEY?.trim();
  if (!apiKey) {
    throw new Error(
      "未配置 OPENAI_API_KEY，请在 .env.local 或 Docker env_file 中设置",
    );
  }
  return apiKey;
}

/** 在 Server Action 内 createStreamableValue 后调用，不可从子函数返回 stream.value */
export function pumpAIStream(
  stream: AIStream,
  system: string,
  user: string,
): void {
  const openai = createOpenAI({
    apiKey: assertOpenAiKey(),
    baseURL: getOpenAiBaseUrl(),
  });

  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;
    });
}