next-tjwater-agent/src/routes/chatTokenSmoother.ts

import { smoothStream, type TextStreamPart } from "ai";

type TokenSmootherOptions = {
  enabled: boolean;
  delayMs: number;
  locale: string;
  sessionId: string;
  write: (event: string, data: Record<string, unknown>) => void;
};

type SmoothTextStreamPart = TextStreamPart<Record<string, never>>;
type SmoothTextDeltaPart = Extract<SmoothTextStreamPart, { type: "text-delta" }>;
type SmoothTextEndPart = Extract<SmoothTextStreamPart, { type: "text-end" }>;

const segmenters = new Map<string, Intl.Segmenter | null>();
const CJK_SCRIPT_PATTERN = /[\p{Script=Han}\p{Script=Hiragana}\p{Script=Katakana}\p{Script=Hangul}]/u;
const CJK_PUNCTUATION_PATTERN = /[。！？!?，、；：,;:]/u;
const CJK_FALLBACK_TRIGGER_CHARS = 24;
const CJK_FALLBACK_MIN_CHARS = 12;
const CJK_FALLBACK_TARGET_CHARS = 18;

function getTokenSmoothingSegmenter(locale: string) {
  if (segmenters.has(locale)) {
    return segmenters.get(locale) ?? null;
  }

  const segmenter =
    typeof Intl !== "undefined" && "Segmenter" in Intl
      ? new Intl.Segmenter(locale, { granularity: "word" })
      : null;
  segmenters.set(locale, segmenter);
  return segmenter;
}

function sliceCodePoints(value: string, count: number) {
  return Array.from(value).slice(0, count).join("");
}

function codePointLength(value: string) {
  return Array.from(value).length;
}

function isCjkSmoothingContent(content: string, locale: string) {
  return (
    CJK_SCRIPT_PATTERN.test(content) ||
    (/^(zh|ja|ko)(?:-|$)/i.test(locale) && CJK_PUNCTUATION_PATTERN.test(content))
  );
}

function detectCjkFallbackChunk(content: string, locale: string) {
  if (codePointLength(content) < CJK_FALLBACK_TRIGGER_CHARS) {
    return null;
  }

  const segmenter = getTokenSmoothingSegmenter(locale);
  if (segmenter) {
    let candidate: string | null = null;
    for (const segment of segmenter.segment(content)) {
      const nextCandidate = content.slice(0, segment.index + segment.segment.length);
      const nextLength = codePointLength(nextCandidate);
      if (nextLength > CJK_FALLBACK_TARGET_CHARS) {
        break;
      }
      if (nextLength >= CJK_FALLBACK_MIN_CHARS) {
        candidate = nextCandidate;
      }
    }
    if (candidate) {
      return candidate;
    }
  }

  return sliceCodePoints(content, CJK_FALLBACK_TARGET_CHARS);
}

export function detectTokenSmoothingChunk(content: string, locale = "zh") {
  if (!content) {
    return null;
  }

  if (isCjkSmoothingContent(content, locale)) {
    return detectCjkFallbackChunk(content, locale);
  }

  const wordChunk = content.match(/^\S+\s*/u)?.[0];
  if (wordChunk) {
    return wordChunk;
  }

  return sliceCodePoints(content, 1);
}

export function createTokenSmoother({
  enabled,
  delayMs,
  locale,
  sessionId,
  write,
}: TokenSmootherOptions) {
  const emitChunk = (content: string) => {
    write("token", {
      session_id: sessionId,
      content,
    });
  };

  if (!enabled) {
    return {
      writeToken(content: string) {
        if (content) {
          emitChunk(content);
        }
      },
      async flush() {
        return;
      },
    };
  }

  const textPartId = `smooth-text-${sessionId}`;
  const stream = smoothStream<Record<string, never>>({
    delayInMs: delayMs,
    chunking: (buffer) => detectTokenSmoothingChunk(buffer, locale),
  })({ tools: {} });
  const writer = stream.writable.getWriter();
  const reader = stream.readable.getReader();
  let closed = false;
  let writeChain = Promise.resolve();
  const readPromise = (async () => {
    while (true) {
      const result = await reader.read();
      if (result.done) {
        return;
      }
      if (result.value.type === "text-delta") {
        emitChunk(result.value.text);
      }
    }
  })();

  return {
    writeToken(content: string) {
      if (!content) {
        return;
      }
      if (closed) {
        return;
      }
      writeChain = writeChain.then(() =>
        writer.write({
          type: "text-delta",
          id: textPartId,
          text: content,
        } satisfies SmoothTextDeltaPart),
      );
    },
    async flush() {
      if (closed) {
        await readPromise;
        return;
      }
      closed = true;
      await writeChain;
      await writer.write({
        type: "text-end",
        id: textPartId,
      } satisfies SmoothTextEndPart);
      await writer.close();
      await readPromise;
    },
  };
}