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

src/lib/buildPrompt.ts

@@ -28,6 +28,7 @@ export async function buildPrompt({
 			preprompt,
 			tools,
 			toolResults,
+			continueMessage,
 		})
 		// Not super precise, but it's truncated in the model's backend anyway
 		.split(" ")
@@ -35,12 +36,20 @@ export async function buildPrompt({
 		.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;
 };