From 124af70dcf4ebca88e6c0454672de08f231d4cbf Mon Sep 17 00:00:00 2001
From: Stefan Merettig <Stefan@Merettig.io>
Date: Mon, 7 Oct 2024 14:54:19 +0200
Subject: [PATCH] feature: Allow the user to configure fully custom LLM prompts

---
 packages/shared/config.ts  | 16 +++++++----
 packages/shared/prompts.ts | 59 +++++++++++++++++++++++++++++---------
 2 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/packages/shared/config.ts b/packages/shared/config.ts
index 288becab0..69e704367 100644
--- a/packages/shared/config.ts
+++ b/packages/shared/config.ts
@@ -24,6 +24,8 @@ const allEnv = z.object({
   INFERENCE_JOB_TIMEOUT_SEC: z.coerce.number().default(30),
   INFERENCE_TEXT_MODEL: z.string().default("gpt-4o-mini"),
   INFERENCE_IMAGE_MODEL: z.string().default("gpt-4o-mini"),
+  IMAGE_TAG_PROMPT: z.string().optional(),
+  TEXT_TAG_PROMPT: z.string().optional(),
   CRAWLER_HEADLESS_BROWSER: stringBool("true"),
   BROWSER_WEB_URL: z.string().url().optional(),
   BROWSER_WEBSOCKET_URL: z.string().url().optional(),
@@ -74,6 +76,8 @@ const serverConfigSchema = allEnv.transform((val) => {
       textModel: val.INFERENCE_TEXT_MODEL,
       imageModel: val.INFERENCE_IMAGE_MODEL,
       inferredTagLang: val.INFERENCE_LANG,
+      imageTagPrompt: val.IMAGE_TAG_PROMPT,
+      textTagPrompt: val.TEXT_TAG_PROMPT,
     },
     crawler: {
       numWorkers: val.CRAWLER_NUM_WORKERS,
@@ -90,16 +94,16 @@ const serverConfigSchema = allEnv.transform((val) => {
     },
     meilisearch: val.MEILI_ADDR
       ? {
-          address: val.MEILI_ADDR,
-          key: val.MEILI_MASTER_KEY,
-        }
+        address: val.MEILI_ADDR,
+        key: val.MEILI_MASTER_KEY,
+      }
       : undefined,
     logLevel: val.LOG_LEVEL,
     demoMode: val.DEMO_MODE
       ? {
-          email: val.DEMO_MODE_EMAIL,
-          password: val.DEMO_MODE_PASSWORD,
-        }
+        email: val.DEMO_MODE_EMAIL,
+        password: val.DEMO_MODE_PASSWORD,
+      }
       : undefined,
     dataDir: val.DATA_DIR,
     maxAssetSizeMb: val.MAX_ASSET_SIZE_MB,
diff --git a/packages/shared/prompts.ts b/packages/shared/prompts.ts
index cf6d48b66..be8c1634b 100644
--- a/packages/shared/prompts.ts
+++ b/packages/shared/prompts.ts
@@ -1,33 +1,64 @@
-export function buildImagePrompt(lang: string, customPrompts: string[]) {
-  return `
+import serverConfig from "./config";
+
+const DEFAULT_IMAGE_PROMPT = `
 You are a bot in a read-it-later app and your responsibility is to help with automatic tagging.
 Please analyze the attached image and suggest relevant tags that describe its key themes, topics, and main ideas. The rules are:
 - Aim for a variety of tags, including broad categories, specific keywords, and potential sub-genres.
-- The tags language must be in ${lang}.
+- The tags language must be in {{lang}}.
 - If the tag is not generic enough, don't include it.
 - Aim for 10-15 tags.
 - If there are no good tags, don't emit any.
-${customPrompts && customPrompts.map((p) => `- ${p}`).join("\n")}
+{{customPrompts}}
 You must respond in valid JSON with the key "tags" and the value is list of tags. Don't wrap the response in a markdown code.`;
-}
 
-export function buildTextPrompt(
-  lang: string,
-  customPrompts: string[],
-  content: string,
-) {
-  return `
+const DEFAULT_TEXT_PROMPT = `
 You are a bot in a read-it-later app and your responsibility is to help with automatic tagging.
 Please analyze the text between the sentences "CONTENT START HERE" and "CONTENT END HERE" and suggest relevant tags that describe its key themes, topics, and main ideas. The rules are:
 - Aim for a variety of tags, including broad categories, specific keywords, and potential sub-genres.
-- The tags language must be in ${lang}.
+- The tags language must be in {{lang}}.
 - If it's a famous website you may also include a tag for the website. If the tag is not generic enough, don't include it.
 - The content can include text for cookie consent and privacy policy, ignore those while tagging.
 - Aim for 3-5 tags.
 - If there are no good tags, leave the array empty.
-${customPrompts && customPrompts.map((p) => `- ${p}`).join("\n")}
+{{customPrompts}}
+
 CONTENT START HERE
-${content}
+{{content}}
 CONTENT END HERE
 You must respond in JSON with the key "tags" and the value is an array of string tags.`;
+
+function renderTemplate(template: string, variables: Map<string, string>): string {
+  let result = template;
+
+  for (const [name, value] of variables.entries()) {
+    const placeholder = "\\{\\{" + name + "\\}\\}"; // Don't parse the {{}} as regex!
+    result = result.replace(new RegExp(placeholder, 'g'), value);
+  }
+
+  return result;
+}
+
+export function buildImagePrompt(lang: string, customPrompts: string[]) {
+  const promptTemplate = serverConfig.inference.imageTagPrompt || DEFAULT_IMAGE_PROMPT;
+  const customPromptsRendered = [...customPrompts].map((p) => `- ${p}`).join("\n");
+
+  return renderTemplate(promptTemplate, new Map([
+    ['lang', lang],
+    ['customPrompts', customPromptsRendered],
+  ]));
+}
+
+export function buildTextPrompt(
+  lang: string,
+  customPrompts: string[],
+  content: string,
+) {
+  const promptTemplate = serverConfig.inference.textTagPrompt || DEFAULT_TEXT_PROMPT;
+  const customPromptsRendered = [...customPrompts].map((p) => `- ${p}`).join("\n");
+
+  return renderTemplate(promptTemplate, new Map([
+    ['lang', lang],
+    ['customPrompts', customPromptsRendered],
+    ['content', content],
+  ]));
 }