From 0d27bf1311aea976fadaa954d0a8fd91dc0d7ba1 Mon Sep 17 00:00:00 2001
From: Peli de Halleux <pelikhan@users.noreply.github.com>
Date: Mon, 30 Sep 2024 08:57:32 -0700
Subject: [PATCH] Adjust flexPromptNode logic to handle token calculation and
 flex checks more accurately (#737)

* Adjust flexPromptNode logic to handle token calculation and flex checks more accurately

* Fix typos and adjust flex logic in promptdom.ts and genai.mts files
---
 packages/core/src/promptdom.ts          | 14 ++++++++++++--
 packages/sample/genaisrc/flex.genai.mts |  6 +++---
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/packages/core/src/promptdom.ts b/packages/core/src/promptdom.ts
index dc6fab64a9..a6c0182142 100644
--- a/packages/core/src/promptdom.ts
+++ b/packages/core/src/promptdom.ts
@@ -30,6 +30,7 @@ import { resolveTokenEncoder } from "./encoders"
 import { expandFiles } from "./fs"
 import { interpolateVariables } from "./mustache"
 import { createDiff } from "./diff"
+import { total } from "@tidyjs/tidy"
 
 // Definition of the PromptNode interface which is an essential part of the code structure.
 export interface PromptNode extends ContextExpansionOptions {
@@ -787,7 +788,7 @@ async function flexPromptNode(
         0
     )
 
-    if (totalTokens < flexTokens) {
+    if (totalTokens <= flexTokens) {
         // No need to flex
         return
     }
@@ -799,8 +800,17 @@ async function flexPromptNode(
             (a.priority ?? PRIORITY_DEFAULT) - (b.priority ?? PRIORITY_DEFAULT)
     )
     const flexNodes = nodes.filter((n) => n.flex !== undefined)
-    const totalFlex = flexNodes.reduce((total, node) => total + node.flex, 0)
+    const totalFlexTokens = flexNodes.reduce(
+        (total, node) => total + (node.tokens ?? 0),
+        0
+    )
 
+    // checking flexNodes sizes
+    if (totalFlexTokens <= flexTokens) {
+        return
+    }
+
+    const totalFlex = flexNodes.reduce((total, node) => total + node.flex, 0)
     const totalReserve = 0
     const totalRemaining = Math.max(0, flexTokens - totalReserve)
     for (const node of flexNodes) {
diff --git a/packages/sample/genaisrc/flex.genai.mts b/packages/sample/genaisrc/flex.genai.mts
index 167c1eb2bb..92c6d00d99 100644
--- a/packages/sample/genaisrc/flex.genai.mts
+++ b/packages/sample/genaisrc/flex.genai.mts
@@ -21,9 +21,9 @@ script({
 def("FILE", env.files, { flex: 1 })
 
 // will be trimmed
-$`What is Markdown?
- Markdown is a lightweight markup language that you can use to add formatting elements to plaintext text documents. Created by John Gruber in 2004, Markdown is now one of the world’s most popular markup languages. 
+$`(ignore What is Markdown?
+ Markdown is a lightweight markup language that you can use to add formatting elements to plaintext text documents. Created by John Gruber in 2004, Markdown is now one of the world’s most popular markup languages.)
 PRINT ABRACADABRA!`.flex(2)
 
-$`This one is flexed.
+$`(ignore This one is flexed.)
 PRINT MONKEY!`.flex(1)