feat: smooth and combine token output (#936)

* feat: smooth and combine token output

* fix: stop generating button not triggering message updates abort

src/lib/utils/messageUpdates.ts

@@ -0,0 +1,214 @@
+import type { MessageUpdate, TextStreamUpdate } from "$lib/types/MessageUpdate";
+
+type MessageUpdateRequestOptions = {
+	base: string;
+	inputs?: string;
+	messageId?: string;
+	isRetry: boolean;
+	isContinue: boolean;
+	webSearch: boolean;
+	files?: string[];
+};
+export async function fetchMessageUpdates(
+	conversationId: string,
+	opts: MessageUpdateRequestOptions,
+	abortSignal: AbortSignal
+): Promise<AsyncGenerator<MessageUpdate>> {
+	const abortController = new AbortController();
+	abortSignal.addEventListener("abort", () => abortController.abort());
+
+	const response = await fetch(`${opts.base}/conversation/${conversationId}`, {
+		method: "POST",
+		headers: { "Content-Type": "application/json" },
+		body: JSON.stringify({
+			inputs: opts.inputs,
+			id: opts.messageId,
+			is_retry: opts.isRetry,
+			is_continue: opts.isContinue,
+			web_search: opts.webSearch,
+			files: opts.files,
+		}),
+		signal: abortController.signal,
+	});
+
+	if (!response.ok) {
+		const errorMessage = await response
+			.json()
+			.then((obj) => obj.message)
+			.catch(() => `Request failed with status code ${response.status}: ${response.statusText}`);
+		throw Error(errorMessage);
+	}
+	if (!response.body) {
+		throw Error("Body not defined");
+	}
+	return smoothAsyncIterator(
+		streamMessageUpdatesToFullWords(endpointStreamToIterator(response, abortController))
+	);
+}
+
+async function* endpointStreamToIterator(
+	response: Response,
+	abortController: AbortController
+): AsyncGenerator<MessageUpdate> {
+	const reader = response.body?.pipeThrough(new TextDecoderStream()).getReader();
+	if (!reader) throw Error("Response for endpoint had no body");
+
+	// Handle any cases where we must abort
+	reader.closed.then(() => abortController.abort());
+
+	// Handle logic for aborting
+	abortController.signal.addEventListener("abort", () => reader.cancel());
+
+	// ex) If the last response is => {"type": "stream", "token":
+	// It should be => {"type": "stream", "token": "Hello"} = prev_input_chunk + "Hello"}
+	let prevChunk = "";
+	while (!abortController.signal.aborted) {
+		const { done, value } = await reader.read();
+		if (done) {
+			abortController.abort();
+			break;
+		}
+		if (!value) continue;
+
+		const { messageUpdates, remainingText } = parseMessageUpdates(prevChunk + value);
+		prevChunk = remainingText;
+		for (const messageUpdate of messageUpdates) yield messageUpdate;
+	}
+}
+
+function parseMessageUpdates(value: string): {
+	messageUpdates: MessageUpdate[];
+	remainingText: string;
+} {
+	const inputs = value.split("\n");
+	const messageUpdates: MessageUpdate[] = [];
+	for (const input of inputs) {
+		try {
+			messageUpdates.push(JSON.parse(input) as MessageUpdate);
+		} catch (error) {
+			// in case of parsing error, we return what we were able to parse
+			if (error instanceof SyntaxError) {
+				return {
+					messageUpdates,
+					remainingText: inputs.at(-1) ?? "",
+				};
+			}
+		}
+	}
+	return { messageUpdates, remainingText: "" };
+}
+
+/**
+ * Emits all the message updates immediately that aren't "stream" type
+ * Emits a concatenated "stream" type message update once it detects a full word
+ * Example: "what" " don" "'t" => "what" " don't"
+ * Only supports latin languages, ignores others
+ */
+async function* streamMessageUpdatesToFullWords(
+	iterator: AsyncGenerator<MessageUpdate>
+): AsyncGenerator<MessageUpdate> {
+	let bufferedStreamUpdates: TextStreamUpdate[] = [];
+
+	const endAlphanumeric = /[a-zA-Z0-9À-ž'`]+$/;
+	const beginnningAlphanumeric = /^[a-zA-Z0-9À-ž'`]+/;
+
+	for await (const messageUpdate of iterator) {
+		if (messageUpdate.type !== "stream") {
+			yield messageUpdate;
+			continue;
+		}
+		bufferedStreamUpdates.push(messageUpdate);
+
+		let lastIndexEmitted = 0;
+		for (let i = 1; i < bufferedStreamUpdates.length; i++) {
+			const prevEndsAlphanumeric = endAlphanumeric.test(bufferedStreamUpdates[i - 1].token);
+			const currBeginsAlphanumeric = beginnningAlphanumeric.test(bufferedStreamUpdates[i].token);
+			const shouldCombine = prevEndsAlphanumeric && currBeginsAlphanumeric;
+			const combinedTooMany = i - lastIndexEmitted >= 5;
+			if (shouldCombine && !combinedTooMany) continue;
+
+			// Combine tokens together and emit
+			yield {
+				type: "stream",
+				token: bufferedStreamUpdates
+					.slice(lastIndexEmitted, i)
+					.map((_) => _.token)
+					.join(""),
+			};
+			lastIndexEmitted = i;
+		}
+		bufferedStreamUpdates = bufferedStreamUpdates.slice(lastIndexEmitted);
+	}
+	for (const messageUpdate of bufferedStreamUpdates) yield messageUpdate;
+}
+
+/**
+ * Attempts to smooth out the time between values emitted by an async iterator
+ * by waiting for the average time between values to emit the next value
+ */
+async function* smoothAsyncIterator<T>(iterator: AsyncGenerator<T>): AsyncGenerator<T> {
+	const eventTarget = new EventTarget();
+	let done = false;
+	const valuesBuffer: T[] = [];
+	const valueTimesMS: number[] = [];
+
+	const next = async () => {
+		const obj = await iterator.next();
+		if (obj.done) {
+			done = true;
+		} else {
+			valuesBuffer.push(obj.value);
+			valueTimesMS.push(performance.now());
+			next();
+		}
+		eventTarget.dispatchEvent(new Event("next"));
+	};
+	next();
+
+	let timeOfLastEmitMS = performance.now();
+	while (!done || valuesBuffer.length > 0) {
+		// Only consider the last X times between tokens
+		const sampledTimesMS = valueTimesMS.slice(-30);
+
+		// Get the total time spent in abnormal periods
+		const anomalyThresholdMS = 2000;
+		const anomalyDurationMS = sampledTimesMS
+			.map((time, i, times) => time - times[i - 1])
+			.slice(1)
+			.filter((time) => time > anomalyThresholdMS)
+			.reduce((a, b) => a + b, 0);
+
+		// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+		const totalTimeMSBetweenValues = sampledTimesMS.at(-1)! - sampledTimesMS[0];
+		const timeMSBetweenValues = totalTimeMSBetweenValues - anomalyDurationMS;
+
+		const averageTimeMSBetweenValues = Math.min(
+			200,
+			timeMSBetweenValues / (sampledTimesMS.length - 1)
+		);
+		const timeSinceLastEmitMS = performance.now() - timeOfLastEmitMS;
+
+		// Emit after waiting duration or cancel if "next" event is emitted
+		const gotNext = await Promise.race([
+			sleep(Math.max(5, averageTimeMSBetweenValues - timeSinceLastEmitMS)),
+			waitForEvent(eventTarget, "next"),
+		]);
+
+		// Go to next iteration so we can re-calculate when to emit
+		if (gotNext) continue;
+
+		// Nothing in buffer to emit
+		if (valuesBuffer.length === 0) continue;
+
+		// Emit
+		timeOfLastEmitMS = performance.now();
+		// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+		yield valuesBuffer.shift()!;
+	}
+}
+
+const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+const waitForEvent = (eventTarget: EventTarget, eventName: string) =>
+	new Promise<boolean>((resolve) =>
+		eventTarget.addEventListener(eventName, () => resolve(true), { once: true })
+	);

src/routes/conversation/[id]/+page.svelte

@@ -16,6 +16,7 @@
 	import file2base64 from "$lib/utils/file2base64";
 	import { addChildren } from "$lib/utils/tree/addChildren";
 	import { addSibling } from "$lib/utils/tree/addSibling";
+	import { fetchMessageUpdates } from "$lib/utils/messageUpdates";
 	import { createConvTreeStore } from "$lib/stores/convTree";
 	import type { v4 } from "uuid";
 
@@ -181,125 +182,71 @@
 
 			messages = [...messages];
 			const messageToWriteTo = messages.find((message) => message.id === messageToWriteToId);
-
 			if (!messageToWriteTo) {
 				throw new Error("Message to write to not found");
 			}
+
 			// disable websearch if assistant is present
 			const hasAssistant = !!$page.data.assistant;
-
-			const response = await fetch(`${base}/conversation/${$page.params.id}`, {
-				method: "POST",
-				headers: { "Content-Type": "application/json" },
-				body: JSON.stringify({
+			const messageUpdatesAbortController = new AbortController();
+			const messageUpdatesIterator = await fetchMessageUpdates(
+				$page.params.id,
+				{
+					base,
 					inputs: prompt,
-					id: messageId,
-					is_retry: isRetry,
-					is_continue: isContinue,
-					web_search: !hasAssistant && $webSearchParameters.useSearch,
+					messageId,
+					isRetry,
+					isContinue,
+					webSearch: !hasAssistant && $webSearchParameters.useSearch,
 					files: isRetry ? undefined : resizedImages,
-				}),
+				},
+				messageUpdatesAbortController.signal
+			).catch((err) => {
+				error.set(err.message);
 			});
+			if (messageUpdatesIterator === undefined) return;
 
 			files = [];
-			if (!response.body) {
-				throw new Error("Body not defined");
-			}
-
-			if (!response.ok) {
-				error.set((await response.json())?.message);
-				return;
-			}
 
-			// eslint-disable-next-line no-undef
-			const encoder = new TextDecoderStream();
-			const reader = response?.body?.pipeThrough(encoder).getReader();
-			let finalAnswer = "";
 			const messageUpdates: MessageUpdate[] = [];
-
-			// set str queue
-			// ex) if the last response is => {"type": "stream", "token":
-			// It should be => {"type": "stream", "token": "Hello"} = prev_input_chunk + "Hello"}
-			let prev_input_chunk = [""];
-
-			// this is a bit ugly
-			// we read the stream until we get the final answer
-
-			let readerClosed = false;
-
-			reader.closed.then(() => {
-				readerClosed = true;
-			});
-
-			while (finalAnswer === "") {
-				// check for abort
-				if ($isAborted || $error || readerClosed) {
-					reader?.cancel();
+			for await (const update of messageUpdatesIterator) {
+				if ($isAborted) {
+					messageUpdatesAbortController.abort();
+					return;
+				}
+				if (update.type === "finalAnswer") {
+					loading = false;
+					pending = false;
 					break;
 				}
 
-				// if there is something to read
-				await reader?.read().then(async ({ done, value }) => {
-					// we read, if it's done we cancel
-					if (done) {
-						reader.cancel();
-					}
-
-					if (!value) {
-						return;
-					}
-
-					value = prev_input_chunk.pop() + value;
-
-					// if it's not done we parse the value, which contains all messages
-					const inputs = value.split("\n");
-					inputs.forEach(async (el: string) => {
-						try {
-							const update = JSON.parse(el) as MessageUpdate;
-
-							if (update.type !== "stream") {
-								messageUpdates.push(update);
-							}
-
-							if (update.type === "finalAnswer") {
-								finalAnswer = update.text;
-								loading = false;
-								pending = false;
-							} else if (update.type === "stream") {
-								pending = false;
-								messageToWriteTo.content += update.token;
-								messages = [...messages];
-							} else if (update.type === "webSearch") {
-								messageToWriteTo.updates = [...(messageToWriteTo.updates ?? []), update];
-								messages = [...messages];
-							} else if (update.type === "status") {
-								if (update.status === "title" && update.message) {
-									const convInData = data.conversations.find(({ id }) => id === $page.params.id);
-									if (convInData) {
-										convInData.title = update.message;
-
-										$titleUpdate = {
-											title: update.message,
-											convId: $page.params.id,
-										};
-									}
-								} else if (update.status === "error") {
-									$error = update.message ?? "An error has occurred";
-								}
-							} else if (update.type === "error") {
-								error.set(update.message);
-								reader.cancel();
-							}
-						} catch (parseError) {
-							// in case of parsing error we wait for the next message
-
-							if (el === inputs[inputs.length - 1]) {
-								prev_input_chunk.push(el);
-							}
-							return;
+				messageUpdates.push(update);
+
+				if (update.type === "stream") {
+					pending = false;
+					messageToWriteTo.content += update.token;
+					messages = [...messages];
+				} else if (update.type === "webSearch") {
+					messageToWriteTo.updates = [...(messageToWriteTo.updates ?? []), update];
+					messages = [...messages];
+				} else if (update.type === "status") {
+					if (update.status === "title" && update.message) {
+						const convInData = data.conversations.find(({ id }) => id === $page.params.id);
+						if (convInData) {
+							convInData.title = update.message;
+
+							$titleUpdate = {
+								title: update.message,
+								convId: $page.params.id,
+							};
 						}
-					});
-				});
+					} else if (update.status === "error") {
+						$error = update.message ?? "An error has occurred";
+					}
+				} else if (update.type === "error") {
+					error.set(update.message);
+					messageUpdatesAbortController.abort();
+				}
 			}
 
 			messageToWriteTo.updates = messageUpdates;