{"version":3,"sources":["../../src/memory/tokenMemory.ts"],"names":["TokenMemory","BaseMemory","messages","llm","threshold","syncThreshold","maxTokens","tokensByMessage","WeakMap","handlers","constructor","config","capacityThreshold","estimate","msg","Math","ceil","role","length","text","removalSelector","R","clamp","min","max","TypeError","register","tokensUsed","sum","map","get","tokensCount","isDirty","some","dirty","add","message","index","meta","tokenLimit","Infinity","has","MemoryFatalError","messageToDelete","exists","delete","set","ensureRange","splice","sync","removeFromArray","Promise","all","cache","result","tokenize","addMany","reset","stats","messagesCount","createSnapshot","shallowCopy","filter","_","value","undefined","loadSnapshot","state","Object","assign"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AA2CO,MAAMA,oBAAoBC,mBAAAA,CAAAA;EA3CjC;;;AA4CkBC,EAAAA,QAAAA,GAA0B,EAAA;AAEhCC,EAAAA,GAAAA;AACAC,EAAAA,SAAAA;AACAC,EAAAA,aAAAA;EACAC,SAA2B,GAAA,IAAA;AAC3BC,EAAAA,eAAAA,uBAAsBC,OAAAA,EAAAA;AAChBC,EAAAA,QAAAA;AAEhBC,EAAAA,WAAAA,CAAYC,MAA0B,EAAA;AACpC,IAAK,KAAA,EAAA;AACL,IAAA,IAAA,CAAKR,MAAMQ,MAAOR,CAAAA,GAAAA;AAClB,IAAKG,IAAAA,CAAAA,SAAAA,GAAYK,OAAOL,SAAa,IAAA,IAAA;AACrC,IAAKF,IAAAA,CAAAA,SAAAA,GAAYO,OAAOC,iBAAqB,IAAA,IAAA;AAC7C,IAAKP,IAAAA,CAAAA,aAAAA,GAAgBM,OAAON,aAAiB,IAAA,IAAA;AAC7C,IAAA,IAAA,CAAKI,QAAW,GAAA;AACd,MAAA,GAAGE,MAAQF,EAAAA,QAAAA;AACXI,MAAAA,QAAAA,EACEF,MAAQF,EAAAA,QAAAA,EAAUI,QAAa,KAAA,CAACC,GAAQC,KAAAA,IAAAA,CAAKC,IAAMF,CAAAA,CAAAA,GAAAA,CAAIG,IAAKC,CAAAA,MAAAA,GAASJ,GAAIK,CAAAA,IAAAA,CAAKD,UAAU,CAAA,CAAA,CAAA;AAC1FE,MAAAA,eAAAA,EAAiBT,OAAOF,QAAUW,EAAAA,eAAAA,KAAoB,CAAClB,QAAAA,KAAaA,SAAS,CAAA,CAAA;AAC/E,KAAA;AACA,IAAI,IAAA,CAACmB,aAAEC,KAAM,CAAA;MAAEC,GAAK,EAAA,CAAA;MAAGC,GAAK,EAAA;KAAE,CAAA,CAAG,IAAKpB,CAAAA,SAAS,CAAG,EAAA;AAChD,MAAM,MAAA,IAAIqB,UAAU,sDAAA,CAAA;AACtB;AACF;EAEA;AACE,IAAA,IAAA,CAAKC,QAAQ,EAAA;AACf;AAEA,EAAA,IAAIC,UAAqB,GAAA;AACvB,IAAA,OAAOC,KAAI,CAAA,IAAA,CAAK1B,QAAS2B,CAAAA,GAAAA,CAAI,CAACf,GAAAA,KAAQ,IAAKP,CAAAA,eAAAA,CAAgBuB,GAAIhB,CAAAA,GAAAA,CAAMiB,CAAAA,WAAW,CAAA,CAAA;AAClF;AAEA,EAAA,IAAIC,OAAmB,GAAA;AACrB,IAAO,OAAA,IAAA,CAAK9B,QAAS+B,CAAAA,IAAAA,CAAK,CAACnB,GAAAA,KAAQ,IAAKP,CAAAA,eAAAA,CAAgBuB,GAAIhB,CAAAA,GAAAA,CAAMoB,EAAAA,KAAAA,KAAU,KAAA,CAAA;AAC9E;EAEA,MAAMC,GAAAA,CAAIC,SAAsBC,KAAgB,EAAA;AAC9C,IAAI,IAAA,IAAA,CAAK/B,cAAc,IAAM,EAAA;AAC3B,MAAA,MAAMgC,KAAO,GAAA,MAAM,IAAKnC,CAAAA,GAAAA,CAAImC,IAAI,EAAA;AAChC,MAAA,IAAA,CAAKhC,YAAYS,IAAKC,CAAAA,IAAAA,CAAAA,CAAMsB,MAAKC,UAAcC,IAAAA,QAAAA,IAAY,KAAKpC,SAAS,CAAA;AAC3E;AAEA,IAAMkC,MAAAA,IAAAA,GAAO,IAAK/B,CAAAA,eAAAA,CAAgBkC,GAAIL,CAAAA,OAAAA,IAClC,IAAK7B,CAAAA,eAAAA,CAAgBuB,GAAIM,CAAAA,OAAAA,CACzB,GAAA;MAAEL,WAAa,EAAA,IAAA,CAAKtB,QAASI,CAAAA,QAAAA,CAASuB,OAAAA,CAAAA;MAAUF,KAAO,EAAA;AAAK,KAAA;AAEhE,IAAII,IAAAA,IAAAA,CAAKP,WAAc,GAAA,IAAA,CAAKzB,SAAW,EAAA;AACrC,MAAM,MAAA,IAAIoC,0BACR,CAAsBJ,mBAAAA,EAAAA,IAAAA,CAAKP,WAAW,CAA8C,2CAAA,EAAA,IAAA,CAAKzB,SAAS,CAAU,QAAA,CAAA,CAAA;AAEhH;AAEA,IAAA,OAAO,IAAKqB,CAAAA,UAAAA,GAAa,IAAKrB,CAAAA,SAAAA,GAAYgC,KAAKP,WAAa,EAAA;AAC1D,MAAA,MAAMY,eAAkB,GAAA,IAAA,CAAKlC,QAASW,CAAAA,eAAAA,CAAgB,KAAKlB,QAAQ,CAAA;AACnE,MAAA,MAAM0C,MAAS,GAAA,MAAM,IAAKC,CAAAA,MAAAA,CAAOF,eAAAA,CAAAA;AAEjC,MAAI,IAAA,CAACA,eAAmB,IAAA,CAACC,MAAQ,EAAA;AAC/B,QAAM,MAAA,IAAIF,0BAAiB,4DAAA,CAAA;AAC7B;AACF;AAEA,IAAKnC,IAAAA,CAAAA,eAAAA,CAAgBuC,GAAIV,CAAAA,OAAAA,EAASE,IAAAA,CAAAA;AAElCD,IAAAA,KAAAA,GAAQU,sBAAYV,CAAAA,KAAAA,IAAS,IAAKnC,CAAAA,QAAAA,CAASgB,MAAQ,EAAA;MAAEK,GAAK,EAAA,CAAA;AAAGC,MAAAA,GAAAA,EAAK,KAAKtB,QAASgB,CAAAA;KAAO,CAAA;AACvF,IAAA,IAAA,CAAKhB,QAAS8C,CAAAA,MAAAA,CAAOX,KAAO,EAAA,CAAA,EAAGD,OAAAA,CAAAA;AAE/B,IAAA,IAAI,KAAKJ,OAAW,IAAA,IAAA,CAAKL,aAAa,IAAKrB,CAAAA,SAAAA,IAAa,KAAKD,aAAe,EAAA;AAC1E,MAAA,MAAM,KAAK4C,IAAI,EAAA;AACjB;AACF;AAEA,EAAA,MAAMJ,OAAOT,OAAsB,EAAA;AACjC,IAAOc,OAAAA,yBAAAA,CAAgB,IAAKhD,CAAAA,QAAAA,EAAUkC,OAAAA,CAAAA;AACxC;AAEA,EAAA,MAAMa,IAAO,GAAA;AACX,IAAM/C,MAAAA,QAAAA,GAAW,MAAMiD,OAAQC,CAAAA,GAAAA,CAC7B,KAAKlD,QAAS2B,CAAAA,GAAAA,CAAI,OAAOf,GAAAA,KAAAA;AACvB,MAAA,MAAMuC,KAAQ,GAAA,IAAA,CAAK9C,eAAgBuB,CAAAA,GAAAA,CAAIhB,GAAAA,CAAAA;AACvC,MAAIuC,IAAAA,KAAAA,EAAOnB,UAAU,KAAO,EAAA;AAC1B,QAAA,MAAMoB,MAAS,GAAA,MAAM,IAAKnD,CAAAA,GAAAA,CAAIoD,QAAS,CAAA;AAACzC,UAAAA;AAAI,SAAA,CAAA;AAC5C,QAAKP,IAAAA,CAAAA,eAAAA,CAAgBuC,IAAIhC,GAAK,EAAA;AAAEiB,UAAAA,WAAAA,EAAauB,MAAOvB,CAAAA,WAAAA;UAAaG,KAAO,EAAA;SAAM,CAAA;AAChF;AACA,MAAOpB,OAAAA,GAAAA;AACT,KAAA,CAAA,CAAA;AAGF,IAAA,IAAA,CAAKZ,SAASgB,MAAS,GAAA,CAAA;AACvB,IAAM,MAAA,IAAA,CAAKsC,QAAQtD,QAAAA,CAAAA;AACrB;EAEAuD,KAAQ,GAAA;AACN,IAAW3C,KAAAA,MAAAA,GAAAA,IAAO,KAAKZ,QAAU,EAAA;AAC/B,MAAKK,IAAAA,CAAAA,eAAAA,CAAgBsC,OAAO/B,GAAAA,CAAAA;AAC9B;AACA,IAAA,IAAA,CAAKZ,SAASgB,MAAS,GAAA,CAAA;AACzB;EAEAwC,KAAQ,GAAA;AACN,IAAO,OAAA;AACL/B,MAAAA,UAAAA,EAAY,IAAKA,CAAAA,UAAAA;AACjBrB,MAAAA,SAAAA,EAAW,IAAKA,CAAAA,SAAAA;AAChBqD,MAAAA,aAAAA,EAAe,KAAKzD,QAASgB,CAAAA,MAAAA;AAC7Bc,MAAAA,OAAAA,EAAS,IAAKA,CAAAA;AAChB,KAAA;AACF;EAEA4B,cAAiB,GAAA;AACf,IAAO,OAAA;AACLzD,MAAAA,GAAAA,EAAK,IAAKA,CAAAA,GAAAA;AACVG,MAAAA,SAAAA,EAAW,IAAKA,CAAAA,SAAAA;AAChBF,MAAAA,SAAAA,EAAW,IAAKA,CAAAA,SAAAA;AAChBC,MAAAA,aAAAA,EAAe,IAAKA,CAAAA,aAAAA;MACpBH,QAAU2D,EAAAA,qBAAAA,CAAY,KAAK3D,QAAQ,CAAA;MACnCO,QAAUoD,EAAAA,qBAAAA,CAAY,KAAKpD,QAAQ,CAAA;AACnCF,MAAAA,eAAAA,EAAiB,IAAKL,CAAAA,QAAAA,CACnB2B,GAAI,CAAA,CAACO,OAAY,KAAA;AAACA,QAAAA,OAAAA;QAAS,IAAK7B,CAAAA,eAAAA,CAAgBuB,IAAIM,OAAAA;AAAS,OAAA,CAAA,CAC7D0B,OAAO,CAAC,CAACC,GAAGC,KAAAA,CAAAA,KAAWA,UAAUC,MAAAA;AACtC,KAAA;AACF;AAEAC,EAAAA,YAAAA,CAAa,EAAE3D,eAAAA,EAAiB,GAAG4D,KAAAA,EAAiD,EAAA;AAClFC,IAAOC,MAAAA,CAAAA,MAAAA,CAAO,MAAMF,KAAO,EAAA;MACzB5D,eAAiB,EAAA,IAAIC,QAAQD,eAAAA;KAC/B,CAAA;AACF;AACF","file":"tokenMemory.cjs","sourcesContent":["/**\n * Copyright 2025 IBM Corp.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport { BaseMessage } from \"@/llms/primitives/message.js\";\nimport { BaseMemory, MemoryFatalError } from \"@/memory/base.js\";\nimport { ChatLLM, ChatLLMOutput } from \"@/llms/chat.js\";\nimport * as R from \"remeda\";\nimport { shallowCopy } from \"@/serializer/utils.js\";\nimport { removeFromArray } from \"@/internals/helpers/array.js\";\nimport { sum } from \"remeda\";\nimport { ensureRange } from \"@/internals/helpers/number.js\";\n\nexport interface Handlers {\n estimate: (messages: BaseMessage) => number;\n removalSelector: (messages: BaseMessage[]) => BaseMessage;\n}\n\nexport interface TokenMemoryInput {\n llm: ChatLLM;\n maxTokens?: number;\n syncThreshold?: number;\n capacityThreshold?: number;\n handlers?: Partial;\n}\n\ninterface TokenByMessage {\n tokensCount: number;\n dirty: boolean;\n}\n\nexport class TokenMemory extends BaseMemory {\n public readonly messages: BaseMessage[] = [];\n\n protected llm: ChatLLM;\n protected threshold;\n protected syncThreshold;\n protected maxTokens: number | null = null;\n protected tokensByMessage = new WeakMap();\n public readonly handlers: Handlers;\n\n constructor(config: TokenMemoryInput) {\n super();\n this.llm = config.llm;\n this.maxTokens = config.maxTokens ?? null;\n this.threshold = config.capacityThreshold ?? 0.75;\n this.syncThreshold = config.syncThreshold ?? 0.25;\n this.handlers = {\n ...config?.handlers,\n estimate:\n config?.handlers?.estimate || ((msg) => Math.ceil((msg.role.length + msg.text.length) / 4)),\n removalSelector: config.handlers?.removalSelector || ((messages) => messages[0]),\n };\n if (!R.clamp({ min: 0, max: 1 })(this.threshold)) {\n throw new TypeError('\"capacityThreshold\" must be a number in range (0, 1>');\n }\n }\n\n static {\n this.register();\n }\n\n get tokensUsed(): number {\n return sum(this.messages.map((msg) => this.tokensByMessage.get(msg)!.tokensCount!));\n }\n\n get isDirty(): boolean {\n return this.messages.some((msg) => this.tokensByMessage.get(msg)?.dirty !== false);\n }\n\n async add(message: BaseMessage, index?: number) {\n if (this.maxTokens === null) {\n const meta = await this.llm.meta();\n this.maxTokens = Math.ceil((meta.tokenLimit ?? Infinity) * this.threshold);\n }\n\n const meta = this.tokensByMessage.has(message)\n ? this.tokensByMessage.get(message)!\n : { tokensCount: this.handlers.estimate(message), dirty: true };\n\n if (meta.tokensCount > this.maxTokens) {\n throw new MemoryFatalError(\n `Retrieved message (${meta.tokensCount} tokens) cannot fit inside current memory (${this.maxTokens} tokens)`,\n );\n }\n\n while (this.tokensUsed > this.maxTokens - meta.tokensCount) {\n const messageToDelete = this.handlers.removalSelector(this.messages);\n const exists = await this.delete(messageToDelete);\n\n if (!messageToDelete || !exists) {\n throw new MemoryFatalError('The \"removalSelector\" handler must return a valid message!');\n }\n }\n\n this.tokensByMessage.set(message, meta);\n\n index = ensureRange(index ?? this.messages.length, { min: 0, max: this.messages.length });\n this.messages.splice(index, 0, message);\n\n if (this.isDirty && this.tokensUsed / this.maxTokens >= this.syncThreshold) {\n await this.sync();\n }\n }\n\n async delete(message: BaseMessage) {\n return removeFromArray(this.messages, message);\n }\n\n async sync() {\n const messages = await Promise.all(\n this.messages.map(async (msg) => {\n const cache = this.tokensByMessage.get(msg);\n if (cache?.dirty !== false) {\n const result = await this.llm.tokenize([msg]);\n this.tokensByMessage.set(msg, { tokensCount: result.tokensCount, dirty: false });\n }\n return msg;\n }),\n );\n\n this.messages.length = 0;\n await this.addMany(messages);\n }\n\n reset() {\n for (const msg of this.messages) {\n this.tokensByMessage.delete(msg);\n }\n this.messages.length = 0;\n }\n\n stats() {\n return {\n tokensUsed: this.tokensUsed,\n maxTokens: this.maxTokens,\n messagesCount: this.messages.length,\n isDirty: this.isDirty,\n };\n }\n\n createSnapshot() {\n return {\n llm: this.llm,\n maxTokens: this.maxTokens,\n threshold: this.threshold,\n syncThreshold: this.syncThreshold,\n messages: shallowCopy(this.messages),\n handlers: shallowCopy(this.handlers),\n tokensByMessage: this.messages\n .map((message) => [message, this.tokensByMessage.get(message)])\n .filter(([_, value]) => value !== undefined) as [BaseMessage, number][],\n };\n }\n\n loadSnapshot({ tokensByMessage, ...state }: ReturnType) {\n Object.assign(this, state, {\n tokensByMessage: new WeakMap(tokensByMessage),\n });\n }\n}\n"]}