node-llama-cpp 3.0.0-beta.2 → 3.0.0-beta.20

package/dist/{llamaEvaluator → evaluator}/LlamaChat/LlamaChat.js

@@ -1,10 +1,13 @@
 import { DisposeAggregator, DisposedError, EventRelay } from "lifecycle-utils";
-import { resolveChatWrapper } from "../../utils/resolveChatWrapper.js";
 import { removeNullFields } from "../../utils/removeNullFields.js";
 import { LlamaGrammarEvaluationState } from "../LlamaGrammarEvaluationState.js";
-import { AbortError } from "../../AbortError.js";
 import { StopGenerationDetector } from "../../utils/StopGenerationDetector.js";
 import { TokenStreamRegulator } from "../../utils/TokenStreamRegulator.js";
+import { UNKNOWN_UNICODE_CHAR } from "../../consts.js";
+import { getQueuedTokensBeforeStopTrigger } from "../../utils/getQueuedTokensBeforeStopTrigger.js";
+import { resolveChatWrapper } from "../../chatWrappers/utils/resolveChatWrapper.js";
+import { GeneralChatWrapper } from "../../chatWrappers/GeneralChatWrapper.js";
+import { getConsoleLogPrefix } from "../../utils/getConsoleLogPrefix.js";
 import { eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy } from "./utils/contextShiftStrategies/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.js";
 import { FunctionCallGrammar, LlamaFunctionCallValidationError } from "./utils/FunctionCallGrammar.js";
 const defaultContextShiftOptions = {
@@ -12,7 +15,6 @@ const defaultContextShiftOptions = {
     strategy: "eraseFirstResponseAndKeepFirstSystem",
     lastEvaluationMetadata: null
 };
-const UNKNOWN_UNICODE_CHAR = "\ufffd";
 export class LlamaChat {
     /** @internal */ _chatWrapper;
     /** @internal */ _disposeAggregator = new DisposeAggregator();
@@ -30,7 +32,14 @@ export class LlamaChat {
             this.dispose();
         }));
         this._disposeAggregator.add(this.onDispose.dispatchEvent);
