372 lines · 12.4 KB · TypeScript
RawDownload
import { randomUUID } from "node:crypto";import {  parseToolCallsFromContent,  stripToolCallsFromContent,  type ToolCall,} from "./tool-call-parser";import {  REASONING_FIELDS,  firstReasoningField,  createThinkRewriter,  thinkingTagPrefixIsPartial,} from "./reasoning"; export interface StreamUsage {  prompt_tokens: number;  completion_tokens: number;  reasoning_tokens?: number;  cache_read_tokens?: number;  cache_write_tokens?: number;} export interface ToolCallStreamOptions {  bufferImplicitReasoningContent?: boolean;} // DeepSeek V4 occasionally emits these control tokens as visible text around// long, tool-heavy conversations. They are prompt delimiters, never intended// for the client, and replaying them in the next assistant message amplifies// the leak. The tokenizer normally hides them, but strip the literal fallback// here as well so every controller client gets a clean stream.const stripLeakedDeepSeekControlTokens = (text: string): string =>  text    .replaceAll("<|begin▁of▁sentence|>", "")    .replaceAll("<|end▁of▁sentence|>", ""); export const createToolCallStream = (  source: ReadableStream<Uint8Array>,  onUsage?: (usage: StreamUsage) => void,  onFirstToken?: () => void,  options: ToolCallStreamOptions = {},): ReadableStream<Uint8Array> => {  const decoder = new TextDecoder();  const encoder = new TextEncoder();  let buffer = "";  let pendingEventLines: string[] = [];  let visibleContentBuffer = "";  let toolCallsFound = false;  let usageTracked = false;  let firstTokenTracked = false;  const contentHistory = new Map<string, { text: string; snapshot: boolean }>();  const reasoningHistory = new Map<string, { text: string; snapshot: boolean }>();  const replayCursors = new Map<string, number>();  const stripToolXmlDelta = (text: string): string => {    return stripToolCallsFromContent(stripLeakedDeepSeekControlTokens(text));  };   const normalizeTextDelta = (    history: Map<string, { text: string; snapshot: boolean }>,    key: string,    text: string,    forceSnapshot = false,  ): string => {    if (!text) return text;    const previous = history.get(key) ?? { text: "", snapshot: forceSnapshot };    const replayCursor = replayCursors.get(key);    if (replayCursor !== undefined) {      const expected = previous.text.slice(replayCursor, replayCursor + text.length);      if (expected === text) {        const nextCursor = replayCursor + text.length;        if (nextCursor >= previous.text.length) replayCursors.delete(key);        else replayCursors.set(key, nextCursor);        return "";      }      replayCursors.delete(key);      const resurrected = previous.text.slice(0, replayCursor);      const merged = resurrected + text;      history.set(key, { text: previous.text + merged, snapshot: false });      return merged;    }    const isCumulative =      previous.text.length > 0 &&      text.length > previous.text.length &&      text.startsWith(previous.text);    const shouldSlice = forceSnapshot || previous.snapshot || isCumulative;     if (shouldSlice) {      history.set(key, { text, snapshot: true });      return isCumulative ? text.slice(previous.text.length) : text;    }     if (      text.trim() !== "" &&      previous.text.length > text.length &&      previous.text.startsWith(text)    ) {      replayCursors.set(key, text.length);      return "";    }     history.set(key, { text: previous.text + text, snapshot: false });    return text;  };   const contentThink = createThinkRewriter({    bufferImplicitReasoningContent: Boolean(options.bufferImplicitReasoningContent),  });  const reasoningThink = createThinkRewriter();   const enqueueLine = (    controller: TransformStreamDefaultController<Uint8Array>,    line: string,  ): void => {    controller.enqueue(encoder.encode(`${line}\n`));  };  const enqueueDataEvent = (    controller: TransformStreamDefaultController<Uint8Array>,    dataLine: string,  ): void => {    enqueueLine(controller, dataLine);    enqueueLine(controller, "");  };   const buildToolCallChunk = (toolCalls: ToolCall[]): string => {    const payload = {      id: `chatcmpl-${randomUUID().slice(0, 8)}`,      choices: [        {          index: 0,          delta: { tool_calls: toolCalls },          finish_reason: "tool_calls",        },      ],    };    return `data: ${JSON.stringify(payload)}`;  };   const buildFlushChunk = (payload: {    content?: string;    reasoning_content?: string;  }): string | null => {    const content = payload.content ?? "";    const reasoning = payload.reasoning_content ?? "";    if (!content && !reasoning) return null;    const delta: Record<string, string> = {};    if (content) delta["content"] = content;    if (reasoning) delta["reasoning_content"] = reasoning;    return `data: ${JSON.stringify({ id: `chatcmpl-${randomUUID().slice(0, 8)}`, choices: [{ index: 0, delta }] })}`;  };   const emitVisibleContent = (    controller: TransformStreamDefaultController<Uint8Array>,    content: string,  ): void => {    if (!content) return;    const controlTokensStripped = stripLeakedDeepSeekControlTokens(content);    visibleContentBuffer += controlTokensStripped;    const cleaned = stripToolXmlDelta(controlTokensStripped);    const chunk = buildFlushChunk({ content: cleaned });    if (chunk) enqueueDataEvent(controller, chunk);  };   const flushThinkCarry = (controller: TransformStreamDefaultController<Uint8Array>): void => {    emitVisibleContent(controller, contentThink.drainPendingContent());    const tail = contentThink.drainCarry();    if (!tail) return;    const carryLooksLikeThink = thinkingTagPrefixIsPartial(tail.trim());    const chunk =      contentThink.inThink() || carryLooksLikeThink        ? buildFlushChunk({ reasoning_content: stripToolXmlDelta(tail) })        : buildFlushChunk({ content: stripToolXmlDelta(tail) });    if (chunk) enqueueDataEvent(controller, chunk);  };   const parseUsage = (data: Record<string, unknown>): void => {    if (usageTracked || !onUsage) return;    const usage = data["usage"] as Record<string, number> | undefined;    if (usage && (usage["prompt_tokens"] || usage["completion_tokens"])) {      onUsage({        prompt_tokens: usage["prompt_tokens"] ?? 0,        completion_tokens: usage["completion_tokens"] ?? 0,        reasoning_tokens:          (usage["reasoning_tokens"] as number | undefined) ??          (usage["completion_tokens_details"] as Record<string, number> | undefined)?.[            "reasoning_tokens"          ] ??          0,        cache_read_tokens:          (usage["prompt_tokens_details"] as Record<string, number> | undefined)?.[            "cached_tokens"          ] ?? 0,        cache_write_tokens: 0,      });      usageTracked = true;    }  };   const trackFirstToken = (): void => {    if (firstTokenTracked) return;    firstTokenTracked = true;    onFirstToken?.();  };   const maybeInjectToolCalls = (controller: TransformStreamDefaultController<Uint8Array>): void => {    if (toolCallsFound || !visibleContentBuffer) return;    const parsed = parseToolCallsFromContent(visibleContentBuffer);    if (parsed.length > 0) {      enqueueDataEvent(controller, buildToolCallChunk(parsed));      toolCallsFound = true;    }  };   const flushEvent = (    controller: TransformStreamDefaultController<Uint8Array>,    lines: string[],  ): void => {    if (lines.length === 0) return;     const dataLines: string[] = [];    const otherLines: string[] = [];    for (const rawLine of lines) {      const trimmedStart = rawLine.trimStart();      if (trimmedStart.startsWith("data:")) {        dataLines.push(trimmedStart.slice("data:".length).trimStart());      } else if (rawLine.length > 0) {        otherLines.push(rawLine);      }    }     if (dataLines.length === 0) {      for (const outLine of lines) {        enqueueLine(controller, outLine);      }      return;    }     const data = dataLines.join("\n").trim();    if (data === "[DONE]") {      flushThinkCarry(controller);      maybeInjectToolCalls(controller);      for (const outLine of otherLines) {        enqueueLine(controller, outLine);      }      enqueueDataEvent(controller, "data: [DONE]");      return;    }     let parsed: Record<string, unknown> | null = null;    try {      parsed = JSON.parse(data) as Record<string, unknown>;    } catch {      parsed = null;    }    if (!parsed) {      for (const outLine of lines) {        enqueueLine(controller, outLine);      }      return;    }     parseUsage(parsed);    const choices = parsed["choices"];    if (Array.isArray(choices)) {      for (const [choiceIndex, choice] of choices.entries()) {        const choiceRecord = choice as Record<string, unknown>;        const hasDelta = choiceRecord["delta"] && typeof choiceRecord["delta"] === "object";        const delta = (hasDelta ? choiceRecord["delta"] : choiceRecord["message"]) as          | Record<string, unknown>          | undefined;        if (!delta) continue;        const toolCalls = delta["tool_calls"];        const hasActiveToolCalls = Array.isArray(toolCalls) && toolCalls.length > 0;        if (hasActiveToolCalls) {          toolCallsFound = true;          trackFirstToken();        }        const rawContent = typeof delta["content"] === "string" ? String(delta["content"]) : "";        const content = normalizeTextDelta(          contentHistory,          `${choiceIndex}:content`,          rawContent,          !hasDelta,        );        const rawReasoning = firstReasoningField(delta);        const reasoningRaw = rawReasoning          ? normalizeTextDelta(              reasoningHistory,              `${choiceIndex}:reasoning`,              rawReasoning,              !hasDelta,            )          : "";        if (rawReasoning) {          emitVisibleContent(controller, contentThink.resolveImplicitPrefixAsContent());        }        if (content || reasoningRaw) trackFirstToken();        let reasoning = "";        let reasoningFromContent = "";        if (content) {          const rewritten = contentThink.rewrite(content, false);          const controlTokensStripped = stripLeakedDeepSeekControlTokens(rewritten.content);          if (controlTokensStripped) {            visibleContentBuffer += controlTokensStripped;          }          const cleanedContent = stripToolXmlDelta(controlTokensStripped);          if (cleanedContent) {            delta["content"] = cleanedContent;          } else if ("content" in delta) {            delete delta["content"];          }          reasoningFromContent = rewritten.reasoningAppend;        } else if (rawContent && "content" in delta) {          delete delta["content"];        }         if (reasoningRaw) {          const rewrittenReasoning = reasoningThink.rewrite(reasoningRaw, true);          reasoning = `${reasoning}${rewrittenReasoning.reasoningAppend}`;        }         if (reasoningFromContent) {          reasoning = `${reasoning}${reasoningFromContent}`;        }         if (reasoning) {          delta["reasoning_content"] = stripToolXmlDelta(reasoning);        } else if (REASONING_FIELDS.some((field) => field in delta)) {          delete delta["reasoning_content"];        }        delete delta["reasoning"];        delete delta["reasoning_text"];      }    }     for (const outLine of otherLines) {      enqueueLine(controller, outLine);    }    enqueueDataEvent(controller, `data: ${JSON.stringify(parsed)}`);  };   const transform = new TransformStream<Uint8Array, Uint8Array>({    transform(chunk, controller): void {      buffer += decoder.decode(chunk, { stream: true });      const lines = buffer.split("\n");      buffer = lines.pop() ?? "";      for (const line of lines) {        const normalized = line.endsWith("\r") ? line.slice(0, -1) : line;        if (normalized === "") {          flushEvent(controller, pendingEventLines);          pendingEventLines = [];          enqueueLine(controller, "");        } else {          pendingEventLines.push(normalized);        }      }    },    flush(controller): void {      buffer += decoder.decode();      if (buffer) {        const trailing = buffer.endsWith("\r") ? buffer.slice(0, -1) : buffer;        if (trailing) pendingEventLines.push(trailing);        buffer = "";      }      if (pendingEventLines.length > 0) {        flushEvent(controller, pendingEventLines);        pendingEventLines = [];      }      flushThinkCarry(controller);      maybeInjectToolCalls(controller);    },  });  return source.pipeThrough(transform);};