next-tjwater-drainage-frontend/features/agent/components/streaming-token-response-state.ts

export const DEFAULT_STREAMING_TOKEN_LOCALE = "zh";
export const STREAMING_TOKEN_BLOCK_SIZE = 4;

const FALLBACK_CHUNK_SIZE = 2;

export function splitStreamingText(text: string, locale = DEFAULT_STREAMING_TOKEN_LOCALE) {
  if (!text) {
    return [];
  }

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

  if (segmenter) {
    return Array.from(segmenter.segment(text), (item) => item.segment).filter(Boolean);
  }

  return splitStreamingTextFallback(text);
}

export function splitStreamingTextFallback(text: string) {
  const codePoints = Array.from(text);
  const chunks: string[] = [];
  for (let index = 0; index < codePoints.length; index += FALLBACK_CHUNK_SIZE) {
    chunks.push(codePoints.slice(index, index + FALLBACK_CHUNK_SIZE).join(""));
  }
  return chunks;
}

function takeStreamingTokenBlock(pending: string[], blockSize = STREAMING_TOKEN_BLOCK_SIZE) {
  return pending.splice(0, blockSize).join("");
}

export type StreamingMarkdownBufferState = {
  content: string;
  displayContent: string;
  locale: string;
};

type FenceMarker = {
  char: "`" | "~";
  length: number;
};

const FENCE_START_PATTERN = /^ {0,3}(`{3,}|~{3,})/;
const HEADING_OPENER_PATTERN = /^ {0,3}#{1,6}[ \t]+\S/;
const PARTIAL_HEADING_OPENER_PATTERN = /^ {0,3}#{1,6}[ \t]*$/;
const LIST_OPENER_PATTERN = /^ {0,3}(?:[-+*]|\d{1,9}[.)])[ \t]+\S/;
const PARTIAL_LIST_OPENER_PATTERN = /^ {0,3}(?:[-+*]|\d{1,9}[.)])[ \t]*$/;
const BLOCKQUOTE_OPENER_PATTERN = /^ {0,3}>[ \t]?\S/;
const PARTIAL_BLOCKQUOTE_OPENER_PATTERN = /^ {0,3}>[ \t]*$/;
const HORIZONTAL_RULE_PATTERN = /^ {0,3}(?:-{3,}|\*{3,}|_{3,})[ \t]*(?:\r?\n|$)/;
const TABLE_SEPARATOR_PATTERN = /^ {0,3}\|?(?:\s*:?-{3,}:?\s*\|)+\s*:?-{3,}:?\s*\|?\s*$/;
const POSSIBLE_TABLE_HEADER_PATTERN = /^ {0,3}\|.*\||^ {0,3}\S.*\|.*\|/;
const LINE_START_PATTERN = /(?:^|\n)(?= {0,3}(?:#{1,6}|[-+*]|\d{1,9}[.)]|>|`{3,}|~{3,}|\|))/g;

function readFenceStart(line: string): FenceMarker | null {
  const match = FENCE_START_PATTERN.exec(line);
  if (!match) {
    return null;
  }

  const marker = match[1];
  return {
    char: marker[0] as FenceMarker["char"],
    length: marker.length,
  };
}

function readLineEnd(content: string) {
  const newlineIndex = content.indexOf("\n");
  if (newlineIndex === -1) {
    return content.length;
  }
  return newlineIndex + 1;
}

function isTableLine(line: string) {
  return line.includes("|") && line.trim().length > 0;
}

function isTableSeparatorLine(line: string) {
  return TABLE_SEPARATOR_PATTERN.test(line);
}

type MarkdownOpenerMatch = {
  chunk: string;
};

export function createStreamingMarkdownBufferState(
  content = "",
  displayContent = "",
  locale = DEFAULT_STREAMING_TOKEN_LOCALE
): StreamingMarkdownBufferState {
  return { content, displayContent, locale };
}

export function resolveStreamingMarkdownDisplayContent(
  content: string,
  displayContent: string,
  streamDone: boolean
) {
  if (streamDone) {
    return content;
  }

  return content.startsWith(displayContent) ? displayContent : "";
}

function isLineStart(content: string, offset: number) {
  return offset === 0 || content[offset - 1] === "\n";
}

function takeMarkdownTableOpener(content: string): MarkdownOpenerMatch | null {
  const headerEnd = readLineEnd(content);
  const headerLine = content.slice(0, headerEnd).replace(/\r?\n$/, "");

  if (!POSSIBLE_TABLE_HEADER_PATTERN.test(headerLine)) {
    return null;
  }

  if (headerEnd >= content.length && !content.endsWith("\n")) {
    return { chunk: "" };
  }

  const afterHeader = content.slice(headerEnd);
  const separatorEnd = readLineEnd(afterHeader);
  const separatorLine = afterHeader.slice(0, separatorEnd).replace(/\r?\n$/, "");

  if (!separatorLine) {
    return { chunk: "" };
  }

  if (!isTableSeparatorLine(separatorLine)) {
    return null;
  }

  return {
    chunk: content.slice(0, headerEnd + separatorEnd),
  };
}

