{"version":3,"sources":["../../src/memory/tokenMemory.ts"],"names":["simpleEstimate","msg","Math","ceil","text","length","simpleTokenize","msgs","sum","map","TokenMemory","BaseMemory","messages","threshold","syncThreshold","maxTokens","tokensByMessage","WeakMap","handlers","config","capacityThreshold","estimate","role","tokenize","removalSelector","R","clamp","min","max","TypeError","register","tokensUsed","get","tokensCount","isDirty","some","dirty","add","message","index","meta","has","MemoryFatalError","messageToDelete","exists","delete","set","ensureRange","splice","sync","removeFromArray","Promise","all","cache","addMany","reset","stats","messagesCount","createSnapshot","shallowCopy","filter","_","value","undefined","loadSnapshot","state","Object","assign"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AAmBA,MAAMA,cAAAA,2BAAwCC,GAAAA,KAAiBC,IAAAA,CAAKC,KAAKF,GAAAA,CAAIG,IAAAA,CAAKC,MAAAA,GAAS,CAAA,CAAA,EAA9C,gBAAA,CAAA;AAC7C,MAAMC,cAAAA,iCAA8CC,IAAAA,KAClDC,KAAAA,CAAIC,MAAIF,IAAAA,EAAMP,cAAAA,CAAAA,CAAAA,EAD6B,gBAAA,CAAA;AAetC,MAAMU,oBAAoBC,mBAAAA,CAAAA;EAnCjC;;;AAoCkBC,EAAAA,QAAAA,GAAsB,EAAA;AAE5BC,EAAAA,SAAAA;AACAC,EAAAA,aAAAA;EACAC,SAAAA,GAA2B,IAAA;AAC3BC,EAAAA,eAAAA,uBAAsBC,OAAAA,EAAAA;AAChBC,EAAAA,QAAAA;EAEhB,WAAA,CAAYC,MAAAA,GAA2B,EAAC,EAAG;AACzC,IAAA,KAAA,EAAK;AACL,IAAA,IAAA,CAAKJ,SAAAA,GAAYI,OAAOJ,SAAAA,IAAa,IAAA;AACrC,IAAA,IAAA,CAAKF,SAAAA,GAAYM,OAAOC,iBAAAA,IAAqB,IAAA;AAC7C,IAAA,IAAA,CAAKN,aAAAA,GAAgBK,OAAOL,aAAAA,IAAiB,IAAA;AAC7C,IAAA,IAAA,CAAKI,QAAAA,GAAW;AACd,MAAA,GAAGC,MAAAA,EAAQD,QAAAA;AACXG,MAAAA,QAAAA,EACEF,MAAAA,EAAQD,QAAAA,EAAUG,QAAAA,KAAa,CAACpB,GAAAA,KAAQC,IAAAA,CAAKC,IAAAA,CAAAA,CAAMF,GAAAA,CAAIqB,IAAAA,CAAKjB,MAAAA,GAASJ,GAAAA,CAAIG,IAAAA,CAAKC,UAAU,CAAA,CAAA,CAAA;MAC1FkB,QAAAA,EAAUJ,MAAAA,EAAQD,UAAUK,QAAAA,IAAYjB,cAAAA;AACxCkB,MAAAA,eAAAA,EAAiBL,OAAOD,QAAAA,EAAUM,eAAAA,KAAoB,CAACZ,QAAAA,KAAaA,SAAS,CAAA,CAAA;AAC/E,KAAA;AACA,IAAA,IAAI,CAACa,aAAEC,KAAAA,CAAM;MAAEC,GAAAA,EAAK,CAAA;MAAGC,GAAAA,EAAK;KAAE,CAAA,CAAG,IAAA,CAAKf,SAAS,CAAA,EAAG;AAChD,MAAA,MAAM,IAAIgB,UAAU,sDAAA,CAAA;AACtB,IAAA;AACF,EAAA;EAEA;AACE,IAAA,IAAA,CAAKC,QAAAA,EAAQ;AACf;AAEA,EAAA,IAAIC,UAAAA,GAAqB;AACvB,IAAA,OAAOvB,KAAAA,CAAI,IAAA,CAAKI,QAAAA,CAASH,GAAAA,CAAI,CAACR,GAAAA,KAAQ,IAAA,CAAKe,eAAAA,CAAgBgB,GAAAA,CAAI/B,GAAAA,CAAAA,CAAMgC,WAAW,CAAA,CAAA;AAClF,EAAA;AAEA,EAAA,IAAIC,OAAAA,GAAmB;AACrB,IAAA,OAAO,IAAA,CAAKtB,QAAAA,CAASuB,IAAAA,CAAK,CAAClC,GAAAA,KAAQ,IAAA,CAAKe,eAAAA,CAAgBgB,GAAAA,CAAI/B,GAAAA,CAAAA,EAAMmC,KAAAA,KAAU,KAAA,CAAA;AAC9E,EAAA;EAEA,MAAMC,GAAAA,CAAIC,SAAkBC,KAAAA,EAAgB;AAC1C,IAAA,IAAI,IAAA,CAAKxB,cAAc,IAAA,EAAM;AAE3B,MAAA,IAAA,CAAKA,SAAAA,GAAY,KAAA;AACnB,IAAA;AAEA,IAAA,MAAMyB,IAAAA,GAAO,IAAA,CAAKxB,eAAAA,CAAgByB,GAAAA,CAAIH,OAAAA,IAClC,IAAA,CAAKtB,eAAAA,CAAgBgB,GAAAA,CAAIM,OAAAA,CAAAA,GACzB;MAAEL,WAAAA,EAAa,IAAA,CAAKf,QAAAA,CAASG,QAAAA,CAASiB,OAAAA,CAAAA;MAAUF,KAAAA,EAAO;AAAK,KAAA;AAEhE,IAAA,IAAII,IAAAA,CAAKP,WAAAA,GAAc,IAAA,CAAKlB,SAAAA,EAAW;AACrC,MAAA,MAAM,IAAI2B,0BACR,CAAA,mBAAA,EAAsBF,IAAAA,CAAKP,WAAW,CAAA,2CAAA,EAA8C,IAAA,CAAKlB,SAAS,CAAA,QAAA,CAAU,CAAA;AAEhH,IAAA;AAEA,IAAA,OAAO,IAAA,CAAKgB,UAAAA,GAAa,IAAA,CAAKhB,SAAAA,GAAYyB,KAAKP,WAAAA,EAAa;AAC1D,MAAA,MAAMU,eAAAA,GAAkB,IAAA,CAAKzB,QAAAA,CAASM,eAAAA,CAAgB,KAAKZ,QAAQ,CAAA;AACnE,MAAA,MAAMgC,MAAAA,GAAS,MAAM,IAAA,CAAKC,MAAAA,CAAOF,eAAAA,CAAAA;AAEjC,MAAA,IAAI,CAACA,eAAAA,IAAmB,CAACC,MAAAA,EAAQ;AAC/B,QAAA,MAAM,IAAIF,0BAAiB,4DAAA,CAAA;AAC7B,MAAA;AACF,IAAA;AAEA,IAAA,IAAA,CAAK1B,eAAAA,CAAgB8B,GAAAA,CAAIR,OAAAA,EAASE,IAAAA,CAAAA;AAElCD,IAAAA,KAAAA,GAAQQ,sBAAAA,CAAYR,KAAAA,IAAS,IAAA,CAAK3B,QAAAA,CAASP,MAAAA,EAAQ;MAAEsB,GAAAA,EAAK,CAAA;AAAGC,MAAAA,GAAAA,EAAK,KAAKhB,QAAAA,CAASP;KAAO,CAAA;AACvF,IAAA,IAAA,CAAKO,QAAAA,CAASoC,MAAAA,CAAOT,KAAAA,EAAO,CAAA,EAAGD,OAAAA,CAAAA;AAE/B,IAAA,IAAI,KAAKJ,OAAAA,IAAW,IAAA,CAAKH,aAAa,IAAA,CAAKhB,SAAAA,IAAa,KAAKD,aAAAA,EAAe;AAC1E,MAAA,MAAM,KAAKmC,IAAAA,EAAI;AACjB,IAAA;AACF,EAAA;AAEA,EAAA,MAAMJ,OAAOP,OAAAA,EAAkB;AAC7B,IAAA,OAAOY,yBAAAA,CAAgB,IAAA,CAAKtC,QAAAA,EAAU0B,OAAAA,CAAAA;AACxC,EAAA;AAEA,EAAA,MAAMW,IAAAA,GAAO;AACX,IAAA,MAAMrC,QAAAA,GAAW,MAAMuC,OAAAA,CAAQC,GAAAA,CAC7B,KAAKxC,QAAAA,CAASH,GAAAA,CAAI,OAAOR,GAAAA,KAAAA;AACvB,MAAA,MAAMoD,KAAAA,GAAQ,IAAA,CAAKrC,eAAAA,CAAgBgB,GAAAA,CAAI/B,GAAAA,CAAAA;AACvC,MAAA,IAAIoD,KAAAA,EAAOjB,UAAU,KAAA,EAAO;AAC1B,QAAA,MAAMH,WAAAA,GAAc,MAAM,IAAA,CAAKf,QAAAA,CAASK,QAAAA,CAAS;AAACtB,UAAAA;AAAI,SAAA,CAAA;AACtD,QAAA,IAAA,CAAKe,eAAAA,CAAgB8B,IAAI7C,GAAAA,EAAK;AAAEgC,UAAAA,WAAAA;UAAaG,KAAAA,EAAO;SAAM,CAAA;AAC5D,MAAA;AACA,MAAA,OAAOnC,GAAAA;AACT,IAAA,CAAA,CAAA,CAAA;AAGF,IAAA,IAAA,CAAKW,SAASP,MAAAA,GAAS,CAAA;AACvB,IAAA,MAAM,IAAA,CAAKiD,QAAQ1C,QAAAA,CAAAA;AACrB,EAAA;EAEA2C,KAAAA,GAAQ;AACN,IAAA,KAAA,MAAWtD,GAAAA,IAAO,KAAKW,QAAAA,EAAU;AAC/B,MAAA,IAAA,CAAKI,eAAAA,CAAgB6B,OAAO5C,GAAAA,CAAAA;AAC9B,IAAA;AACA,IAAA,IAAA,CAAKW,SAASP,MAAAA,GAAS,CAAA;AACzB,EAAA;EAEAmD,KAAAA,GAAQ;AACN,IAAA,OAAO;AACLzB,MAAAA,UAAAA,EAAY,IAAA,CAAKA,UAAAA;AACjBhB,MAAAA,SAAAA,EAAW,IAAA,CAAKA,SAAAA;AAChB0C,MAAAA,aAAAA,EAAe,KAAK7C,QAAAA,CAASP,MAAAA;AAC7B6B,MAAAA,OAAAA,EAAS,IAAA,CAAKA;AAChB,KAAA;AACF,EAAA;EAEAwB,cAAAA,GAAiB;AACf,IAAA,OAAO;AACL7C,MAAAA,SAAAA,EAAW,IAAA,CAAKA,SAAAA;AAChBC,MAAAA,aAAAA,EAAe,IAAA,CAAKA,aAAAA;MACpBF,QAAAA,EAAU+C,qBAAAA,CAAY,KAAK/C,QAAQ,CAAA;MACnCM,QAAAA,EAAUyC,qBAAAA,CAAY,KAAKzC,QAAQ,CAAA;AACnCH,MAAAA,SAAAA,EAAW,IAAA,CAAKA,SAAAA;AAChBC,MAAAA,eAAAA,EAAiB,IAAA,CAAKJ,QAAAA,CACnBH,GAAAA,CAAI,CAAC6B,OAAAA,KAAY;AAACA,QAAAA,OAAAA;QAAS,IAAA,CAAKtB,eAAAA,CAAgBgB,IAAIM,OAAAA;AAAS,OAAA,CAAA,CAC7DsB,OAAO,CAAC,CAACC,GAAGC,KAAAA,CAAAA,KAAWA,UAAUC,MAAAA;AACtC,KAAA;AACF,EAAA;AAEAC,EAAAA,YAAAA,CAAa,EAAEhD,eAAAA,EAAiB,GAAGiD,KAAAA,EAAAA,EAAiD;AAClFC,IAAAA,MAAAA,CAAOC,MAAAA,CAAO,MAAMF,KAAAA,EAAO;MACzBjD,eAAAA,EAAiB,IAAIC,QAAQD,eAAAA;KAC/B,CAAA;AACF,EAAA;AACF","file":"tokenMemory.cjs","sourcesContent":["/**\n * Copyright 2025 © BeeAI a Series of LF Projects, LLC\n * SPDX-License-Identifier: Apache-2.0\n */\n\nimport { BaseMemory, MemoryFatalError } from \"@/memory/base.js\";\nimport * as R from \"remeda\";\nimport { shallowCopy } from \"@/serializer/utils.js\";\nimport { removeFromArray } from \"@/internals/helpers/array.js\";\nimport { map, sum } from \"remeda\";\nimport { ensureRange } from \"@/internals/helpers/number.js\";\nimport { Message } from \"@/backend/message.js\";\n\nexport interface Handlers {\n estimate: (messages: Message) => number;\n tokenize: (messages: Message[]) => Promise;\n removalSelector: (messages: Message[]) => Message;\n}\n\nconst simpleEstimate: Handlers[\"estimate\"] = (msg: Message) => Math.ceil(msg.text.length / 4);\nconst simpleTokenize: Handlers[\"tokenize\"] = async (msgs: Message[]) =>\n sum(map(msgs, simpleEstimate)); // TODO\n\nexport interface TokenMemoryInput {\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: Message[] = [];\n\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.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 tokenize: config?.handlers?.tokenize || simpleTokenize,\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: Message, index?: number) {\n if (this.maxTokens === null) {\n // TODO: improve\n this.maxTokens = 128_000;\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: Message) {\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 tokensCount = await this.handlers.tokenize([msg]);\n this.tokensByMessage.set(msg, { 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 threshold: this.threshold,\n syncThreshold: this.syncThreshold,\n messages: shallowCopy(this.messages),\n handlers: shallowCopy(this.handlers),\n maxTokens: this.maxTokens,\n tokensByMessage: this.messages\n .map((message) => [message, this.tokensByMessage.get(message)])\n .filter(([_, value]) => value !== undefined) as [Message, number][],\n };\n }\n\n loadSnapshot({ tokensByMessage, ...state }: ReturnType) {\n Object.assign(this, state, {\n tokensByMessage: new WeakMap(tokensByMessage),\n });\n }\n}\n"]}