Make all prompt templates configurable (#400)

* Make all prompt templates configurable

This PR enables all prompts to be configurable using handlebar
templates as described in #382.

For backward compatibility the old hardcoded templates are
reimplemented as the default templates. The old template parameters
such as `preprompt`, `userMessageToken`, `userMessageEndToken`,
`assistantMessageToken`, `assistantMessageEndToken` are now considered
legacy. They still work as they are exposed as variables to the default
template. However, new prompt configurations should not use these. And
it is recommended that the legacy support is eventually removed.

As an example, this is how the default chat prompt template is
implemented:

```
{{preprompt}}
{{#each messages}}
{{#ifUser}}{{@root.userMessageToken}}{{content}}{{@root.userMessageEndToken}}{{/ifUser}}
{{#ifAssistant}}{{@root.assistantMessageToken}}{{content}}{{@root.assistantMessageEndToken}}{{/ifAssistant}}
{{/each}}
{{assistantMessageToken}}
```

In addition, this PR fixes an issue where the `model` configuration was used
to generate the prompts in WebSearch. However, the `defaultModel` was used
to query. This caused issues when the `model` and `defaultModel` uses
different prompt configurations. This has now been changed to always use
the `defaultModel`.

Note, when developing this PR, it has been observed that the WebSearch
prompts can violate typical model assumptions. For example, a query may
be generated as:

```
Assistant: The following context ...
User: content ...
```

```
User: user message 1
User: user message 2
Assistant:
```

Models typically assume the prompts to be User -> Assistant -> User. For
best compatability with existing configurations, this issues was not
fixed. Instead, the old behavior is maintained with the default
templates. This is also the reason why `defaultModel` was chosen as the
WebSearch model instead of `model`. As `defaultModel` may allow the
WebSearch format, while `model` might not.

This behavior, as well as the overall aritecture of chat-ui,
necessitated that the template input maintained the format

```
messages: [{from: 'user' | 'assistant', content: string }]
```

For the template to be able to detect which source a message comes from
a `ifUser` and a `ifAssistant` handlebar block-helper was implemented.
The original proposed format in #382 was:

```
history: [{ user: string, assistant: string }]
```

However, using such format would require significant changes to the
project and would make it impossible to implement the existing
websearch templates.

Finally, there may be minor differences in how truncation is
implemented. As in some cases, truncation is now applied to the entire
prompt, rather than part of the prompt.

Fixes: #382

* Add Sagemaker support (#401)


* work on sagemaker support

* fix sagemaker integration

* remove unnecessary deps

* fix default endpoint

* remove unneeded deps, fixed types

* Use conditional validation for endpoints

This was needed because the discriminated union couldn't handle the legacy case where `host` is undefined.

* add note in readme about aws sagemaker

* lint

* -summery +summary

---------

Co-authored-by: Nathan Sarrazin <[email protected]>

README.md

@@ -120,9 +120,8 @@ MODELS=`[
     "websiteUrl": "https://open-assistant.io",
     "userMessageToken": "<|prompter|>", # This does not need to be a token, can be any string
     "assistantMessageToken": "<|assistant|>", # This does not need to be a token, can be any string
-    "messageEndToken": "<|endoftext|>", # This does not need to be a token, can be any string
-    # "userMessageEndToken": "", # Applies only to user messages, messageEndToken has no effect if specified. Can be any string.
-    # "assistantMessageEndToken": "", # Applies only to assistant messages, messageEndToken has no effect if specified. Can be any string.
+    "userMessageEndToken": "<|endoftext|>", # Applies only to user messages. Can be any string.
+    "assistantMessageEndToken": "<|endoftext|>", # Applies only to assistant messages. Can be any string.
     "preprompt": "Below are a series of dialogues between various people and an AI assistant. The AI tries to be helpful, polite, honest, sophisticated, emotionally aware, and humble-but-knowledgeable. The assistant is happy to help with almost anything, and will do its best to understand exactly what is needed. It also tries to avoid giving false or misleading information, and it caveats when it isn't entirely sure about the right answer. That said, the assistant is practical and really does its best, and doesn't let caution get too much in the way of being useful.\n-----\n",
     "promptExamples": [
       {
@@ -152,7 +151,72 @@ MODELS=`[
 
 You can change things like the parameters, or customize the preprompt to better suit your needs. You can also add more models by adding more objects to the array, with different preprompts for example.
 
-### Running your own models using a custom endpoint
+#### Custom prompt templates:
+
+By default the prompt is constructed using `userMessageToken`, `assistantMessageToken`, `userMessageEndToken`, `assistantMessageEndToken`, `preprompt` parameters and a series of default templates.
+
+However, these templates can be modified by setting the `chatPromptTemplate`, `webSearchSummaryPromptTemplate`, and `webSearchQueryPromptTemplate` parameters. Note that if WebSearch is not enabled, only `chatPromptTemplate` needs to be set. The template language is https://handlebarsjs.com. The templates have access to the model's prompt parameters (`preprompt`, etc.). However, if the templates are specified it is recommended to inline the prompt parameters, as using the references (`{{preprompt}}`) is deprecated.
+
+For example:
+
+```
+<System>You are an AI, called ChatAI.</System>
+{{#each messages}}
+  {{#ifUser}}<User>{{content}}</User>{{/ifUser}}
+  {{#ifAssistant}}<Assistant>{{content}}</Assistant>{{/ifAssistant}}
+{{/each}}
+<Assistant>
+```
+
+**chatPromptTemplate**
+
+When quering the model for a chat response, the `chatPromptTemplate` template is used. `messages` is an array of chat messages, it has the format `[{ content: string }, ...]`. To idenify if a message is a user message or an assistant message the `ifUser` and `ifAssistant` block helpers can be used.
+
+The following is the default `chatPromptTemplate`, although newlines and indentiation have been added for readability.
+
+```
+{{preprompt}}
+{{#each messages}}
+  {{#ifUser}}{{@root.userMessageToken}}{{content}}{{@root.userMessageEndToken}}{{/ifUser}}
+  {{#ifAssistant}}{{@root.assistantMessageToken}}{{content}}{{@root.assistantMessageEndToken}}{{/ifAssistant}}
+{{/each}}
+{{assistantMessageToken}}
+```
+
+**webSearchQueryPromptTemplate**
+
+When performing a websearch, the search query is constructed using the `webSearchQueryPromptTemplate` template. It is recommended that that the prompt instructs the chat model to only return a few keywords.
+
+The following is the default `webSearchQueryPromptTemplate`. Note that not all models supports consecutive user-messages which this template uses.
+
+```
+{{userMessageToken}}
+  The following messages were written by a user, trying to answer a question.
+{{userMessageEndToken}}
+{{#each messages}}
+  {{#ifUser}}{{@root.userMessageToken}}{{content}}{{@root.userMessageEndToken}}{{/ifUser}}
+{{/each}}
+{{userMessageToken}}
+  What plain-text english sentence would you input into Google to answer the last question? Answer with a short (10 words max) simple sentence.
+{{userMessageEndToken}}
+{{assistantMessageToken}}Query:
+```
+
+**webSearchSummaryPromptTemplate**
+
+The search-engine response (`answer`) is summarized using the following prompt template. However, when `HF_ACCESS_TOKEN` is provided, a dedicated summary model is used instead. Additionally, the model's `query` response to `webSearchQueryPromptTemplate` is also available to this template.
+
+The following is the default `webSearchSummaryPromptTemplate`. Note that not all models supports consecutive user-messages which this template uses.
+
+```
+{{userMessageToken}}{{answer}}{{userMessageEndToken}}
+{{userMessageToken}}
+  The text above should be summarized to best answer the query: {{query}}.
+{{userMessageEndToken}}
+{{assistantMessageToken}}Summary:
+```
+
+#### Running your own models using a custom endpoint
 
 If you want to, instead of hitting models on the Hugging Face Inference API, you can run your own models locally.
 

package-lock.json

@@ -14,6 +14,7 @@
 				"aws4fetch": "^1.0.17",
 				"date-fns": "^2.29.3",
 				"dotenv": "^16.0.3",
+				"handlebars": "^4.7.8",
 				"highlight.js": "^11.7.0",
 				"jsdom": "^22.0.0",
 				"marked": "^4.3.0",
@@ -2671,6 +2672,26 @@
 			"integrity": "sha512-bzh50DW9kTPM00T8y4o8vQg89Di9oLJVLW/KaOGIXJWP/iqCN6WKYkbNOF04vFLJhwcpYUh9ydh/+5vpOqV4YQ==",
 			"dev": true
 		},
+		"node_modules/handlebars": {
+			"version": "4.7.8",
+			"resolved": "https://registry.npmjs.org/handlebars/-/handlebars-4.7.8.tgz",
+			"integrity": "sha512-vafaFqs8MZkRrSX7sFVUdo3ap/eNiLnb4IakshzvP56X5Nr1iGKAIqdX6tMlm6HcNRIkr6AxO5jFEoJzzpT8aQ==",
+			"dependencies": {
+				"minimist": "^1.2.5",
+				"neo-async": "^2.6.2",
+				"source-map": "^0.6.1",
+				"wordwrap": "^1.0.0"
+			},
+			"bin": {
+				"handlebars": "bin/handlebars"
+			},
+			"engines": {
+				"node": ">=0.4.7"
+			},
+			"optionalDependencies": {
+				"uglify-js": "^3.1.4"
+			}
+		},
 		"node_modules/has": {
 			"version": "1.0.3",
 			"resolved": "https://registry.npmjs.org/has/-/has-1.0.3.tgz",
@@ -3317,7 +3338,6 @@
 			"version": "1.2.8",
 			"resolved": "https://registry.npmjs.org/minimist/-/minimist-1.2.8.tgz",
 			"integrity": "sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA==",
-			"dev": true,
 			"funding": {
 				"url": "https://github.com/sponsors/ljharb"
 			}
@@ -3449,6 +3469,11 @@
 			"integrity": "sha512-Tj+HTDSJJKaZnfiuw+iaF9skdPpTo2GtEly5JHnWV/hfv2Qj/9RKsGISQtLh2ox3l5EAGw487hnBee0sIJ6v2g==",
 			"dev": true
 		},
+		"node_modules/neo-async": {
+			"version": "2.6.2",
+			"resolved": "https://registry.npmjs.org/neo-async/-/neo-async-2.6.2.tgz",
+			"integrity": "sha512-Yd3UES5mWCSqR+qNT93S3UoYUkqAZ9lLg8a7g9rimsWmYGK8cVToA4/sF3RrshdyV3sAGMXVUmpMYOw+dLpOuw=="
+		},
 		"node_modules/node-int64": {
 			"version": "0.4.0",
 			"resolved": "https://registry.npmjs.org/node-int64/-/node-int64-0.4.0.tgz",
@@ -4438,6 +4463,14 @@
 				"sorcery": "bin/sorcery"
 			}
 		},
+		"node_modules/source-map": {
+			"version": "0.6.1",
+			"resolved": "https://registry.npmjs.org/source-map/-/source-map-0.6.1.tgz",
+			"integrity": "sha512-UjgapumWlbMhkBgzT7Ykc5YXUT46F0iKu8SGXq0bcwP5dz/h0Plj6enJqjz1Zbq2l5WaqYnrVbwWOWMyF3F47g==",
+			"engines": {
+				"node": ">=0.10.0"
+			}
+		},
 		"node_modules/source-map-js": {
 			"version": "1.0.2",
 			"resolved": "https://registry.npmjs.org/source-map-js/-/source-map-js-1.0.2.tgz",
@@ -5042,6 +5075,18 @@
 			"integrity": "sha512-TrY6DsjTQQgyS3E3dBaOXf0TpPD8u9FVrVYmKVegJuFw51n/YB9XPt+U6ydzFG5ZIN7+DIjPbNmXoBj9esYhgQ==",
 			"dev": true
 		},
+		"node_modules/uglify-js": {
+			"version": "3.17.4",
+			"resolved": "https://registry.npmjs.org/uglify-js/-/uglify-js-3.17.4.tgz",
+			"integrity": "sha512-T9q82TJI9e/C1TAxYvfb16xO120tMVFZrGA3f9/P4424DNu6ypK103y0GPFVa17yotwSyZW5iYXgjYHkGrJW/g==",
+			"optional": true,
+			"bin": {
+				"uglifyjs": "bin/uglifyjs"
+			},
+			"engines": {
+				"node": ">=0.8.0"
+			}
+		},
 		"node_modules/undici": {
 			"version": "5.22.0",
 			"resolved": "https://registry.npmjs.org/undici/-/undici-5.22.0.tgz",
@@ -5441,6 +5486,11 @@
 				"node": ">=0.10.0"
 			}
 		},
