From db817fcf29cdeede00244c8eb8bbbc726fdf4b39 Mon Sep 17 00:00:00 2001
From: Jun Siang Cheah <git@jscheah.me>
Date: Sat, 13 Apr 2024 18:26:50 +0100
Subject: [PATCH] feat: add {{prompt:start:length}} and {{prompt:end:length}}
 to title gen

---
 src/lib/apis/ollama/index.ts |  3 ++-
 src/lib/apis/openai/index.ts |  3 ++-
 src/lib/utils/index.ts       | 25 +++++++++++++++++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/lib/apis/ollama/index.ts b/src/lib/apis/ollama/index.ts
index 7c4e809e5b..237c16425a 100644
--- a/src/lib/apis/ollama/index.ts
+++ b/src/lib/apis/ollama/index.ts
@@ -1,4 +1,5 @@
 import { OLLAMA_API_BASE_URL } from '$lib/constants';
+import { templatePrompt } from '$lib/utils';
 
 export const getOllamaUrls = async (token: string = '') => {
 	let error = null;
@@ -144,7 +145,7 @@ export const generateTitle = async (
 ) => {
 	let error = null;
 
-	template = template.replace(/{{prompt}}/g, prompt);
+	template = templatePrompt(template, prompt);
 
 	console.log(template);
 
diff --git a/src/lib/apis/openai/index.ts b/src/lib/apis/openai/index.ts
index 7d0d9415e0..1502b4a8ca 100644
--- a/src/lib/apis/openai/index.ts
+++ b/src/lib/apis/openai/index.ts
@@ -1,4 +1,5 @@
 import { OPENAI_API_BASE_URL } from '$lib/constants';
+import { templatePrompt } from '$lib/utils';
 
 export const getOpenAIUrls = async (token: string = '') => {
 	let error = null;
@@ -273,7 +274,7 @@ export const generateTitle = async (
 ) => {
 	let error = null;
 
-	template = template.replace(/{{prompt}}/g, prompt);
+	template = templatePrompt(template, prompt);
 
 	console.log(template);
 
diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts
index e9a4e229b2..8a3cdf163d 100644
--- a/src/lib/utils/index.ts
+++ b/src/lib/utils/index.ts
@@ -467,3 +467,28 @@ export const blobToFile = (blob, fileName) => {
 	const file = new File([blob], fileName, { type: blob.type });
 	return file;
 };
+
+// templatePrompt replaces any occurrences of the following in the template with the prompt
+// {{prompt}} will be replaced with the prompt
+// {{prompt:start:<length>}} will be replaced with the first <length> characters of the prompt
+// {{prompt:end:<length>}} will be replaced with the last <length> characters of the prompt
+// Character length is used as we don't have the ability to tokenize the prompt
+export const templatePrompt = (template, prompt) => {
+	template = template.replace(/{{prompt}}/g, prompt);
+
+	// Replace {{prompt:start:<length>}} with the first <length> characters of the prompt
+	const startMatch = template.match(/{{prompt:start:(\d+)}}/);
+	if (startMatch) {
+		const length = parseInt(startMatch[1]);
+		template = template.replace(startMatch[0], prompt.substring(0, length));
+	}
+
+	// Replace {{prompt:end:<length>}} with the last <length> characters of the prompt
+	const endMatch = template.match(/{{prompt:end:(\d+)}}/);
+	if (endMatch) {
+		const length = parseInt(endMatch[1]);
+		template = template.replace(endMatch[0], prompt.substring(prompt.length - length));
+	}
+
+	return template;
+};