From 1b09682685f54f29957163be9b9f9fc2de3b49cc Mon Sep 17 00:00:00 2001
From: MohamedBassem <me@mbassem.com>
Date: Sat, 12 Oct 2024 17:25:01 +0000
Subject: feature: Allow customizing the inference's context length

---
 packages/shared/config.ts  |  2 ++
 packages/shared/prompts.ts | 23 +++++++++++++++++++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

(limited to 'packages/shared')

diff --git a/packages/shared/config.ts b/packages/shared/config.ts
index 44b7e26d..325d9ffa 100644
--- a/packages/shared/config.ts
+++ b/packages/shared/config.ts
@@ -24,6 +24,7 @@ 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"),
+  INFERENCE_CONTEXT_LENGTH: z.coerce.number().default(2048),
   CRAWLER_HEADLESS_BROWSER: stringBool("true"),
   BROWSER_WEB_URL: z.string().url().optional(),
   BROWSER_WEBSOCKET_URL: z.string().url().optional(),
@@ -74,6 +75,7 @@ const serverConfigSchema = allEnv.transform((val) => {
       textModel: val.INFERENCE_TEXT_MODEL,
       imageModel: val.INFERENCE_IMAGE_MODEL,
       inferredTagLang: val.INFERENCE_LANG,
+      contextLength: val.INFERENCE_CONTEXT_LENGTH,
     },
     crawler: {
       numWorkers: val.CRAWLER_NUM_WORKERS,
diff --git a/packages/shared/prompts.ts b/packages/shared/prompts.ts
index cf6d48b6..91bfba3f 100644
--- a/packages/shared/prompts.ts
+++ b/packages/shared/prompts.ts
@@ -1,3 +1,17 @@
+// TODO: Use a proper tokenizer
+function calculateNumTokens(text: string) {
+  return text.split(" ").length;
+}
+
+function truncateContent(content: string, length: number) {
+  let words = content.split(" ");
+  if (words.length > length) {
+    words = words.slice(0, length);
+    content = words.join(" ");
+  }
+  return content;
+}
+
 export function buildImagePrompt(lang: string, customPrompts: string[]) {
   return `
 You are a bot in a read-it-later app and your responsibility is to help with automatic tagging.
@@ -15,8 +29,9 @@ export function buildTextPrompt(
   lang: string,
   customPrompts: string[],
   content: string,
+  contextLength: number,
 ) {
-  return `
+  const constructPrompt = (c: string) => `
 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.
@@ -27,7 +42,11 @@ Please analyze the text between the sentences "CONTENT START HERE" and "CONTENT
 - If there are no good tags, leave the array empty.
 ${customPrompts && customPrompts.map((p) => `- ${p}`).join("\n")}
 CONTENT START HERE
-${content}
+${c}
 CONTENT END HERE
 You must respond in JSON with the key "tags" and the value is an array of string tags.`;
+
+  const promptSize = calculateNumTokens(constructPrompt(""));
+  const truncatedContent = truncateContent(content, contextLength - promptSize);
+  return constructPrompt(truncatedContent);
 }
-- 
cgit v1.2.3-70-g09d2