Add ellipse on token limit exceed and flex layout options in PromptTe…

…mplateString (#694) * Add ellipse on token limit exceed and create new script file in sample package * Add flex layout options and priority setting in PromptTemplateString for better prompt rendering control * Refactor variable names and add flexing logic in prompt rendering functions * Refactor token truncation and streamlining flex options for prompt rendering logic * Add maxTokens option to flexPromptNode and adjust token estimation with overhead constant * basic flex * add test
microsoft · Sep 6, 2024 · a85828c · a85828c
1 parent 615b38e
commit a85828c
Show file tree

Hide file tree

Showing 25 changed files with 551 additions and 56 deletions.
diff --git a/docs/genaisrc/genaiscript.d.ts b/docs/genaisrc/genaiscript.d.ts
diff --git a/genaisrc/genaiscript.d.ts b/genaisrc/genaiscript.d.ts
diff --git a/packages/core/src/chat.ts b/packages/core/src/chat.ts
@@ -470,16 +470,17 @@ async function processChatMessage(
                 const node = ctx.node
                 checkCancelled(cancellationToken)
                 // expand template
-                const { errors, prompt } = await renderPromptNode(
+                const { errors, userPrompt } = await renderPromptNode(
                     options.model,
                     node,
                     {
+                        flexTokens: options.flexTokens,
                         trace,
                     }
                 )
-                if (prompt?.trim().length) {
-                    trace.detailsFenced(`💬 message`, prompt, "markdown")
-                    messages.push({ role: "user", content: prompt })
+                if (userPrompt?.trim().length) {
+                    trace.detailsFenced(`💬 message`, userPrompt, "markdown")
+                    messages.push({ role: "user", content: userPrompt })
                     needsNewTurn = true
                 } else trace.item("no message")
                 if (errors?.length) {

diff --git a/packages/core/src/constants.ts b/packages/core/src/constants.ts
@@ -227,3 +227,5 @@ export const CONSOLE_COLOR_WARNING = 95
 export const CONSOLE_COLOR_ERROR = 91
 
 export const PLAYWRIGHT_DEFAULT_BROWSER = "chromium"
+export const MAX_TOKENS_ELLIPSE = "..."
+export const ESTIMATE_TOKEN_OVERHEAD = 2
diff --git a/packages/core/src/expander.ts b/packages/core/src/expander.ts
@@ -73,7 +73,7 @@ export async function callExpander(
         const node = ctx.node
         if (provider !== MODEL_PROVIDER_AICI) {
             const {
-                prompt,
+                userPrompt,
                 assistantPrompt,
                 images: imgs,
                 errors,
@@ -83,8 +83,11 @@ export async function callExpander(
                 outputProcessors: ops,
                 chatParticipants: cps,
                 fileOutputs: fos,
-            } = await renderPromptNode(model, node, { trace })
-            text = prompt
+            } = await renderPromptNode(model, node, {
+                flexTokens: options.flexTokens,
+                trace,
+            })
+            text = userPrompt
             assistantText = assistantPrompt
             images = imgs
             schemas = schs
@@ -175,7 +178,7 @@ export async function expandTemplate(
     const systems = resolveSystems(prj, template)
     const systemTemplates = systems.map((s) => prj.getTemplate(s))
     // update options
-    options.lineNumbers =
+    const lineNumbers =
         options.lineNumbers ??
         template.lineNumbers ??
         systemTemplates.some((s) => s?.lineNumbers)
@@ -186,7 +189,7 @@ export async function expandTemplate(
         host.defaultModelOptions.temperature
     const topP =
         options.topP ?? normalizeFloat(env.vars["top_p"]) ?? template.topP
-    const max_tokens =
+    const maxTokens =
         options.maxTokens ??
         normalizeInt(env.vars["maxTokens"]) ??
         normalizeInt(env.vars["max_tokens"]) ??
@@ -197,6 +200,11 @@ export async function expandTemplate(
         normalizeInt(env.vars["max_tool_calls"]) ??
         template.maxToolCalls ??
         MAX_TOOL_CALLS
+    const flexTokens =
+        options.flexTokens ??
+        normalizeInt(env.vars["flexTokens"]) ??
+        normalizeInt(env.vars["flex_tokens"]) ??
+        template.flexTokens
     let seed = options.seed ?? normalizeInt(env.vars["seed"]) ?? template.seed
     if (seed !== undefined) seed = seed >> 0
 
@@ -207,7 +215,16 @@ export async function expandTemplate(
     trace.startDetails("🧬 prompt")
     trace.detailsFenced("📓 script source", template.jsSource, "js")
 
-    const prompt = await callExpander(prj, template, env, trace, options)
+    const prompt = await callExpander(prj, template, env, trace, {
+        ...options,
+        maxTokens,
+        maxToolCalls,
+        flexTokens,
+        seed,
+        topP,
+        temperature,
+        lineNumbers,
+    })
 
     const images = prompt.images
     const schemas = prompt.schemas
@@ -339,7 +356,7 @@ ${schemaTs}
         model,
         temperature,
         topP,
-        max_tokens,
+        maxTokens,
         maxToolCalls,
         seed,
         responseType,

diff --git a/packages/core/src/genaisrc/genaiscript.d.ts b/packages/core/src/genaisrc/genaiscript.d.ts
diff --git a/packages/core/src/promptcontext.ts b/packages/core/src/promptcontext.ts
@@ -275,6 +275,7 @@ export async function createPromptContext(
                         messages: msgs,
                         chatParticipants: cps,
                     } = await renderPromptNode(genOptions.model, node, {
+                        flexTokens: genOptions.flexTokens,
                         trace,
                     })