function isStreamingTableOpen(displayContent: string) {
  const lines = displayContent.replace(/\r\n/g, "\n").split("\n");

  while (lines.at(-1) === "") {
    lines.pop();
  }

  const blockLines: string[] = [];
  for (let index = lines.length - 1; index >= 0; index -= 1) {
    const line = lines[index];
    if (!line.trim()) {
      break;
    }
    if (!isTableLine(line)) {
      break;
    }
    blockLines.unshift(line);
  }

  return (
    blockLines.length >= 2 &&
    blockLines.every(isTableLine) &&
    blockLines.some(isTableSeparatorLine)
  );
}

function takeMarkdownOpener(
  content: string,
  { allowTable = true }: { allowTable?: boolean } = {}
): MarkdownOpenerMatch | null {
  if (!content) {
    return null;
  }

  const lineEnd = readLineEnd(content);
  const firstLine = content.slice(0, lineEnd);
  const lineWithoutEnding = firstLine.replace(/\r?\n$/, "");
  const fence = readFenceStart(lineWithoutEnding);

  if (fence) {
    return firstLine.endsWith("\n") ? { chunk: firstLine } : { chunk: "" };
  }

  const tableOpener = allowTable ? takeMarkdownTableOpener(content) : null;
  if (tableOpener) {
    return tableOpener;
  }

  const firstNonSyntaxIndex = (pattern: RegExp) => {
    const match = pattern.exec(content);
    return match ? match[0].length : 0;
  };

  const headingLength = firstNonSyntaxIndex(HEADING_OPENER_PATTERN);
  if (headingLength > 0) {
    return { chunk: content.slice(0, headingLength) };
  }
  if (PARTIAL_HEADING_OPENER_PATTERN.test(lineWithoutEnding)) {
    return { chunk: "" };
  }

  const listLength = firstNonSyntaxIndex(LIST_OPENER_PATTERN);
  if (listLength > 0) {
    return { chunk: content.slice(0, listLength) };
  }
  if (PARTIAL_LIST_OPENER_PATTERN.test(lineWithoutEnding)) {
    return { chunk: "" };
  }

  const blockquoteLength = firstNonSyntaxIndex(BLOCKQUOTE_OPENER_PATTERN);
  if (blockquoteLength > 0) {
    return { chunk: content.slice(0, blockquoteLength) };
  }
  if (PARTIAL_BLOCKQUOTE_OPENER_PATTERN.test(lineWithoutEnding)) {
    return { chunk: "" };
  }

  const ruleMatch = HORIZONTAL_RULE_PATTERN.exec(content);
  if (ruleMatch) {
    return { chunk: ruleMatch[0] };
  }

  return null;
}

export function isMarkdownOpenerBoundary(content: string, offset = 0) {
  if (!isLineStart(content, offset)) {
    return false;
  }

  const opener = takeMarkdownOpener(content.slice(offset));
  return Boolean(opener?.chunk);
}

function findNextPotentialMarkdownOpenerIndex(content: string) {
  LINE_START_PATTERN.lastIndex = 0;

  for (const match of content.matchAll(LINE_START_PATTERN)) {
    const index = match.index + (match[0] === "\n" ? 1 : 0);
    if (index > 0) {
      return index;
    }
  }

  return -1;
}

function takeSegmentedMarkdownChunk(
  content: string,
  locale: string,
  blockSize = STREAMING_TOKEN_BLOCK_SIZE
) {
  const pendingSegments = splitStreamingText(content, locale);
  const candidate = takeStreamingTokenBlock(pendingSegments, blockSize);

  if (!candidate) {
    return "";
  }

  const nextOpenerIndex = findNextPotentialMarkdownOpenerIndex(candidate);
  return nextOpenerIndex > 0 ? candidate.slice(0, nextOpenerIndex) : candidate;
}

export function takeNextStreamingMarkdownChunk(
  state: StreamingMarkdownBufferState,
  blockSize = STREAMING_TOKEN_BLOCK_SIZE
) {
  const displayContent = state.content.startsWith(state.displayContent)
    ? state.displayContent
    : "";
  const pendingContent = state.content.slice(displayContent.length);

  if (!pendingContent) {
    return {
      chunk: "",
      state: { ...state, displayContent },
    };
  }

  const opener = isLineStart(state.content, displayContent.length)
    ? takeMarkdownOpener(pendingContent, {
        allowTable: !isStreamingTableOpen(displayContent),
      })
    : null;
  const chunk = opener
    ? opener.chunk
    : takeSegmentedMarkdownChunk(pendingContent, state.locale, blockSize);
  const nextDisplayContent = `${displayContent}${chunk}`;

  return {
    chunk,
    state: {
      ...state,
      displayContent: nextDisplayContent,
    },
  };
}