✨ Save messages in backend (#31)

.eslintrc.cjs

@@ -12,6 +12,9 @@ module.exports = {
 		sourceType: 'module',
 		ecmaVersion: 2020
 	},
+	rules: {
+		'no-shadow': ['error']
+	},
 	env: {
 		browser: true,
 		es2017: true,

src/lib/buildPrompt.ts

@@ -0,0 +1,25 @@
+import {
+	PUBLIC_ASSISTANT_MESSAGE_TOKEN,
+	PUBLIC_SEP_TOKEN,
+	PUBLIC_USER_MESSAGE_TOKEN
+} from '$env/static/public';
+import type { Message } from './types/Message';
+
+/**
+ * Convert [{user: "assistant", content: "hi"}, {user: "user", content: "hello"}] to:
+ *
+ * <|assistant|>hi<|endoftext|><|prompter|>hello<|endoftext|><|assistant|>
+ */
+export function buildPrompt(messages: Message[]): string {
+	return (
+		messages
+			.map(
+				(m) =>
+					(m.from === 'user'
+						? PUBLIC_USER_MESSAGE_TOKEN + m.content
+						: PUBLIC_ASSISTANT_MESSAGE_TOKEN + m.content) +
+					(m.content.endsWith(PUBLIC_SEP_TOKEN) ? '' : PUBLIC_SEP_TOKEN)
+			)
+			.join('') + PUBLIC_ASSISTANT_MESSAGE_TOKEN
+	);
+}

src/lib/utils/streamToAsyncIterable.ts

@@ -0,0 +1,15 @@
+// https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Statements/for-await...of#iterating_over_async_generators
+export async function* streamToAsyncIterable(
+	stream: ReadableStream<Uint8Array>
+): AsyncIterableIterator<Uint8Array> {
+	const reader = stream.getReader();
+	try {
+		while (true) {
+			const { done, value } = await reader.read();
+			if (done) return;
+			yield value;
+		}
+	} finally {
+		reader.releaseLock();
+	}
+}

src/lib/utils/sum.ts

@@ -0,0 +1,3 @@
+export function sum(nums: number[]): number {
+	return nums.reduce((a, b) => a + b, 0);
+}

src/routes/+page.svelte

@@ -1,5 +1,5 @@
 <script lang="ts">
-	import { goto, invalidate, invalidateAll } from '$app/navigation';
+	import { goto } from '$app/navigation';
 	import ChatWindow from '$lib/components/chat/ChatWindow.svelte';
 	import { pendingMessage } from '$lib/stores/pendingMessage';
 

src/routes/api/conversation/+server.ts

@@ -1,19 +0,0 @@
-import { HF_TOKEN } from '$env/static/private';
-import { PUBLIC_MODEL_ENDPOINT } from '$env/static/public';
-
-export async function POST({ request, fetch }) {
-	const resp = await fetch(PUBLIC_MODEL_ENDPOINT, {
-		headers: {
-			'Content-Type': request.headers.get('Content-Type') ?? 'application/json',
-			Authorization: `Basic ${HF_TOKEN}`
-		},
-		method: 'POST',
-		body: await request.text()
-	});
-
-	return new Response(resp.body, {
-		headers: Object.fromEntries(resp.headers.entries()),
-		status: resp.status,
-		statusText: resp.statusText
-	});
-}

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

@@ -4,23 +4,14 @@
 	import { onMount } from 'svelte';
 	import type { PageData } from './$types';
 	import { page } from '$app/stores';
-	import {
-		PUBLIC_ASSISTANT_MESSAGE_TOKEN,
-		PUBLIC_SEP_TOKEN,
-		PUBLIC_USER_MESSAGE_TOKEN
-	} from '$env/static/public';
 	import { HfInference } from '@huggingface/inference';
 
 	export let data: PageData;
 
 	$: messages = data.messages;
 
-	const userToken = PUBLIC_USER_MESSAGE_TOKEN;
-	const assistantToken = PUBLIC_ASSISTANT_MESSAGE_TOKEN;
-	const sepToken = PUBLIC_SEP_TOKEN;
-
 	const hf = new HfInference();
-	const model = hf.endpoint(`${$page.url.origin}/api/conversation`);
+	const model = hf.endpoint($page.url.href);
 
 	let loading = false;
 
@@ -76,16 +67,7 @@
 
 			messages = [...messages, { from: 'user', content: message }];
 
-			const inputs =
-				messages
-					.map(
-						(m) =>
-							(m.from === 'user' ? userToken + m.content : assistantToken + m.content) +
-							(m.content.endsWith(sepToken) ? '' : sepToken)
-					)
-					.join('') + assistantToken;
-
-			await getTextGenerationStream(inputs);
+			await getTextGenerationStream(message);
 		} finally {
 			loading = false;
 		}

src/routes/conversation/[id]/+server.ts

@@ -0,0 +1,110 @@
+import { HF_TOKEN } from '$env/static/private';
+import { PUBLIC_MODEL_ENDPOINT } from '$env/static/public';
+import { buildPrompt } from '$lib/buildPrompt.js';
+import { collections } from '$lib/server/database.js';
+import type { Message } from '$lib/types/Message.js';
+import { streamToAsyncIterable } from '$lib/utils/streamToAsyncIterable';
+import { sum } from '$lib/utils/sum';
+import { error } from '@sveltejs/kit';
+import { ObjectId } from 'mongodb';
+
+export async function POST({ request, fetch, locals, params }) {
+	// todo: add validation on params.id
+	const convId = new ObjectId(params.id);
+
+	const conv = await collections.conversations.findOne({
+		_id: convId,
+		sessionId: locals.sessionId
+	});
+
+	if (!conv) {
+		throw error(404, 'Conversation not found');
+	}
+
+	// Todo: validate prompt with zod? or aktype
+	const json = await request.json();
+
+	const messages = [...conv.messages, { from: 'user', content: json.inputs }] satisfies Message[];
+
+	json.inputs = buildPrompt(messages);
+
+	const resp = await fetch(PUBLIC_MODEL_ENDPOINT, {
+		headers: {
+			'Content-Type': request.headers.get('Content-Type') ?? 'application/json',
+			Authorization: `Basic ${HF_TOKEN}`
+		},
+		method: 'POST',
+		body: JSON.stringify(json)
+	});
+
+	const [stream1, stream2] = resp.body!.tee();
+
+	async function saveMessage() {
+		const generated_text = await parseGeneratedText(stream2);
+
+		messages.push({ from: 'assistant', content: generated_text });
+
+		console.log('updating conversation', convId, messages);
+
+		await collections.conversations.updateOne(
+			{
+				_id: convId
+			},
+			{
+				$set: {
+					messages,
+					updatedAt: new Date()
+				}
+			}
+		);
+	}
+
+	saveMessage().catch(console.error);
+
+	// Todo: maybe we should wait for the message to be saved before ending the response - in case of errors
+	return new Response(stream1, {
+		headers: Object.fromEntries(resp.headers.entries()),
+		status: resp.status,
+		statusText: resp.statusText
+	});
+}
+
+async function parseGeneratedText(stream: ReadableStream): Promise<string> {
+	const inputs: Uint8Array[] = [];
+	for await (const input of streamToAsyncIterable(stream)) {
+		inputs.push(input);
+	}
+
+	// Merge inputs into a single Uint8Array
+	const completeInput = new Uint8Array(sum(inputs.map((input) => input.length)));
+	let offset = 0;
+	for (const input of inputs) {
+		completeInput.set(input, offset);
+		offset += input.length;
+	}
+
+	// Get last line starting with "data:" and parse it as JSON to get the generated text
+	const message = new TextDecoder().decode(completeInput);
+
+	let lastIndex = message.lastIndexOf('\ndata:');
+	if (lastIndex === -1) {
+		lastIndex = message.indexOf('data');
+	}
+
+	if (lastIndex === -1) {
+		console.error('Could not parse in last message');
+	}
+
+	let lastMessage = message.slice(lastIndex).trim().slice('data:'.length);
+	if (lastMessage.includes('\n')) {
+		lastMessage = lastMessage.slice(0, lastMessage.indexOf('\n'));
+	}
+
+	const res = JSON.parse(lastMessage).generated_text;
+
+	if (typeof res !== 'string') {
+		throw new Error('Could not parse generated text');
+	}
+
+	return res;
+}