-        this._chatWrapper = resolveChatWrapper(chatWrapper, contextSequence.model);
+        this._chatWrapper = chatWrapper === "auto"
+            ? (resolveChatWrapper({
+                bosString: contextSequence.model.tokens.bosString,
+                filename: contextSequence.model.filename,
+                fileInfo: contextSequence.model.fileInfo,
+                tokenizer: contextSequence.model.tokenizer
+            }) ?? new GeneralChatWrapper())
+            : chatWrapper;
     }
     dispose({ disposeSequence = this._autoDisposeSequence } = {}) {
         if (this._sequence == null)
@@ -63,12 +72,12 @@ export class LlamaChat {
     get model() {
         return this.sequence.model;
     }
-    async generateResponse(history, { onToken, signal, maxTokens, temperature, topK, topP, grammar, trimWhitespaceSuffix = false, repeatPenalty = {}, evaluationPriority = 5, functions, documentFunctionParams, contextShift = defaultContextShiftOptions, lastEvaluationContextWindow: { history: lastEvaluationContextWindowHistory, minimumOverlapPercentageToPreventContextShift = 0.5 } = {} } = {}) {
+    async generateResponse(history, { onToken, signal, stopOnAbortSignal = false, maxTokens, temperature, minP, topK, topP, grammar, trimWhitespaceSuffix = false, repeatPenalty = {}, tokenBias, evaluationPriority = 5, functions, documentFunctionParams, contextShift = defaultContextShiftOptions, customStopTriggers, lastEvaluationContextWindow: { history: lastEvaluationContextWindowHistory, minimumOverlapPercentageToPreventContextShift = 0.5 } = {} } = {}) {
         const functionsEnabled = (functions != null && Object.keys(functions).length > 0);
         if (grammar != null && functionsEnabled)
             throw new Error("Using both grammar and functions is not supported yet");
         if (signal?.aborted)
-            throw new AbortError();
+            throw signal.reason;
         if (this._sequence == null)
             throw new DisposedError();
         let resolvedHistory = this._sequence.isLoadedToMemory
@@ -81,7 +90,6 @@ export class LlamaChat {
             });
         const model = this._sequence.model;
         const context = this._sequence.context;
-        const eosToken = model.tokens.eos;
         const resolvedContextShift = {
             ...defaultContextShiftOptions,
             ...removeNullFields(contextShift)
@@ -99,7 +107,7 @@ export class LlamaChat {
             ? new LlamaGrammarEvaluationState({ grammar })
             : undefined;
         let functionsGrammar = functionsEnabled
-            ? new FunctionCallGrammar(functions, this._chatWrapper, false)
+            ? new FunctionCallGrammar(model._llama, functions, this._chatWrapper, false)
             : undefined;
         let functionsEvaluationState = (functionsEnabled && functionsGrammar != null)
             ? new LlamaGrammarEvaluationState({
@@ -108,6 +116,7 @@ export class LlamaChat {
             : undefined;
         const streamRegulator = new TokenStreamRegulator();
         const stopGenerationDetector = new StopGenerationDetector();
+        const customStopGenerationTriggersDetector = new StopGenerationDetector();
         const functionSyntaxStartDetector = new StopGenerationDetector();
         const functionSyntaxEndDetector = new StopGenerationDetector();
         const disengageInitiallyEngagedFunctionMode = new StopGenerationDetector();
@@ -121,8 +130,8 @@ export class LlamaChat {
         let lastContextWindowHistory = resolvedHistory;
         let lastHistoryCompressionMetadata = resolvedContextShift.lastEvaluationMetadata;
         const ensureNotAborted = () => {
-            if (signal?.aborted)
-                throw new AbortError();
+            if (signal?.aborted && (!stopOnAbortSignal || res.length === 0))
+                throw signal.reason;
             if (this._sequence == null)
                 throw new DisposedError();
         };
@@ -132,7 +141,7 @@ export class LlamaChat {
             let punishTokens = res.slice(-repeatPenaltyLastTokens);
             if (punishTokensFilter != null)
                 punishTokens = punishTokensFilter(punishTokens);
-            if (!penalizeNewLine) {
+            if (penalizeNewLine == null || !penalizeNewLine) {
                 const nlToken = model.tokens.nl;
                 if (nlToken != null)
                     punishTokens = punishTokens.filter(token => token !== nlToken);
@@ -173,7 +182,7 @@ export class LlamaChat {
                         ignoredStartTextTokens = mostExhaustiveTriggeredStop.stopTrigger
                             .map((stopTrigger) => {
                             if (typeof stopTrigger === "string")
-                                return model.tokenize(stopTrigger);
+                                return model.tokenize(stopTrigger, false, "trimLeadingSpace");
                             else
                                 return [stopTrigger];
                         })
@@ -181,7 +190,7 @@ export class LlamaChat {
                         const newPendingTokens = mostExhaustiveTriggeredStop.remainingGenerations
                             .map((generation) => {
                             if (typeof generation === "string")
-                                return model.tokenize(generation);
+                                return model.tokenize(generation, false, "trimLeadingSpace");
                             else
                                 return generation;
                         })
@@ -192,8 +201,11 @@ export class LlamaChat {
                 }
             }
         };
+        if (customStopTriggers != null)
+            StopGenerationDetector.resolveStopTriggers(customStopTriggers, model.tokenizer)
+                .map((stopTrigger) => customStopGenerationTriggersDetector.addStopTrigger(stopTrigger));
         if (grammar != null)
-            StopGenerationDetector.resolveStopTriggers(grammar.stopGenerationTriggers, model.tokenize)
+            StopGenerationDetector.resolveStopTriggers(grammar.stopGenerationTriggers, model.tokenizer)
                 .map((stopTrigger) => stopGenerationDetector.addStopTrigger(stopTrigger));
         if (functions != null && Object.keys(functions).length > 0)
             functionSyntaxStartDetector.addStopTrigger([this._chatWrapper.settings.functions.call.prefix]);
@@ -206,7 +218,7 @@ export class LlamaChat {
                 resolvedHistory: getResolvedHistoryWithCurrentModelResponse(),
                 resolvedContextShift,
                 lastHistoryCompressionMetadata,
-                pendingTokensCount: pendingTokens.length + queuedChunkTokens.length,
+                pendingTokensCount: ignoredStartTextTokens.length + pendingTokens.length + queuedChunkTokens.length,
                 isFirstEvaluation,
                 chatWrapper: this._chatWrapper,
                 lastEvaluationContextWindowHistory,
@@ -218,15 +230,15 @@ export class LlamaChat {
             });
             ensureNotAborted();
             if (generatedTokens === 0) {
-                StopGenerationDetector.resolveStopTriggers(ignoreStartText, model.tokenize)
+                StopGenerationDetector.resolveStopTriggers(ignoreStartText, model.tokenizer)
                     .map((stopTrigger) => ignoreStartTextDetector.addStopTrigger(stopTrigger));
                 if (functionsEnabled) {
                     initiallyEngagedFunctionMode = functionCallInitiallyEngaged;
-                    StopGenerationDetector.resolveStopTriggers(disengageInitiallyEngagedFunctionCall, model.tokenize)
+                    StopGenerationDetector.resolveStopTriggers(disengageInitiallyEngagedFunctionCall, model.tokenizer)
                         .map((stopTrigger) => disengageInitiallyEngagedFunctionMode.addStopTrigger(stopTrigger));
                     if (initiallyEngagedFunctionMode) {
                         inFunctionEvaluationMode = true;
-                        functionsGrammar = new FunctionCallGrammar(functions, this._chatWrapper, true);
+                        functionsGrammar = new FunctionCallGrammar(model._llama, functions, this._chatWrapper, true);
                         functionsEvaluationState = new LlamaGrammarEvaluationState({
                             grammar: functionsGrammar
                         });
@@ -239,10 +251,10 @@ export class LlamaChat {
             lastContextWindowHistory = contextWindowHistory;
             const contextWindowLastModelResponse = getLastTextModelResponseFromChatHistory(contextWindowHistory);
             const contextWindowsRes = [];
-            StopGenerationDetector.resolveStopTriggers(stopGenerationTriggers, model.tokenize)
+            StopGenerationDetector.resolveStopTriggers(stopGenerationTriggers, model.tokenizer)
                 .map((stopTrigger) => stopGenerationDetector.addStopTrigger(stopTrigger));
             if (functionsGrammar != null)
-                StopGenerationDetector.resolveStopTriggers(functionsGrammar.stopGenerationTriggers, model.tokenize)
+                StopGenerationDetector.resolveStopTriggers(functionsGrammar.stopGenerationTriggers, model.tokenizer)
                     .map((stopTrigger) => functionSyntaxEndDetector.addStopTrigger(stopTrigger));
             let { firstDifferentIndex } = this._sequence.compareContextTokens(tokens);
             // we need to decode at least one token to generate a response
@@ -257,7 +269,7 @@ export class LlamaChat {
                 ensureNotAborted();
             }
             const evaluationIterator = this._sequence.evaluate(tokens, removeNullFields({
-                temperature, topK, topP,
+                temperature, minP, topK, topP,
                 grammarEvaluationState: () => {
                     if (inFunctionEvaluationMode)
                         return functionsEvaluationState;
@@ -269,194 +281,279 @@ export class LlamaChat {
                     frequencyPenalty,
                     presencePenalty
                 },
+                tokenBias,
                 evaluationPriority,
-                yieldEosToken: true
+                yieldEogToken: true
             }));
-            for await (const token of evaluationIterator) {
-                ensureNotAborted();
-                generatedTokens++;
-                const tokens = [token];
-                const text = model.detokenize([token]);
-                const queuedTokenRelease = streamRegulator.addChunk({ tokens, text });
-                if (initiallyEngagedFunctionMode)
-                    disengageInitiallyEngagedFunctionMode.recordGeneration({ text, tokens, startNewChecks: generatedTokens === 1 });
-                if (text === UNKNOWN_UNICODE_CHAR || ((grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) && text.trim() === "")) {
-                    locksToReleaseOnValidGeneration.push(queuedTokenRelease.createTextIndexLock(0));
-                }
-                else {
-                    while (locksToReleaseOnValidGeneration.length > 0)
-                        locksToReleaseOnValidGeneration.shift().dispose();
-                }
-                functionSyntaxStartDetector.recordGeneration({ text, tokens, queuedTokenRelease });
-                if (initiallyEngagedFunctionMode && disengageInitiallyEngagedFunctionMode.hasTriggeredStops) {
-                    initiallyEngagedFunctionMode = false;
-                    let shouldStopFunctionEvaluationMode = !functionSyntaxStartDetector.hasTriggeredStops;
-                    if (!shouldStopFunctionEvaluationMode && functionsEnabled && functionsGrammar != null) {
-                        const functionCallText = model.detokenize([...functionCallTokens, ...tokens]);
-                        try {
-                            const functionName = functionsGrammar.parseFunctionNameFromPartialCall(functionCallText, {
-                                enableInternalBuiltinFunctions: true,
-                                initialFunctionCallEngaged: true
-                            });
-                            const internalBuiltinFunctions = this._chatWrapper.getInternalBuiltinFunctions({ initialFunctionCallEngaged: true });
-                            if (internalBuiltinFunctions[functionName] != null) {
-                                shouldStopFunctionEvaluationMode = true;
+            try {
+                let currentIteration = await evaluationIterator.next();
+                while (currentIteration.done !== true) {
+                    const token = currentIteration.value;
+                    let replacementToken = undefined;
+                    ensureNotAborted();
+                    generatedTokens++;
+                    const tokens = [token];
+                    const text = model.detokenize([token]);
+                    const queuedTokenRelease = streamRegulator.addChunk({ tokens, text });
+                    if (initiallyEngagedFunctionMode)
+                        disengageInitiallyEngagedFunctionMode.recordGeneration({ text, tokens, startNewChecks: generatedTokens === 1 });
+                    if (text === UNKNOWN_UNICODE_CHAR || ((grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) && text.trim() === "")) {
+                        locksToReleaseOnValidGeneration.push(queuedTokenRelease.createTextIndexLock(0));
+                    }
+                    else {
+                        while (locksToReleaseOnValidGeneration.length > 0)
+                            locksToReleaseOnValidGeneration.shift().dispose();
+                    }
+                    functionSyntaxStartDetector.recordGeneration({ text, tokens, queuedTokenRelease });
+                    if (initiallyEngagedFunctionMode && disengageInitiallyEngagedFunctionMode.hasTriggeredStops) {
+                        initiallyEngagedFunctionMode = false;
+                        let shouldStopFunctionEvaluationMode = !functionSyntaxStartDetector.hasTriggeredStops;
+                        if (!shouldStopFunctionEvaluationMode && functionsEnabled && functionsGrammar != null) {
+                            const functionCallText = model.detokenize([...functionCallTokens, ...tokens]);
+                            try {
+                                const functionName = functionsGrammar.parseFunctionNameFromPartialCall(functionCallText, {
+                                    enableInternalBuiltinFunctions: true,
+                                    initialFunctionCallEngaged: true
+                                });
+                                const internalBuiltinFunctions = this._chatWrapper.getInternalBuiltinFunctions({ initialFunctionCallEngaged: true });
+                                if (internalBuiltinFunctions[functionName] != null) {
+                                    shouldStopFunctionEvaluationMode = true;
+                                }
+                            }
+                            catch (err) {
+                                if (!(err instanceof LlamaFunctionCallValidationError))
+                                    throw err;
                             }
                         }
-                        catch (err) {
-                            if (!(err instanceof LlamaFunctionCallValidationError))
-                                throw err;
+                        if (shouldStopFunctionEvaluationMode) {
+                            inFunctionEvaluationMode = false;
+                            functionsGrammar = new FunctionCallGrammar(model._llama, functions, this._chatWrapper, false);
+                            functionsEvaluationState = new LlamaGrammarEvaluationState({
+                                grammar: functionsGrammar
+                            });
+                            functionCallTokens.length = 0;
+                            while (functionCallTokenSyntaxLocks.length > 0)
+                                functionCallTokenSyntaxLocks.shift().dispose();
+                            functionSyntaxStartDetector.clearInProgressStops();
+                            functionSyntaxStartDetector.clearTriggeredStops();
+                            functionSyntaxEndDetector.clearInProgressStops();
+                            functionSyntaxEndDetector.clearTriggeredStops();
                         }
                     }
-                    if (shouldStopFunctionEvaluationMode) {
-                        inFunctionEvaluationMode = false;
-                        functionsGrammar = new FunctionCallGrammar(functions, this._chatWrapper, false);
-                        functionsEvaluationState = new LlamaGrammarEvaluationState({
-                            grammar: functionsGrammar
-                        });
-                        functionCallTokens.length = 0;
-                        while (functionCallTokenSyntaxLocks.length > 0)
-                            functionCallTokenSyntaxLocks.shift().dispose();
-                        functionSyntaxStartDetector.clearInProgressStops();
-                        functionSyntaxStartDetector.clearTriggeredStops();
-                        functionSyntaxEndDetector.clearInProgressStops();
-                        functionSyntaxEndDetector.clearTriggeredStops();
+                    if (!inFunctionEvaluationMode && functionsEnabled && functionsGrammar != null &&
+                        functionSyntaxStartDetector.hasTriggeredStops && functionsEvaluationState != null) {
+                        inFunctionEvaluationMode = true;
+                        functionCallTokenSyntaxLocks.push(queuedTokenRelease.createTextIndexLock(0));
+                        stopGenerationDetector.clearTriggeredStops();
+                        stopGenerationDetector.clearInProgressStops();
+                        customStopGenerationTriggersDetector.clearTriggeredStops();
+                        customStopGenerationTriggersDetector.clearInProgressStops();
+                        pendingTokens.push(...streamRegulator.popFreeChunkTokens());
+                        const triggeredStops = functionSyntaxStartDetector.getTriggeredStops();
+                        const partiallyFreeTokens = streamRegulator.getPartiallyFreeChunk(model.tokenizer);
+                        const queuedTokensBeforeStopTrigger = getQueuedTokensBeforeStopTrigger(triggeredStops, partiallyFreeTokens, model.tokenizer);
+                        pendingTokens.push(...queuedTokensBeforeStopTrigger);
+                        const [firstRemainingGenerationAfterStop] = triggeredStops
+                            .map((stopTrigger) => stopTrigger.remainingGenerations)
+                            .filter((remainingGenerations) => remainingGenerations.length > 0)
+                            .flat(1);
+                        const remainingTextAfterStop = (firstRemainingGenerationAfterStop == null || firstRemainingGenerationAfterStop.length === 0)
+                            ? ""
+                            : typeof firstRemainingGenerationAfterStop === "string"
+                                ? firstRemainingGenerationAfterStop
+                                : model.detokenize(firstRemainingGenerationAfterStop);
+                        functionCallTokens.push(...model.tokenize(this._chatWrapper.settings.functions.call.prefix, false, "trimLeadingSpace"));
+                        for (const functionCallToken of functionCallTokens)
+                            context._acceptTokenOnGrammarEvaluationState(functionsEvaluationState, functionCallToken);
+                        // these tokens have to be verified that they match the function calling syntax grammar before they can be accepted,
+                        // or the context state should be modified to not include the incompatible tokens
+                        const remainingTextTokens = model.tokenize(remainingTextAfterStop, false, "trimLeadingSpace");
+                        let unfitTokens = [];
+                        for (let i = 0; i < remainingTextTokens.length; i++) {
+                            const remainingToken = remainingTextTokens[i];
+                            const canBeNextToken = context._canBeNextTokenForGrammarEvaluationState(functionsEvaluationState, remainingToken);
+                            if (!canBeNextToken) {
+                                unfitTokens = remainingTextTokens.slice(i);
+                                break;
+                            }
+                            context._acceptTokenOnGrammarEvaluationState(functionsEvaluationState, remainingToken);
+                            functionCallTokens.push(remainingToken);
+                        }
+                        if (unfitTokens.length > 0) {
+                            const unfitTokensText = model.detokenize(unfitTokens); // the current token text must end with it
+                            const currentTokenText = queuedTokenRelease.text;
+                            let replacementTokens;
+                            if (!currentTokenText.endsWith(unfitTokensText)) {
+                                console.warn(getConsoleLogPrefix() + "The current token text does not end with the unfit function call syntax tokens text");
+                                replacementTokens = remainingTextTokens.slice(0, -unfitTokens.length);
+                            }
+                            else {
+                                const newCurrentTokensText = currentTokenText.slice(0, -unfitTokensText.length);
+                                replacementTokens = model.tokenize(newCurrentTokensText, false, "trimLeadingSpace");
+                            }
+                            if (replacementTokens.length > 0) {
+                                replacementToken = replacementTokens[0];
+                                queuedTokenRelease.modifyTokensAndText(replacementTokens, model.detokenize([replacementToken]));
+                            }
+                        }
                     }
-                }
-                if (!inFunctionEvaluationMode && functionsEnabled && functionsGrammar != null &&
-                    functionSyntaxStartDetector.hasTriggeredStops && functionsEvaluationState != null) {
-                    inFunctionEvaluationMode = true;
-                    functionCallTokenSyntaxLocks.push(queuedTokenRelease.createTextIndexLock(0));
-                    stopGenerationDetector.clearTriggeredStops();
-                    stopGenerationDetector.clearInProgressStops();
-                    pendingTokens.push(...streamRegulator.popFreeChunkTokens());
-                    const triggeredStops = functionSyntaxStartDetector.getTriggeredStops();
-                    const partiallyFreeTokens = streamRegulator.getPartiallyFreeChunk();
-                    const queuedTokensBeforeStopTrigger = getQueuedTokensBeforeStopTrigger(triggeredStops, partiallyFreeTokens, model.tokenize);
-                    pendingTokens.push(...queuedTokensBeforeStopTrigger);
-                    const [firstRemainingGenerationAfterStop] = triggeredStops
-                        .map((stopTrigger) => stopTrigger.remainingGenerations)
-                        .filter((remainingGenerations) => remainingGenerations.length > 0)
-                        .flat(1);
-                    const remainingTextAfterStop = (firstRemainingGenerationAfterStop == null || firstRemainingGenerationAfterStop.length === 0)
-                        ? ""
-                        : typeof firstRemainingGenerationAfterStop === "string"
-                            ? firstRemainingGenerationAfterStop
-                            : model.detokenize(firstRemainingGenerationAfterStop);
-                    functionCallTokens.push(...model.tokenize(this._chatWrapper.settings.functions.call.prefix + remainingTextAfterStop));
-                    for (const functionCallToken of functionCallTokens)
-                        context._acceptTokenOnGrammarEvaluationState(functionsEvaluationState, functionCallToken);
-                }
-                else if (inFunctionEvaluationMode) {
-                    functionCallTokens.push(...tokens);
-                    functionCallTokenSyntaxLocks.push(queuedTokenRelease.createTextIndexLock(0));
-                    functionSyntaxEndDetector.recordGeneration({ text, tokens, queuedTokenRelease });
-                }
-                if (inFunctionEvaluationMode && functionSyntaxEndDetector.hasTriggeredStops && functionsGrammar != null) {
-                    const functionCallText = model.detokenize(functionCallTokens);
-                    const functionCall = functionsGrammar.parseFunctionCall(functionCallText);
-                    let modelResponse = model.detokenize(res);
-                    let contextWindowModelResponse = model.detokenize(contextWindowsRes);
-                    if (grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) {
-                        modelResponse = modelResponse.trimEnd();
-                        contextWindowModelResponse = contextWindowModelResponse.trimEnd();
+                    else if (inFunctionEvaluationMode) {
+                        functionCallTokens.push(...tokens);
+                        functionCallTokenSyntaxLocks.push(queuedTokenRelease.createTextIndexLock(0));
+                        functionSyntaxEndDetector.recordGeneration({ text, tokens, queuedTokenRelease });
                     }
-                    return {
-                        response: modelResponse,
-                        lastEvaluation: {
-                            contextWindow: setLastModelTextResponseInChatHistory(lastContextWindowHistory, contextWindowLastModelResponse + contextWindowModelResponse),
-                            cleanHistory: setLastModelTextResponseInChatHistory(resolvedHistory, lastModelResponse + modelResponse),
-                            contextShiftMetadata: lastHistoryCompressionMetadata
-                        },
-                        // prevent infinite TS type instantiation
-                        functionCall: functionCall,
-                        metadata: {
-                            stopReason: "functionCall"
+                    if (inFunctionEvaluationMode && functionSyntaxEndDetector.hasTriggeredStops && functionsGrammar != null) {
+                        const functionCallText = model.detokenize(functionCallTokens);
+                        const functionCall = functionsGrammar.parseFunctionCall(functionCallText);
+                        let modelResponse = model.detokenize(res);
+                        let contextWindowModelResponse = model.detokenize(contextWindowsRes);
+                        if (grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) {
+                            modelResponse = modelResponse.trimEnd();
+                            contextWindowModelResponse = contextWindowModelResponse.trimEnd();
                         }
-                    };
-                }
-                if (!inFunctionEvaluationMode)
-                    stopGenerationDetector.recordGeneration({ text, tokens, queuedTokenRelease });
-                pendingTokens.push(...streamRegulator.popFreeChunkTokens());
-                removeFoundStartIgnoreTextsFromPendingTokens();
-                if (stopGenerationDetector.hasTriggeredStops || token === eosToken) {
-                    const triggeredStops = stopGenerationDetector.getTriggeredStops();
-                    const partiallyFreeTokens = streamRegulator.getPartiallyFreeChunk();
-                    const queuedTokensBeforeStopTrigger = getQueuedTokensBeforeStopTrigger(triggeredStops, partiallyFreeTokens, model.tokenize);
-                    pendingTokens.push(...queuedTokensBeforeStopTrigger);
-                    const [firstRemainingGenerationAfterStop] = triggeredStops
-                        .map((stopTrigger) => stopTrigger.remainingGenerations)
-                        .filter((remainingGenerations) => remainingGenerations.length > 0)
-                        .flat(1);
-                    removeFoundStartIgnoreTextsFromPendingTokens();
-                    if (pendingTokens.length > 0)
-                        onToken?.(pendingTokens.slice());
-                    res.push(...pendingTokens);
-                    contextWindowsRes.push(...pendingTokens);
-                    pendingTokens.length = 0;
-                    let modelResponse = model.detokenize(res);
-                    let contextWindowModelResponse = model.detokenize(contextWindowsRes);
-                    if (grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) {
-                        modelResponse = modelResponse.trimEnd();
-                        contextWindowModelResponse = contextWindowModelResponse.trimEnd();
+                        return {
+                            response: modelResponse,
+                            lastEvaluation: {
+                                contextWindow: setLastModelTextResponseInChatHistory(lastContextWindowHistory, contextWindowLastModelResponse + contextWindowModelResponse),
+                                cleanHistory: setLastModelTextResponseInChatHistory(resolvedHistory, lastModelResponse + modelResponse),
+                                contextShiftMetadata: lastHistoryCompressionMetadata
+                            },
+                            // prevent infinite TS type instantiation
+                            functionCall: functionCall,
+                            metadata: {
+                                stopReason: "functionCall"
+                            }
+                        };
                     }
-                    return {
-                        response: modelResponse,
-                        lastEvaluation: {
-                            contextWindow: setLastModelTextResponseInChatHistory(lastContextWindowHistory, contextWindowLastModelResponse + contextWindowModelResponse),
-                            cleanHistory: setLastModelTextResponseInChatHistory(resolvedHistory, lastModelResponse + modelResponse),
-                            contextShiftMetadata: lastHistoryCompressionMetadata
-                        },
-                        metadata: {
-                            remainingGenerationAfterStop: firstRemainingGenerationAfterStop,
-                            stopReason: token === eosToken
-                                ? "eosToken"
-                                : "stopGenerationTrigger"
-                        }
-                    };
-                }
-                const maxTokensTriggered = maxTokens != null && maxTokens > 0 && generatedTokens >= maxTokens;
-                if (res.length === 0) {
-                    ignoreStartTextDetector.clearInProgressStops();
-                    ignoreStartTextDetector.clearTriggeredStops();
-                    ignoreStartTextDetector.recordGeneration({
-                        text: model.detokenize(pendingTokens),
-                        tokens: pendingTokens
-                    });
-                }
-                if (pendingTokens.length > 0 && (maxTokensTriggered || !ignoreStartTextDetector.hasInProgressStops)) {
+                    if (!inFunctionEvaluationMode) {
+                        stopGenerationDetector.recordGeneration({ text, tokens, queuedTokenRelease });
+                        customStopGenerationTriggersDetector.recordGeneration({ text, tokens, queuedTokenRelease });
+                    }
+                    pendingTokens.push(...streamRegulator.popFreeChunkTokens());
                     removeFoundStartIgnoreTextsFromPendingTokens();
-                    if (pendingTokens.length > 0) {
-                        onToken?.(pendingTokens.slice());
+                    if (stopGenerationDetector.hasTriggeredStops || customStopGenerationTriggersDetector.hasTriggeredStops ||
+                        model.isEogToken(token)) {
+                        stopGenerationDetector.clearInProgressStops();
+                        customStopGenerationTriggersDetector.clearInProgressStops();
+                        pendingTokens.push(...streamRegulator.popFreeChunkTokens());
+                        const triggeredStops = stopGenerationDetector.hasTriggeredStops
+                            ? stopGenerationDetector.getTriggeredStops()
+                            : customStopGenerationTriggersDetector.getTriggeredStops();
+                        const partiallyFreeTokens = streamRegulator.getPartiallyFreeChunk(model.tokenizer);
+                        const queuedTokensBeforeStopTrigger = getQueuedTokensBeforeStopTrigger(triggeredStops, partiallyFreeTokens, model.tokenizer);
+                        pendingTokens.push(...queuedTokensBeforeStopTrigger);
+                        const [firstRemainingGenerationAfterStop] = triggeredStops
+                            .map((stopTrigger) => stopTrigger.remainingGenerations)
+                            .filter((remainingGenerations) => remainingGenerations.length > 0)
+                            .flat(1);
+                        removeFoundStartIgnoreTextsFromPendingTokens();
+                        if (pendingTokens.length > 0)
+                            onToken?.(pendingTokens.slice());
                         res.push(...pendingTokens);
                         contextWindowsRes.push(...pendingTokens);
                         pendingTokens.length = 0;
-                    }
-                }
-                if (maxTokensTriggered) {
-                    let modelResponse = model.detokenize(res);
-                    let contextWindowModelResponse = model.detokenize(contextWindowsRes);
-                    if (grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) {
-                        modelResponse = modelResponse.trimEnd();
-                        contextWindowModelResponse = contextWindowModelResponse.trimEnd();
-                    }
-                    return {
-                        response: modelResponse,
-                        lastEvaluation: {
+                        let modelResponse = model.detokenize(res);
+                        let contextWindowModelResponse = model.detokenize(contextWindowsRes);
+                        if (grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) {
+                            modelResponse = modelResponse.trimEnd();
+                            contextWindowModelResponse = contextWindowModelResponse.trimEnd();
+                        }
+                        const lastEvaluation = {
                             contextWindow: setLastModelTextResponseInChatHistory(lastContextWindowHistory, contextWindowLastModelResponse + contextWindowModelResponse),
                             cleanHistory: setLastModelTextResponseInChatHistory(resolvedHistory, lastModelResponse + modelResponse),
                             contextShiftMetadata: lastHistoryCompressionMetadata
-                        },
-                        metadata: {
-                            stopReason: "maxTokens"
+                        };
+                        const isEogToken = model.isEogToken(token);
+                        if (isEogToken || stopGenerationDetector.hasTriggeredStops) {
+                            return {
+                                response: modelResponse,
+                                lastEvaluation,
+                                metadata: {
+                                    remainingGenerationAfterStop: firstRemainingGenerationAfterStop,
+                                    stopReason: isEogToken
+                                        ? "eogToken"
+                                        : "stopGenerationTrigger"
+                                }
+                            };
                         }
-                    };
-                }
-                if (this._sequence.nextTokenIndex >= context.contextSize) {
-                    shouldContextShift = true;
-                    break;
+                        return {
+                            response: modelResponse,
+                            lastEvaluation,
+                            metadata: {
+                                remainingGenerationAfterStop: firstRemainingGenerationAfterStop,
+                                stopReason: "customStopTrigger",
+                                customStopTrigger: triggeredStops[0].stopTrigger
+                            }
+                        };
+                    }
+                    const maxTokensTriggered = maxTokens != null && maxTokens > 0 && generatedTokens >= maxTokens;
+                    if (res.length === 0) {
+                        ignoreStartTextDetector.clearInProgressStops();
+                        ignoreStartTextDetector.clearTriggeredStops();
+                        ignoreStartTextDetector.recordGeneration({
+                            text: model.detokenize(pendingTokens),
+                            tokens: pendingTokens
+                        });
+                    }
+                    if (pendingTokens.length > 0 && (maxTokensTriggered || !ignoreStartTextDetector.hasInProgressStops)) {
+                        removeFoundStartIgnoreTextsFromPendingTokens();
+                        if (pendingTokens.length > 0) {
+                            onToken?.(pendingTokens.slice());
+                            res.push(...pendingTokens);
+                            contextWindowsRes.push(...pendingTokens);
+                            pendingTokens.length = 0;
+                        }
+                    }
+                    if (maxTokensTriggered) {
+                        let modelResponse = model.detokenize(res);
+                        let contextWindowModelResponse = model.detokenize(contextWindowsRes);
+                        if (grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) {
+                            modelResponse = modelResponse.trimEnd();
+                            contextWindowModelResponse = contextWindowModelResponse.trimEnd();
+                        }
+                        return {
+                            response: modelResponse,
+                            lastEvaluation: {
+                                contextWindow: setLastModelTextResponseInChatHistory(lastContextWindowHistory, contextWindowLastModelResponse + contextWindowModelResponse),
+                                cleanHistory: setLastModelTextResponseInChatHistory(resolvedHistory, lastModelResponse + modelResponse),
+                                contextShiftMetadata: lastHistoryCompressionMetadata
+                            },
+                            metadata: {
+                                stopReason: "maxTokens"
+                            }
+                        };
+                    }
+                    if (this._sequence.nextTokenIndex >= context.contextSize - 1) {
+                        shouldContextShift = true;
+                        break;
+                    }
+                    if (signal?.aborted && stopOnAbortSignal) {
+                        if (res.length === 0)
+                            throw signal.reason;
+                        let modelResponse = model.detokenize(res);
+                        let contextWindowModelResponse = model.detokenize(contextWindowsRes);
+                        if (grammar?.trimWhitespaceSuffix || trimWhitespaceSuffix) {
+                            modelResponse = modelResponse.trimEnd();
+                            contextWindowModelResponse = contextWindowModelResponse.trimEnd();
+                        }
+                        return {
+                            response: modelResponse,
+                            lastEvaluation: {
+                                contextWindow: setLastModelTextResponseInChatHistory(lastContextWindowHistory, contextWindowLastModelResponse + contextWindowModelResponse),
+                                cleanHistory: setLastModelTextResponseInChatHistory(resolvedHistory, lastModelResponse + modelResponse),
+                                contextShiftMetadata: lastHistoryCompressionMetadata
+                            },
+                            metadata: {
+                                stopReason: "abort"
+                            }
+                        };
+                    }
+                    currentIteration = await evaluationIterator.next(replacementToken);
                 }
             }
+            finally {
+                await evaluationIterator.return();
+            }
             isFirstEvaluation = false;
             if (shouldContextShift)
                 continue;
@@ -530,33 +627,13 @@ async function compressHistoryToFitContextSize({ history, contextShiftSize, cont
         lastShiftMetadata: contextShiftLastEvaluationMetadata
     });
     if (!checkIfHistoryFitsContext(chatHistory))
-        throw new Error("The default context shift strategy did not return a history that fits the context size");
+        throw new Error("The default context shift strategy did not return a history that fits the context size. " +
+            "This may happen due to the system prompt being too long");
     return {
         compressedHistory: chatHistory,
         metadata
     };
 }
-function getQueuedTokensBeforeStopTrigger(triggeredStops, partiallyFreeTokens, tokenizer) {
-    if (partiallyFreeTokens.tokens.length === 0 && partiallyFreeTokens.text.length === 0)
-        return [];
-    else if (partiallyFreeTokens.tokens.length !== 0 && partiallyFreeTokens.text.length === 0)
-        return partiallyFreeTokens.tokens;
-    else if (partiallyFreeTokens.tokens.length === 0 && partiallyFreeTokens.text.length !== 0)
-        return tokenizer(partiallyFreeTokens.text);
-    const triggerThatStartsWithStringIndex = triggeredStops.findIndex((trigger) => trigger.stopTrigger.length > 0 && typeof trigger.stopTrigger[0] === "string");
-    const triggerThatStartsWithTokenIndex = triggeredStops.findIndex((trigger) => trigger.stopTrigger.length > 0 && typeof trigger.stopTrigger[0] !== "string");
-    if (triggerThatStartsWithTokenIndex > 0 && triggerThatStartsWithStringIndex < 0)
-        return partiallyFreeTokens.tokens;
-    else if (triggerThatStartsWithStringIndex > 0 && triggerThatStartsWithTokenIndex < 0)
-        return tokenizer(partiallyFreeTokens.text);
-    const stringTokens = tokenizer(partiallyFreeTokens.text);
-    if (stringTokens.length === partiallyFreeTokens.tokens.length &&
-        stringTokens.every((value, index) => value === partiallyFreeTokens.tokens[index]))
-        return stringTokens;
-    else if (triggerThatStartsWithStringIndex < triggerThatStartsWithTokenIndex)
-        return stringTokens;
-    return partiallyFreeTokens.tokens;
-}
 function getLastTextModelResponseFromChatHistory(chatHistory) {
     if (chatHistory.length === 0 || chatHistory[chatHistory.length - 1].type !== "model")
         return "";
@@ -604,7 +681,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
             availableFunctions: functions,
             documentFunctionParams
         });
-        const tokens = contextText.tokenize(model.tokenize);
+        const tokens = contextText.tokenize(model.tokenizer);
         if (tokens.length + pendingTokensCount + minFreeContextTokens < context.contextSize) {
             const { firstDifferentIndex } = sequence.compareContextTokens(tokens);
             const existingEvaluationPercentage = firstDifferentIndex / tokens.length;
@@ -630,11 +707,11 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
             : resolvedContextShift.size;
         const { compressedHistory, metadata } = await compressHistoryToFitContextSize({
             history: resolvedHistory,
-            contextShiftSize: Math.max(contextShiftSize, minFreeContextTokens) + pendingTokensCount,
+            contextShiftSize: Math.max(minFreeContextTokens, Math.min(contextShiftSize, context.contextSize - pendingTokensCount)) + pendingTokensCount,
             contextShiftStrategy: resolvedContextShift.strategy,
             contextShiftLastEvaluationMetadata: resolvedContextShift.lastEvaluationMetadata,
             contextSize: context.contextSize,
-            tokenizer: model.tokenize,
+            tokenizer: model.tokenizer,
             chatWrapper: chatWrapper,
             functions,
             documentFunctionParams
@@ -646,7 +723,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
         return {
             history: compressedHistory,
             stopGenerationTriggers,
-            tokens: contextText.tokenize(model.tokenize),
+            tokens: contextText.tokenize(model.tokenizer),
             newResolvedHistory: resolvedHistory,
             newHistoryCompressionMetadata: metadata,
             ignoreStartText: ignoreStartText ?? [],
@@ -659,7 +736,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
             availableFunctions: functions,
             documentFunctionParams
         });
-        const tokens = contextText.tokenize(model.tokenize);
+        const tokens = contextText.tokenize(model.tokenizer);
         if (tokens.length + pendingTokensCount + minFreeContextTokens < context.contextSize)
             return {
                 history: resolvedHistory,
@@ -672,16 +749,16 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
                 disengageInitiallyEngagedFunctionCall: functionCall?.disengageInitiallyEngaged ?? []
             };
     }
-    const contextShiftSize = resolvedContextShift.size instanceof Function
+    const contextShiftSize = Math.min(context.contextSize, Math.max(1, Math.floor(resolvedContextShift.size instanceof Function
         ? await resolvedContextShift.size(sequence)
-        : resolvedContextShift.size;
+        : resolvedContextShift.size)));
     const { compressedHistory, metadata } = await compressHistoryToFitContextSize({
         history: resolvedHistory,
-        contextShiftSize: Math.max(contextShiftSize, minFreeContextTokens) + pendingTokensCount,
+        contextShiftSize: Math.max(minFreeContextTokens, Math.min(contextShiftSize, context.contextSize - pendingTokensCount)) + pendingTokensCount,
         contextShiftStrategy: resolvedContextShift.strategy,
         contextShiftLastEvaluationMetadata: resolvedContextShift.lastEvaluationMetadata,
         contextSize: context.contextSize,
-        tokenizer: model.tokenize,
+        tokenizer: model.tokenizer,
         chatWrapper: chatWrapper,
         functions,
         documentFunctionParams
@@ -693,7 +770,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
     return {
         history: compressedHistory,
         stopGenerationTriggers,
-        tokens: contextText.tokenize(model.tokenize),
+        tokens: contextText.tokenize(model.tokenizer),
         newResolvedHistory: resolvedHistory,
         newHistoryCompressionMetadata: metadata,
         ignoreStartText: ignoreStartText ?? [],