+		"node_modules/wordwrap": {
+			"version": "1.0.0",
+			"resolved": "https://registry.npmjs.org/wordwrap/-/wordwrap-1.0.0.tgz",
+			"integrity": "sha512-gvVzJFlPycKc5dZN4yPkP8w7Dc37BtP1yczEneOb4uq34pXZcvrtRTmWV8W+Ume+XCxKgbjM+nevkyFPMybd4Q=="
+		},
 		"node_modules/wrappy": {
 			"version": "1.0.2",
 			"resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz",

package.json

@@ -46,6 +46,7 @@
 		"aws4fetch": "^1.0.17",
 		"date-fns": "^2.29.3",
 		"dotenv": "^16.0.3",
+		"handlebars": "^4.7.8",
 		"highlight.js": "^11.7.0",
 		"jsdom": "^22.0.0",
 		"marked": "^4.3.0",

src/lib/buildPrompt.ts

@@ -13,24 +13,6 @@ export async function buildPrompt(
 	model: BackendModel,
 	webSearchId?: string
 ): Promise<string> {
-	const userEndToken = model.userMessageEndToken ?? model.messageEndToken;
-	const assistantEndToken = model.assistantMessageEndToken ?? model.messageEndToken;
-
-	const prompt =
-		messages
-			.map((m) =>
-				m.from === "user"
-					? model.userMessageToken +
-					  m.content +
-					  (m.content.endsWith(userEndToken) ? "" : userEndToken)
-					: model.assistantMessageToken +
-					  m.content +
-					  (m.content.endsWith(assistantEndToken) ? "" : assistantEndToken)
-			)
-			.join("") + model.assistantMessageToken;
-
-	let webPrompt = "";
-
 	if (webSearchId) {
 		const webSearch = await collections.webSearches.findOne({
 			_id: new ObjectId(webSearchId),
@@ -39,20 +21,22 @@ export async function buildPrompt(
 		if (!webSearch) throw new Error("Web search not found");
 
 		if (webSearch.summary) {
-			webPrompt =
-				model.assistantMessageToken +
-				`The following context was found while searching the internet: ${webSearch.summary}` +
-				model.assistantMessageEndToken;
+			messages = [
+				{
+					from: "assistant",
+					content: `The following context was found while searching the internet: ${webSearch.summary}`,
+				},
+				...messages,
+			];
 		}
 	}
-	const finalPrompt =
-		model.preprompt +
-		webPrompt +
-		prompt
+
+	return (
+		model
+			.chatPromptRender({ messages })
+			// Not super precise, but it's truncated in the model's backend anyway
 			.split(" ")
 			.slice(-(model.parameters?.truncate ?? 0))
-			.join(" ");
-
-	// Not super precise, but it's truncated in the model's backend anyway
-	return finalPrompt;
+			.join(" ")
+	);
 }

src/lib/server/models.ts

@@ -1,4 +1,10 @@
 import { HF_ACCESS_TOKEN, MODELS, OLD_MODELS } from "$env/static/private";
+import type {
+	ChatTemplateInput,
+	WebSearchQueryTemplateInput,
+	WebSearchSummaryTemplateInput,
+} from "$lib/types/Template";
+import { compileTemplate } from "$lib/utils/template";
 import { z } from "zod";
 
 const sagemakerEndpoint = z.object({
@@ -46,13 +52,46 @@ const modelsRaw = z
 			modelUrl: z.string().url().optional(),
 			datasetName: z.string().min(1).optional(),
 			datasetUrl: z.string().url().optional(),
-			userMessageToken: z.string(),
+			userMessageToken: z.string().default(""),
 			userMessageEndToken: z.string().default(""),
-			assistantMessageToken: z.string(),
+			assistantMessageToken: z.string().default(""),
 			assistantMessageEndToken: z.string().default(""),
 			messageEndToken: z.string().default(""),
 			preprompt: z.string().default(""),
 			prepromptUrl: z.string().url().optional(),
+			chatPromptTemplate: z
+				.string()
+				.default(
+					"{{preprompt}}" +
+						"{{#each messages}}" +
+						"{{#ifUser}}{{@root.userMessageToken}}{{content}}{{@root.userMessageEndToken}}{{/ifUser}}" +
+						"{{#ifAssistant}}{{@root.assistantMessageToken}}{{content}}{{@root.assistantMessageEndToken}}{{/ifAssistant}}" +
+						"{{/each}}" +
+						"{{assistantMessageToken}}"
+				),
+			webSearchSummaryPromptTemplate: z
+				.string()
+				.default(
+					"{{userMessageToken}}{{answer}}{{userMessageEndToken}}" +
+						"{{userMessageToken}}" +
+						"The text above should be summarized to best answer the query: {{query}}." +
+						"{{userMessageEndToken}}" +
+						"{{assistantMessageToken}}Summary: "
+				),
+			webSearchQueryPromptTemplate: z
+				.string()
+				.default(
+					"{{userMessageToken}}" +
+						"The following messages were written by a user, trying to answer a question." +
+						"{{userMessageEndToken}}" +
+						"{{#each messages}}" +
+						"{{#ifUser}}{{@root.userMessageToken}}{{content}}{{@root.userMessageEndToken}}{{/ifUser}}" +
+						"{{/each}}" +
+						"{{userMessageToken}}" +
+						"What plain-text english sentence would you input into Google to answer the last question? Answer with a short (10 words max) simple sentence." +
+						"{{userMessageEndToken}}" +
+						"{{assistantMessageToken}}Query: "
+				),
 			promptExamples: z
 				.array(
 					z.object({
@@ -80,6 +119,15 @@ export const models = await Promise.all(
 		...m,
 		userMessageEndToken: m?.userMessageEndToken || m?.messageEndToken,
 		assistantMessageEndToken: m?.assistantMessageEndToken || m?.messageEndToken,
+		chatPromptRender: compileTemplate<ChatTemplateInput>(m.chatPromptTemplate, m),
+		webSearchSummaryPromptRender: compileTemplate<WebSearchSummaryTemplateInput>(
+			m.webSearchSummaryPromptTemplate,
+			m
+		),
+		webSearchQueryPromptRender: compileTemplate<WebSearchQueryTemplateInput>(
+			m.webSearchQueryPromptTemplate,
+			m
+		),
 		id: m.id || m.name,
 		displayName: m.displayName || m.name,
 		preprompt: m.prepromptUrl ? await fetch(m.prepromptUrl).then((r) => r.text()) : m.preprompt,

src/lib/server/websearch/generateQuery.ts

@@ -1,21 +1,9 @@
 import type { Message } from "$lib/types/Message";
 import { generateFromDefaultEndpoint } from "../generateFromDefaultEndpoint";
-import type { BackendModel } from "../models";
-
-export async function generateQuery(messages: Message[], model: BackendModel) {
-	const promptSearchQuery =
-		model.userMessageToken +
-		"The following messages were written by a user, trying to answer a question." +
-		model.userMessageEndToken +
-		messages
-			.filter((message) => message.from === "user")
-			.map((message) => model.userMessageToken + message.content + model.userMessageEndToken) +
-		model.userMessageToken +
-		"What plain-text english sentence would you input into Google to answer the last question? Answer with a short (10 words max) simple sentence." +
-		model.userMessageEndToken +
-		model.assistantMessageToken +
-		"Query: ";
+import { defaultModel } from "../models";
 
+export async function generateQuery(messages: Message[]) {
+	const promptSearchQuery = defaultModel.webSearchQueryPromptRender({ messages });
 	const searchQuery = await generateFromDefaultEndpoint(promptSearchQuery).then((query) => {
 		const arr = query.split(/\r?\n/);
 		return arr[0].length > 0 ? arr[0] : arr[1];

src/lib/server/websearch/summarizeWeb.ts

@@ -1,7 +1,8 @@
 import { HF_ACCESS_TOKEN } from "$env/static/private";
 import { HfInference } from "@huggingface/inference";
-import { generateFromDefaultEndpoint } from "../generateFromDefaultEndpoint";
+import { defaultModel } from "$lib/server/models";
 import type { BackendModel } from "../models";
+import { generateFromDefaultEndpoint } from "../generateFromDefaultEndpoint";
 
 export async function summarizeWeb(content: string, query: string, model: BackendModel) {
 	// if HF_ACCESS_TOKEN is set, we use a HF dedicated endpoint for summarization
@@ -23,19 +24,13 @@ export async function summarizeWeb(content: string, query: string, model: Backen
 	}
 
 	// else we use the LLM to generate a summary
-	const summaryPrompt =
-		model.userMessageToken +
-		content
+	const summaryPrompt = defaultModel.webSearchSummaryPromptRender({
+		answer: content
 			.split(" ")
 			.slice(0, model.parameters?.truncate ?? 0)
-			.join(" ") +
-		model.userMessageEndToken +
-		model.userMessageToken +
-		`The text above should be summarized to best answer the query: ${query}.` +
-		model.userMessageEndToken +
-		model.assistantMessageToken +
-		"Summary: ";
-
+			.join(" "),
+		query: query,
+	});
 	const summary = await generateFromDefaultEndpoint(summaryPrompt).then((txt: string) =>
 		txt.trim()
 	);

src/lib/types/Template.ts

@@ -0,0 +1,22 @@
+import type { Message } from "./Message";
+
+export type LegacyParamatersTemplateInput = {
+	preprompt: string;
+	userMessageToken: string;
+	userMessageEndToken: string;
+	assistantMessageToken: string;
+	assistantMessageEndToken: string;
+};
+
+export type ChatTemplateInput = {
+	messages: Pick<Message, "from" | "content">[];
+};
+
+export type WebSearchSummaryTemplateInput = {
+	answer: string;
+	query: string;
+};
+
+export type WebSearchQueryTemplateInput = {
+	messages: Pick<Message, "from" | "content">[];
+};

src/lib/utils/template.ts

@@ -0,0 +1,28 @@
+import type { Message } from "$lib/types/Message";
+import type { LegacyParamatersTemplateInput } from "$lib/types/Template";
+import Handlebars from "handlebars";
+
+Handlebars.registerHelper("ifUser", function (this: Pick<Message, "from" | "content">, options) {
+	if (this.from == "user") return options.fn(this);
+});
+
+Handlebars.registerHelper(
+	"ifAssistant",
+	function (this: Pick<Message, "from" | "content">, options) {
+		if (this.from == "assistant") return options.fn(this);
+	}
+);
+
+export function compileTemplate<T>(input: string, model: LegacyParamatersTemplateInput) {
+	const template = Handlebars.compile<T & LegacyParamatersTemplateInput>(input, {
+		knownHelpers: { ifUser: true, ifAssistant: true },
+		knownHelpersOnly: true,
+		noEscape: true,
+		strict: true,
+		preventIndent: true,
+	});
+
+	return function render(inputs: T, options?: RuntimeOptions) {
+		return template({ ...model, ...inputs }, options);
+	};
+}

src/routes/conversation/[id]/web-search/+server.ts

@@ -69,7 +69,7 @@ export async function GET({ params, locals, url }) {
 
 			try {
 				appendUpdate("Generating search query");
-				webSearch.searchQuery = await generateQuery(messages, model);
+				webSearch.searchQuery = await generateQuery(messages);
 
 				appendUpdate("Searching Google", [webSearch.searchQuery]);
 				const results = await searchWeb(webSearch.searchQuery);