fix(continue): fix continue feature (#1459)

* fix(continue): more reliable end token stripping

before we would try to strip end tokens once in order from the stop array, now we strip all end tokens from the array until we can't anymore

* fix(tools): remove generation prompt in continue mode

Author: nsarrazin

src/lib/buildPrompt.ts

@@ -28,19 +28,28 @@ export async function buildPrompt({
 			preprompt,
 			tools,
 			toolResults,
+			continueMessage,
 		})
 		// Not super precise, but it's truncated in the model's backend anyway
 		.split(" ")
 		.slice(-(model.parameters?.truncate ?? 0))
 		.join(" ");
 
 	if (continueMessage && model.parameters?.stop) {
-		prompt = model.parameters.stop.reduce((acc: string, curr: string) => {
-			if (acc.endsWith(curr)) {
-				return acc.slice(0, acc.length - curr.length);
+		let trimmedPrompt = prompt.trimEnd();
+		let hasRemovedStop = true;
+		while (hasRemovedStop) {
+			hasRemovedStop = false;
+			for (const stopToken of model.parameters.stop) {
+				if (trimmedPrompt.endsWith(stopToken)) {
+					trimmedPrompt = trimmedPrompt.slice(0, -stopToken.length);
+					hasRemovedStop = true;
+					break;
+				}
 			}
-			return acc;
-		}, prompt.trimEnd());
+			trimmedPrompt = trimmedPrompt.trimEnd();
+		}
+		prompt = trimmedPrompt;
 	}
 
 	return prompt;

src/lib/server/models.ts

@@ -94,7 +94,13 @@ async function getChatPromptRender(
 		process.exit();
 	}
 
-	const renderTemplate = ({ messages, preprompt, tools, toolResults }: ChatTemplateInput) => {
+	const renderTemplate = ({
+		messages,
+		preprompt,
+		tools,
+		toolResults,
+		continueMessage,
+	}: ChatTemplateInput) => {
 		let formattedMessages: { role: string; content: string }[] = messages.map((message) => ({
 			content: message.content,
 			role: message.from,
@@ -222,10 +228,8 @@ async function getChatPromptRender(
 
 		const output = tokenizer.apply_chat_template(formattedMessages, {
 			tokenize: false,
-			add_generation_prompt: true,
+			add_generation_prompt: !continueMessage,
 			chat_template: chatTemplate,
-			// eslint-disable-next-line @typescript-eslint/ban-ts-comment
-			// @ts-ignore
 			tools: mappedTools,
 			documents,
 		});

src/lib/types/Template.ts

@@ -6,4 +6,5 @@ export type ChatTemplateInput = {
 	preprompt?: string;
 	tools?: Tool[];
 	toolResults?: ToolResult[];
+	continueMessage?: boolean;
 };