node-llama-cpp 3.4.3 → 3.6.0

package/dist/evaluator/LlamaChat/LlamaChat.js

@@ -1,4 +1,5 @@
 import { DisposeAggregator, DisposedError, EventRelay, withLock } from "lifecycle-utils";
+import { isChatModelResponseFunctionCall, isChatModelResponseSegment, allSegmentTypes } from "../../types.js";
 import { removeNullFields } from "../../utils/removeNullFields.js";
 import { LlamaGrammarEvaluationState } from "../LlamaGrammarEvaluationState.js";
 import { LlamaText, SpecialToken } from "../../utils/LlamaText.js";
@@ -11,6 +12,7 @@ import { safeEventCallback } from "../../utils/safeEventCallback.js";
 import { pushAll } from "../../utils/pushAll.js";
 import { resolveLastTokens } from "../../utils/resolveLastTokens.js";
 import { LlamaSampler } from "../LlamaContext/LlamaSampler.js";
+import { getChatWrapperSegmentDefinition } from "../../utils/getChatWrapperSegmentDefinition.js";
 import { eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy } from "./utils/contextShiftStrategies/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.js";
 import { FunctionCallNameGrammar } from "./utils/FunctionCallNameGrammar.js";
 import { FunctionCallParamsGrammar } from "./utils/FunctionCallParamsGrammar.js";
@@ -76,11 +78,12 @@ export class LlamaChat {
         return this.sequence.model;
     }
     async generateResponse(history, options = {}) {
-        const { onTextChunk, onToken, signal, stopOnAbortSignal = false, maxTokens, temperature, minP, topK, topP, seed, grammar, trimWhitespaceSuffix = defaultTrimWhitespaceSuffix, repeatPenalty = {}, tokenBias, evaluationPriority = defaultEvaluationPriority, functions, onFunctionCall, documentFunctionParams, maxParallelFunctionCalls, contextShift = defaultContextShiftOptions, customStopTriggers, lastEvaluationContextWindow: { history: lastEvaluationContextWindowHistory, minimumOverlapPercentageToPreventContextShift = 0.5 } = {} } = options;
+        const { onTextChunk, onToken, onResponseChunk, signal, stopOnAbortSignal = false, maxTokens, temperature, minP, topK, topP, seed, grammar, trimWhitespaceSuffix = defaultTrimWhitespaceSuffix, repeatPenalty = {}, tokenBias, evaluationPriority = defaultEvaluationPriority, functions, onFunctionCall, documentFunctionParams, maxParallelFunctionCalls, contextShift = defaultContextShiftOptions, customStopTriggers, lastEvaluationContextWindow: { history: lastEvaluationContextWindowHistory, minimumOverlapPercentageToPreventContextShift = 0.5 } = {} } = options;
         this.sequence.tokenPredictor?.updateInputTokens?.(this.model.tokenize(findLastUserMessageInChatHistory(history)?.text ?? ""));
         const generateResponseState = new GenerateResponseState(this, this._chatWrapper, history, {
             onTextChunk,
             onToken,
+            onResponseChunk,
             signal,
             stopOnAbortSignal,
             maxTokens,
@@ -110,6 +113,7 @@ export class LlamaChat {
         return await withLock(this._chatLock, "evaluate", signal, async () => {
             try {
                 generateResponseState.ensureLastHistoryItemIsModel();
+                generateResponseState.ensureReopenedThoughtSegmentAfterFunctionCallsIfNeeded();
                 const loadContextWindow = async (avoidReloadingHistory = false) => {
                     await generateResponseState.loadContextWindow(generateResponseState.getResolvedHistoryWithCurrentModelResponse(), generateResponseState.getContextWindowsHistoryWithCurrentModelResponse(), false, avoidReloadingHistory);
                 };
@@ -134,23 +138,25 @@ export class LlamaChat {
                     await generateResponseState.alignCurrentSequenceStateWithCurrentTokens();
                     await generateResponseState.createNewEvaluationIterator();
                     while (await generateResponseState.iterateEvaluation()) {
-                        generateResponseState.waitOnPartialCharactersOrWhiteSpaceTokens();
-                        generateResponseState.detectAndHandleFunctionStartSyntax();
-                        if (generateResponseState.functionEvaluationMode !== false) {
-                            generateResponseState.canAvoidReloadingHistory = false;
-                            generateResponseState.releasePartiallyFreeTokensBeforeFunctionCallStart();
-                            const functionsCallsRes = await generateResponseState.enterFunctionCallingLoop(loadContextWindowForFunctionCallingLoop);
-                            if (functionsCallsRes != null)
-                                return functionsCallsRes;
+                        if (!generateResponseState.holdPartialTokensForNextEvaluation()) {
+                            generateResponseState.waitOnPartialCharactersOrWhiteSpaceTokens();
+                            generateResponseState.detectAndHandleFunctionStartSyntax();
+                            if (generateResponseState.functionEvaluationMode !== false) {
+                                generateResponseState.canAvoidReloadingHistory = false;
+                                generateResponseState.releasePartiallyFreeTokensBeforeFunctionCallStart();
+                                const functionsCallsRes = await generateResponseState.enterFunctionCallingLoop(loadContextWindowForFunctionCallingLoop);
+                                if (functionsCallsRes != null)
+                                    return functionsCallsRes;
+                            }
+                            generateResponseState.recordStopGenerationEvaluation();
+                            generateResponseState.popStreamRegulatorFreeTokens();
+                            generateResponseState.removeFoundStartIgnoreTextsFromPendingTokens();
+                            const stopGenerationTriggerRes = generateResponseState.handleStopGenerationTrigger("model");
+                            if (stopGenerationTriggerRes != null)
+                                return stopGenerationTriggerRes;
+                            generateResponseState.spliceIgnoreStartTextDetectedTokens();
+                            generateResponseState.moveFreePendingTokensToRes();
                         }
-                        generateResponseState.recordStopGenerationEvaluation();
-                        generateResponseState.popStreamRegulatorFreeTokens();
-                        generateResponseState.removeFoundStartIgnoreTextsFromPendingTokens();
-                        const stopGenerationTriggerRes = generateResponseState.handleStopGenerationTrigger("model");
-                        if (stopGenerationTriggerRes != null)
-                            return stopGenerationTriggerRes;
-                        generateResponseState.spliceIgnoreStartTextDetectedTokens();
-                        generateResponseState.moveFreePendingTokensToRes();
                         const maxTokensTriggerRes = generateResponseState.handleMaxTokensTrigger("model");
                         if (maxTokensTriggerRes != null)
                             return maxTokensTriggerRes;
@@ -174,16 +180,20 @@ export class LlamaChat {
     }
     async loadChatAndCompleteUserMessage(history, options = {}) {
         const { initialUserPrompt = "", stopOnAbortSignal = false, onTextChunk, onToken, signal, maxTokens = Math.min(256, Math.ceil(this.context.contextSize / 2)), temperature, minP, topK, topP, seed, grammar, trimWhitespaceSuffix = defaultTrimWhitespaceSuffix, repeatPenalty = {}, tokenBias, evaluationPriority = defaultEvaluationPriority, functions, documentFunctionParams, contextShift = defaultContextShiftOptions, customStopTriggers, lastEvaluationContextWindow: { history: lastEvaluationContextWindowHistory, minimumOverlapPercentageToPreventContextShift = 0.8 } = {} } = options;
-        const lastEvaluationContextWindowHistoryItem = lastEvaluationContextWindowHistory == null
-            ? null
-            : lastEvaluationContextWindowHistory[lastEvaluationContextWindowHistory.length - 1];
-        const lastEvaluationContextWindowUserMessage = lastEvaluationContextWindowHistoryItem?.type === "user"
-            ? lastEvaluationContextWindowHistoryItem.text
-            : "";
         this.sequence.tokenPredictor?.updateInputTokens?.(this.model.tokenize((findLastModelMessageInChatHistory(history)?.response ?? [])
-            .filter((item) => typeof item === "string")
+            .map((item) => {
+            if (typeof item === "string")
+                return item;
+            else if (isChatModelResponseFunctionCall(item))
+                return null;
+            else if (isChatModelResponseSegment(item))
+                return item.text;
+            void item;
+            return null;
+        })
+            .filter((item) => item != null)
             .join(" ")));
-        const generateResponseState = new GenerateResponseState(this, this._chatWrapper, history, {
+        const generateResponseState = new GenerateResponseState(this, this._chatWrapper, mergeGeneratedResultWithChatHistory("user", history, [initialUserPrompt]), {
             onTextChunk,
             onToken,
             signal,
@@ -204,27 +214,16 @@ export class LlamaChat {
             contextShift,
             customStopTriggers,
             lastEvaluationContextWindow: {
-                history: lastEvaluationContextWindowHistory == null
-                    ? undefined
-                    : setLastUserTextInChatHistory(lastEvaluationContextWindowHistory, lastEvaluationContextWindowUserMessage + initialUserPrompt),
+                history: mergeGeneratedResultWithChatHistory("user", lastEvaluationContextWindowHistory ?? history, [initialUserPrompt]),
                 minimumOverlapPercentageToPreventContextShift
             }
         });
         return await withLock(this._chatLock, "evaluate", signal, async () => {
             try {
                 generateResponseState.ensureLastHistoryItemIsUser();
-                const getInitialUserMessage = (history) => {
-                    const lastResolvedHistoryItem = history[history.length - 1];
-                    if (lastResolvedHistoryItem?.type === "user")
-                        return lastResolvedHistoryItem.text;
-                    return "";
-                };
-                const initialUserMessage = getInitialUserMessage(generateResponseState.resolvedHistory);
-                const contextWindowInitialUserMessage = getInitialUserMessage(generateResponseState.lastContextWindowHistory);
                 while (true) {
                     generateResponseState.startTokenLoop();
-                    const { userTextSuffix } = await generateResponseState.loadContextWindow(setLastUserTextInChatHistory(generateResponseState.resolvedHistory, initialUserMessage + initialUserPrompt + this.model.detokenize(generateResponseState.res)), setLastUserTextInChatHistory(generateResponseState.lastContextWindowHistory, contextWindowInitialUserMessage + initialUserPrompt +
-                        this.model.detokenize(generateResponseState.contextWindowsRes)), true);
+                    const { userTextSuffix } = await generateResponseState.loadContextWindow(mergeGeneratedResultWithChatHistory("user", generateResponseState.resolvedHistory, generateResponseState.segmentHandler.getModelResponseSegments()), mergeGeneratedResultWithChatHistory("user", generateResponseState.lastContextWindowHistory, generateResponseState.segmentHandler.getContextWindowModelResponseSegments()), true);
                     generateResponseState.functionEvaluationMode = false;
                     generateResponseState.addStopGenerationTriggersFromChatWrapper();
                     if (userTextSuffix != null && userTextSuffix.values.length > 0)
@@ -235,7 +234,7 @@ export class LlamaChat {
                         return {
                             completion: "",
                             lastEvaluation: {
-                                contextWindow: setLastUserTextInChatHistory(generateResponseState.lastContextWindowHistory, initialUserMessage),
+                                contextWindow: mergeGeneratedResultWithChatHistory("user", generateResponseState.lastContextWindowHistory, generateResponseState.segmentHandler.getContextWindowModelResponseSegments()),
                                 contextShiftMetadata: generateResponseState.lastHistoryCompressionMetadata
                             },
                             metadata: {
@@ -245,28 +244,30 @@ export class LlamaChat {
                     }
                     await generateResponseState.createNewEvaluationIterator();
                     while (await generateResponseState.iterateEvaluation()) {
-                        generateResponseState.waitOnPartialCharactersOrWhiteSpaceTokens();
-                        generateResponseState.recordStopGenerationEvaluation();
-                        generateResponseState.popStreamRegulatorFreeTokens();
-                        const stopGenerationTriggerRes = generateResponseState.handleStopGenerationTrigger("user");
-                        if (stopGenerationTriggerRes != null)
-                            return {
-                                completion: stopGenerationTriggerRes.response,
-                                lastEvaluation: {
-                                    contextWindow: setLastUserTextInChatHistory(generateResponseState.lastContextWindowHistory, initialUserMessage),
-                                    contextShiftMetadata: stopGenerationTriggerRes.lastEvaluation.contextShiftMetadata
-                                },
-                                metadata: stopGenerationTriggerRes.metadata.stopReason === "customStopTrigger"
-                                    ? stopGenerationTriggerRes.metadata
-                                    : stopGenerationTriggerRes.metadata
-                            };
-                        generateResponseState.moveFreePendingTokensToRes(false);
+                        if (!generateResponseState.holdPartialTokensForNextEvaluation()) {
+                            generateResponseState.waitOnPartialCharactersOrWhiteSpaceTokens();
+                            generateResponseState.recordStopGenerationEvaluation();
+                            generateResponseState.popStreamRegulatorFreeTokens();
+                            const stopGenerationTriggerRes = generateResponseState.handleStopGenerationTrigger("user");
+                            if (stopGenerationTriggerRes != null)
+                                return {
+                                    completion: stopGenerationTriggerRes.response,
+                                    lastEvaluation: {
+                                        contextWindow: mergeGeneratedResultWithChatHistory("user", generateResponseState.lastContextWindowHistory, generateResponseState.segmentHandler.getContextWindowModelResponseSegments()),
+                                        contextShiftMetadata: stopGenerationTriggerRes.lastEvaluation.contextShiftMetadata
+                                    },
+                                    metadata: stopGenerationTriggerRes.metadata.stopReason === "customStopTrigger"
+                                        ? stopGenerationTriggerRes.metadata
+                                        : stopGenerationTriggerRes.metadata
+                                };
+                            generateResponseState.moveFreePendingTokensToRes(false);
+                        }
                         const maxTokensTriggerRes = generateResponseState.handleMaxTokensTrigger("user");
                         if (maxTokensTriggerRes != null)
                             return {
                                 completion: maxTokensTriggerRes.response,
                                 lastEvaluation: {
-                                    contextWindow: setLastUserTextInChatHistory(generateResponseState.lastContextWindowHistory, initialUserMessage),
+                                    contextWindow: mergeGeneratedResultWithChatHistory("user", generateResponseState.lastContextWindowHistory, generateResponseState.segmentHandler.getContextWindowModelResponseSegments()),
                                     contextShiftMetadata: maxTokensTriggerRes.lastEvaluation.contextShiftMetadata
                                 },
                                 metadata: maxTokensTriggerRes.metadata
@@ -278,7 +279,7 @@ export class LlamaChat {
                             return {
                                 completion: abortRes.response,
                                 lastEvaluation: {
-                                    contextWindow: setLastUserTextInChatHistory(generateResponseState.lastContextWindowHistory, initialUserMessage),
+                                    contextWindow: mergeGeneratedResultWithChatHistory("user", generateResponseState.lastContextWindowHistory, generateResponseState.segmentHandler.getContextWindowModelResponseSegments()),
                                     contextShiftMetadata: abortRes.lastEvaluation.contextShiftMetadata
                                 },
                                 metadata: abortRes.metadata
@@ -303,11 +304,18 @@ function removeRawFromHistoryItem(historyItem) {
         newHistoryItem.response = newHistoryItem.response.map((item) => {
             if (typeof item === "string")
                 return item;
-            else
+            else if (isChatModelResponseFunctionCall(item))
                 return {
                     ...item,
                     rawCall: undefined
                 };
+            else if (isChatModelResponseSegment(item))
+                return {
+                    ...item,
+                    raw: undefined
+                };
+            void item;
+            return item;
         });
         return newHistoryItem;
     }
@@ -370,42 +378,17 @@ async function compressHistoryToFitContextSize({ history, contextShiftSize, cont
         metadata
     };
 }
-function getLastTextModelResponseFromChatHistory(chatHistory) {
-    if (chatHistory.length === 0 || chatHistory[chatHistory.length - 1].type !== "model")
-        return "";
-    const lastModelResponseItem = chatHistory[chatHistory.length - 1];
-    const modelResponse = lastModelResponseItem.response;
-    if (modelResponse.length > 0 && typeof modelResponse[modelResponse.length - 1] === "string")
-        return modelResponse[modelResponse.length - 1];
-    return "";
+function getLastModelMessageFullResponseFromChatHistory(chatHistory) {
+    const lastModelResponseItem = chatHistory.at(-1);
+    if (lastModelResponseItem == null || lastModelResponseItem.type !== "model")
+        return [];
+    return lastModelResponseItem.response;
 }
 function getLastUserTextFromChatHistory(chatHistory) {
     if (chatHistory.length === 0 || chatHistory[chatHistory.length - 1].type !== "user")
         return "";
     return chatHistory[chatHistory.length - 1].text;
 }
-function setLastModelTextResponseInChatHistory(chatHistory, textResponse) {
-    const newChatHistory = chatHistory.slice();
-    if (newChatHistory.length === 0 || newChatHistory[newChatHistory.length - 1].type !== "model")
-        newChatHistory.push({
-            type: "model",
-            response: []
-        });
-    const lastModelResponseItem = newChatHistory[newChatHistory.length - 1];
-    const newLastModelResponseItem = { ...lastModelResponseItem };
-    newChatHistory[newChatHistory.length - 1] = newLastModelResponseItem;
-    const modelResponse = newLastModelResponseItem.response.slice();
-    newLastModelResponseItem.response = modelResponse;
-    if (modelResponse.length > 0 && typeof modelResponse[modelResponse.length - 1] === "string") {
-        if (textResponse === "")
-            modelResponse.pop();
-        else
-            modelResponse[modelResponse.length - 1] = textResponse;
-    }
-    else if (textResponse !== "")
-        modelResponse.push(textResponse);
-    return newChatHistory;
-}
 function setLastUserTextInChatHistory(chatHistory, userText) {
     const newChatHistory = chatHistory.slice();
     if (newChatHistory.length === 0 || newChatHistory[newChatHistory.length - 1].type !== "user")
@@ -419,11 +402,73 @@ function setLastUserTextInChatHistory(chatHistory, userText) {
     newLastUserItem.text = userText;
     return newChatHistory;
 }
-function setLastTextInChatHistory(itemType, chatHistory, text) {
-    if (itemType === "user")
-        return setLastUserTextInChatHistory(chatHistory, text);
-    else
-        return setLastModelTextResponseInChatHistory(chatHistory, text);
+function mergeGeneratedResultWithChatHistory(itemType, chatHistory, generatedResult) {
+    if (generatedResult.length === 0 || (generatedResult.length === 1 && generatedResult[0] === ""))
+        return chatHistory;
+    const newChatHistory = chatHistory.slice();
+    if (itemType === "user") {
+        let lastUserItem = newChatHistory.at(-1);
+        if (lastUserItem?.type !== "user") {
+            lastUserItem = {
+                type: "user",
+                text: ""
+            };
+            newChatHistory.push(lastUserItem);
+        }
+        const newLastUserItem = { ...lastUserItem };
+        newChatHistory[newChatHistory.length - 1] = newLastUserItem;
+        newLastUserItem.text += generatedResult
+            .map((item) => {
+            if (typeof item === "string")
+                return item;
+            return item.text;
+        })
+            .join("");
+        return newChatHistory;
+    }
+    else {
+        let lastModelItem = newChatHistory.at(-1);
+        if (lastModelItem?.type !== "model") {
+            lastModelItem = {
+                type: "model",
+                response: []
+            };
+            newChatHistory.push(lastModelItem);
+        }
+        const newLastModelItem = { ...lastModelItem };
+        newChatHistory[newChatHistory.length - 1] = newLastModelItem;
+        const modelResponse = newLastModelItem.response.slice();
+        newLastModelItem.response = modelResponse;
+        const firstGeneratedResultItem = generatedResult[0];
+        if (firstGeneratedResultItem == null)
+            return newChatHistory;
+        const lastModelResponseItem = modelResponse.at(-1);
+        if (typeof firstGeneratedResultItem === "string" && typeof lastModelResponseItem === "string") {
+            modelResponse[modelResponse.length - 1] = lastModelResponseItem + firstGeneratedResultItem;
+        }
+        else if (typeof firstGeneratedResultItem !== "string" && isChatModelResponseSegment(firstGeneratedResultItem) &&
+            typeof lastModelResponseItem !== "string" && isChatModelResponseSegment(lastModelResponseItem) &&
+            !lastModelResponseItem.ended && lastModelResponseItem.segmentType === firstGeneratedResultItem.segmentType) {
+            modelResponse[modelResponse.length - 1] = {
+                ...lastModelResponseItem,
+                ...firstGeneratedResultItem,
+                text: lastModelResponseItem.text + firstGeneratedResultItem.text,
+                ended: firstGeneratedResultItem.ended,
+                raw: (lastModelResponseItem.raw != null && firstGeneratedResultItem.raw != null)
+                    ? LlamaText([
+                        LlamaText.fromJSON(lastModelResponseItem.raw),
+                        LlamaText.fromJSON(firstGeneratedResultItem.raw)
+                    ]).toJSON()
+                    : undefined,
+                startTime: lastModelResponseItem.startTime,
+                endTime: firstGeneratedResultItem.endTime
+            };
+        }
+        else
+            modelResponse.push(firstGeneratedResultItem);
+        pushAll(modelResponse, generatedResult.slice(1));
+        return newChatHistory;
+    }
 }
 function findLastUserMessageInChatHistory(chatHistory) {
     for (let i = chatHistory.length - 1; i >= 0; i--) {
@@ -486,6 +531,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
         throw new DisposedError();
     const model = sequence.model;
     const context = sequence.context;
+    let removeRawFromHistory = false;
     if (isFirstEvaluation && lastEvaluationContextWindowHistory != null && sequence.isLoadedToMemory) {
         const newContextWindow = lastEvaluationContextWindowHistory.slice();
         if (endWithUserText) {
@@ -514,7 +560,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
                     history: newContextWindow,
                     stopGenerationTriggers,
                     tokens,
-                    newResolvedHistory: resolvedHistory,
+                    removeRawFromHistory,
                     newHistoryCompressionMetadata: lastHistoryCompressionMetadata,
                     ignoreStartText: ignoreStartText ?? [],
                     functionCallInitiallyEngaged: functionCall?.initiallyEngaged ?? false,
@@ -523,9 +569,10 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
                 };
         }
     }
-    resolvedHistory = sequence.isLoadedToMemory
-        ? resolvedHistory.slice()
-        : resolvedHistory.map(removeRawFromHistoryItem);
+    removeRawFromHistory = !sequence.isLoadedToMemory;
+    resolvedHistory = removeRawFromHistory
+        ? resolvedHistory.map(removeRawFromHistoryItem)
+        : resolvedHistory.slice();
     if (resolvedContextShift.lastEvaluationMetadata != null) {
         const contextShiftSize = resolvedContextShift.size instanceof Function
             ? await resolvedContextShift.size(sequence)
@@ -550,7 +597,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
             history: compressedHistory,
             stopGenerationTriggers,
             tokens: contextText.tokenize(model.tokenizer),
-            newResolvedHistory: resolvedHistory,
+            removeRawFromHistory,
             newHistoryCompressionMetadata: metadata,
             ignoreStartText: ignoreStartText ?? [],
             functionCallInitiallyEngaged: functionCall?.initiallyEngaged ?? false,
@@ -570,7 +617,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
                 history: resolvedHistory,
                 stopGenerationTriggers,
                 tokens,
-                newResolvedHistory: resolvedHistory,
+                removeRawFromHistory,
                 newHistoryCompressionMetadata: lastHistoryCompressionMetadata,
                 ignoreStartText: ignoreStartText ?? [],
                 functionCallInitiallyEngaged: functionCall?.initiallyEngaged ?? false,
@@ -601,7 +648,7 @@ async function getContextWindow({ resolvedHistory, resolvedContextShift, lastHis
         history: compressedHistory,
         stopGenerationTriggers,
         tokens: contextText.tokenize(model.tokenizer),
-        newResolvedHistory: resolvedHistory,
+        removeRawFromHistory,
         newHistoryCompressionMetadata: metadata,
         ignoreStartText: ignoreStartText ?? [],
         functionCallInitiallyEngaged: functionCall?.initiallyEngaged ?? false,
@@ -615,6 +662,7 @@ class GenerateResponseState {
     history;
     onTextChunk;
     onToken;
+    onResponseChunk;
     signal;
     stopOnAbortSignal;
     maxTokens;
@@ -638,7 +686,6 @@ class GenerateResponseState {
     repeatPenaltyEnabled;
     resolvedContextShift;
     resolvedRepeatPenalty;
-    lastModelResponse;
     grammarEvaluationState;
     functionNameGrammar;
     functionsGrammar;
@@ -651,10 +698,13 @@ class GenerateResponseState {
     ignoreStartTextDetector = new StopGenerationDetector();
     locksToReleaseOnValidGeneration = [];
     resolvedHistory;
+    noRawInResolvedHistory;
     res = [];
     pendingTokens = [];
     ignoredStartTextTokens = [];
     resFunctionCalls = [];
+    segmentHandler;
+    pendingPartialTokens = [];
     functionEvaluationMode = false;
     currentFunctionCallPreviousText = LlamaText([]);
     currentFunctionCallCurrentPartTokens = [];
@@ -678,8 +728,6 @@ class GenerateResponseState {
     disengageInitiallyEngagedFunctionCall = [];
     userTextSuffix = undefined;
     tokens = [];
-    contextWindowLastModelResponse = "";
-    contextWindowsRes = [];
     // token evaluation loop
     evaluationIterator;
     currentIteration;
@@ -688,12 +736,13 @@ class GenerateResponseState {
     currentTokens = [];
     currentText = "";
     currentQueuedTokenRelease;
-    constructor(llamaChat, chatWrapper, history, { onTextChunk, onToken, signal, stopOnAbortSignal = false, maxTokens, temperature, minP, topK, topP, seed, grammar, trimWhitespaceSuffix = defaultTrimWhitespaceSuffix, repeatPenalty = {}, tokenBias, evaluationPriority = defaultEvaluationPriority, functions, onFunctionCall, documentFunctionParams, maxParallelFunctionCalls, contextShift = defaultContextShiftOptions, customStopTriggers, lastEvaluationContextWindow: { history: lastEvaluationContextWindowHistory, minimumOverlapPercentageToPreventContextShift = 0.5 } = {} } = {}) {
+    constructor(llamaChat, chatWrapper, history, { onTextChunk, onToken, onResponseChunk, signal, stopOnAbortSignal = false, maxTokens, temperature, minP, topK, topP, seed, grammar, trimWhitespaceSuffix = defaultTrimWhitespaceSuffix, repeatPenalty = {}, tokenBias, evaluationPriority = defaultEvaluationPriority, functions, onFunctionCall, documentFunctionParams, maxParallelFunctionCalls, contextShift = defaultContextShiftOptions, customStopTriggers, lastEvaluationContextWindow: { history: lastEvaluationContextWindowHistory, minimumOverlapPercentageToPreventContextShift = 0.5 } = {} } = {}) {
         this.llamaChat = llamaChat;
         this.chatWrapper = chatWrapper;
         this.history = history;
         this.onTextChunk = safeEventCallback(onTextChunk);
         this.onToken = safeEventCallback(onToken);
+        this.onResponseChunk = safeEventCallback(onResponseChunk);
         this.signal = signal;
         this.stopOnAbortSignal = stopOnAbortSignal;
         this.maxTokens = maxTokens;
@@ -718,9 +767,10 @@ class GenerateResponseState {
             throw this.signal.reason;
         if (this.llamaChat.disposed)
             throw new DisposedError();
-        this.resolvedHistory = this.llamaChat.sequence.isLoadedToMemory
-            ? this.history.slice()
-            : this.history.map(removeRawFromHistoryItem);
+        this.noRawInResolvedHistory = !this.llamaChat.sequence.isLoadedToMemory;
+        this.resolvedHistory = this.noRawInResolvedHistory
+            ? this.history.map(removeRawFromHistoryItem)
+            : this.history.slice();
         this.resolvedContextShift = {
             ...defaultContextShiftOptions,
             ...removeNullFields(this.contextShift)
@@ -731,7 +781,6 @@ class GenerateResponseState {
                 ...(repeatPenalty ?? {}),
                 lastTokens: repeatPenalty?.lastTokens ?? defaultRepeatPenaltyLastTokens
             };
-        this.lastModelResponse = getLastTextModelResponseFromChatHistory(this.resolvedHistory);
         this.repeatPenaltyEnabled = this.resolvedRepeatPenalty.lastTokens > 0;
         this.grammarEvaluationState = this.grammar != null
             ? new LlamaGrammarEvaluationState({ model: this.llamaChat.model, grammar: this.grammar })
@@ -742,7 +791,7 @@ class GenerateResponseState {
         this.functionsGrammar = undefined;
         this.functionsEvaluationState = undefined;
         this.lastContextWindowHistory = lastEvaluationContextWindowHistory ?? this.resolvedHistory;
-        this.lastHistoryCompressionMetadata = this.resolvedContextShift;
+        this.lastHistoryCompressionMetadata = this.resolvedContextShift.lastEvaluationMetadata;
         if (this.customStopTriggers != null)
             StopGenerationDetector.resolveStopTriggers(this.customStopTriggers, this.llamaChat.model.tokenizer)
                 .map((stopTrigger) => this.customStopGenerationTriggersDetector.addStopTrigger(stopTrigger));
@@ -754,6 +803,22 @@ class GenerateResponseState {
                 this.chatWrapper.settings.functions?.parallelism?.call?.sectionPrefix ?? "",
                 this.chatWrapper.settings.functions.call.prefix
             ]), this.llamaChat.model.tokenizer));
+        const segmentDefinitions = new Map();
+        for (const segmentType of allSegmentTypes) {
+            const segmentDefinition = getChatWrapperSegmentDefinition(this.chatWrapper.settings, segmentType);
+            if (segmentDefinition != null)
+                segmentDefinitions.set(segmentType, segmentDefinition);
+        }
+        this.segmentHandler = new SegmentHandler({
+            model: this.llamaChat.model,
+            onTextChunk: this.onTextChunk,
+            onToken: this.onToken,
+            onResponseChunk: this.onResponseChunk,
+            previousTokens: this.getLastTokens(),
+            closeAllSegments: this.chatWrapper.settings.segments?.closeAllSegments,
+            segmentDefinitions,
+            initialSegmentStack: SegmentHandler.getStackFromModelResponse(getLastModelMessageFullResponseFromChatHistory(this.resolvedHistory))
+        });
         this.getPenaltyTokens = this.getPenaltyTokens.bind(this);
     }
     async dispose() {
@@ -763,19 +828,47 @@ class GenerateResponseState {
         await this.dispose();
     }
     ensureLastHistoryItemIsModel() {
-        if (this.resolvedHistory.length === 0 || this.resolvedHistory[this.resolvedHistory.length - 1].type !== "model")
+        if (this.resolvedHistory.at(-1)?.type !== "model")
             this.resolvedHistory.push({
                 type: "model",
                 response: []
             });
     }
     ensureLastHistoryItemIsUser() {
-        if (this.resolvedHistory.length === 0 || this.resolvedHistory[this.resolvedHistory.length - 1].type !== "user")
+        if (this.resolvedHistory.at(-1)?.type !== "user")
             this.resolvedHistory.push({
                 type: "user",
                 text: ""
             });
     }
+    ensureReopenedThoughtSegmentAfterFunctionCallsIfNeeded() {
+        if (this.chatWrapper.settings.segments?.thought?.reopenAfterFunctionCalls !== true)
+            return;
+        const lastModelResponseItem = this.resolvedHistory.at(-1);
+        if (lastModelResponseItem == null || lastModelResponseItem.type !== "model")
+            return;
+        const lastResponse = lastModelResponseItem.response.at(-1);
+        if (lastResponse == null)
+            return;
+        const lastResponseIsFunctionCall = typeof lastResponse !== "string" && lastResponse.type === "functionCall";
+        if (!lastResponseIsFunctionCall)
+            return;
+        const currentResponseSegmentsStack = SegmentHandler.getStackFromModelResponse(lastModelResponseItem.response);
+        if (currentResponseSegmentsStack.includes("thought"))
+            return;
+        const hadThoughtSegments = this.resolvedHistory.some((chatItem) => {
+            if (chatItem.type !== "model")
+                return false;
+            return chatItem.response.some((responseItem) => {
+                if (typeof responseItem === "string")
+                    return false;
+                return responseItem.type === "segment" && responseItem.segmentType === "thought";
+            });
+        });
+        if (!hadThoughtSegments)
+            return;
+        this.segmentHandler.openSegment("thought");
+    }
     ensureNotAborted() {
         if (this.signal?.aborted && (!this.stopOnAbortSignal || this.res.length === 0))
             throw this.signal.reason;
@@ -784,7 +877,7 @@ class GenerateResponseState {
     }
     getPenaltyTokens() {
         if (this.llamaChat.disposed)
-            throw new DisposedError();
+            return [];
         let punishTokens = this.res.slice(-this.resolvedRepeatPenalty.lastTokens);
         if (this.resolvedRepeatPenalty.punishTokensFilter != null)
             punishTokens = this.resolvedRepeatPenalty.punishTokensFilter(punishTokens);
@@ -796,24 +889,10 @@ class GenerateResponseState {
         return punishTokens;
     }
     getResolvedHistoryWithCurrentModelResponse() {
-        if (this.res.length === 0)
-            return this.resolvedHistory;
-        let modelResponse = this.llamaChat.model.detokenize(this.res);
-        if (this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix)
-            modelResponse = modelResponse.trimEnd();
-        if (modelResponse === "")
-            return this.resolvedHistory;
-        return setLastModelTextResponseInChatHistory(this.resolvedHistory, this.lastModelResponse + modelResponse);
+        return mergeGeneratedResultWithChatHistory("model", this.resolvedHistory, this.segmentHandler.getModelResponseSegments());
     }
     getContextWindowsHistoryWithCurrentModelResponse() {
-        if (this.contextWindowsRes.length === 0)
-            return this.lastContextWindowHistory;
-        let modelResponse = this.llamaChat.model.detokenize(this.contextWindowsRes);
-        if (this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix)
-            modelResponse = modelResponse.trimEnd();
-        if (modelResponse === "")
-            return this.lastContextWindowHistory;
-        return setLastModelTextResponseInChatHistory(this.lastContextWindowHistory, this.contextWindowLastModelResponse + modelResponse);
+        return mergeGeneratedResultWithChatHistory("model", this.lastContextWindowHistory, this.segmentHandler.getContextWindowModelResponseSegments());
     }
     removeFoundStartIgnoreTextsFromPendingTokens(forceRemove = false) {
         if (!this.removedStartTextToIgnore && this.res.length === 0 && this.pendingTokens.length > 0 &&
@@ -826,14 +905,26 @@ class GenerateResponseState {
                 this.contextWindowTokens,
                 this.ignoredStartTextTokens
             ]);
+            const pendingPartialTokens = [];
             for (let i = 0; i < this.pendingTokens.length; i++) {
+                const currentToken = this.pendingTokens[i];
+                const tokens = [...pendingPartialTokens, currentToken];
+                const text = this.llamaChat.model.detokenize(tokens, false, lastTokensForDetokenizer);
+                if (pendingPartialTokens.length === 0 &&
+                    text.endsWith(UNKNOWN_UNICODE_CHAR) &&
+                    !this.llamaChat.model.isSpecialToken(currentToken) &&
+                    !this.llamaChat.model.isEogToken(currentToken)) {
+                    pendingPartialTokens.length = 0;
+                    pushAll(pendingPartialTokens, tokens);
+                    continue;
+                }
                 this.ignoreStartTextDetector.recordGeneration({
-                    text: this.llamaChat.model.detokenize([this.pendingTokens[i]], false, lastTokensForDetokenizer),
-                    tokens: [this.pendingTokens[i]],
+                    text: this.llamaChat.model.detokenize(tokens, false, lastTokensForDetokenizer),
+                    tokens,
                     startNewChecks: i === 0,
                     triggerMustStartWithGeneration: true
                 });
-                lastTokensForDetokenizer.push(this.pendingTokens[i]);
+                pushAll(lastTokensForDetokenizer, tokens);
                 if (this.ignoreStartTextDetector.hasTriggeredStops) {
                     mostExhaustiveTriggeredStops = this.ignoreStartTextDetector.getTriggeredStops();
                     this.ignoreStartTextDetector.clearTriggeredStops();
@@ -902,11 +993,12 @@ class GenerateResponseState {
         const queuedChunkTokens = this.streamRegulator.getAllQueuedChunkTokens();
         const functionCallsTokens = this.getContextWindowFunctionCallsTokens();
         if (!avoidReloadingHistory || !this.canAvoidReloadingHistory || !this.llamaChat.sequence.isLoadedToMemory) {
-            const { history: contextWindowHistory, stopGenerationTriggers, tokens: contextWindowTokens, newResolvedHistory, newHistoryCompressionMetadata, ignoreStartText, functionCallInitiallyEngaged, disengageInitiallyEngagedFunctionCall, userTextSuffix } = await getContextWindow({
+            const { history: contextWindowHistory, stopGenerationTriggers, tokens: contextWindowTokens, removeRawFromHistory, newHistoryCompressionMetadata, ignoreStartText, functionCallInitiallyEngaged, disengageInitiallyEngagedFunctionCall, userTextSuffix } = await getContextWindow({
                 resolvedHistory: resolvedHistory,
                 resolvedContextShift: this.resolvedContextShift,
                 lastHistoryCompressionMetadata: this.lastHistoryCompressionMetadata,
-                pendingTokensCount: this.pendingTokens.length + queuedChunkTokens.length + functionCallsTokens.length,
+                pendingTokensCount: this.pendingTokens.length + queuedChunkTokens.length + functionCallsTokens.length +
+                    this.pendingPartialTokens.length,
                 isFirstEvaluation: this.isFirstEvaluation,
                 chatWrapper: this.chatWrapper,
                 lastEvaluationContextWindowHistory: resolvedContextWindowsHistory,
@@ -924,19 +1016,22 @@ class GenerateResponseState {
             this.functionCallInitiallyEngaged = functionCallInitiallyEngaged;
             this.disengageInitiallyEngagedFunctionCall = disengageInitiallyEngagedFunctionCall;
             this.userTextSuffix = userTextSuffix;
-            this.resolvedHistory = newResolvedHistory;
             this.lastHistoryCompressionMetadata = newHistoryCompressionMetadata;
             this.lastContextWindowHistory = contextWindowHistory;
-            this.contextWindowLastModelResponse = getLastTextModelResponseFromChatHistory(contextWindowHistory);
-            this.contextWindowsRes = [];
+            this.segmentHandler.resetContextWindow();
             this.canAvoidReloadingHistory = true;
+            if (removeRawFromHistory && !this.noRawInResolvedHistory) {
+                this.noRawInResolvedHistory = true;
+                this.resolvedHistory = this.resolvedHistory.map(removeRawFromHistoryItem);
+            }
         }
         this.tokens = [
             ...this.contextWindowTokens,
             ...this.ignoredStartTextTokens,
             ...this.pendingTokens,
             ...queuedChunkTokens,
-            ...functionCallsTokens
+            ...functionCallsTokens,
+            ...this.pendingPartialTokens
         ];
         if (avoidReloadingHistory && this.tokens.length >= this.llamaChat.sequence.context.contextSize - 1)
             return await this.loadContextWindow(resolvedHistory, resolvedContextWindowsHistory, endWithUserText, false);
@@ -1017,24 +1112,24 @@ class GenerateResponseState {
                         pushAll(prefixDetectorRecordedTokens, tokens);
                     }
                 }
-                for await (const token of this.evaluateWithContextShift(loadContextWindow)) {
+                for await (const tokens of this.evaluateWithContextShift(loadContextWindow)) {
                     const stopGenerationTriggerRes = this.handleStopGenerationTrigger("model");
                     if (stopGenerationTriggerRes != null)
                         return stopGenerationTriggerRes;
-                    this.currentFunctionCallCurrentPartTokens.push(token);
+                    pushAll(this.currentFunctionCallCurrentPartTokens, tokens);
                     this.disengageInitiallyEngagedFunctionMode.recordGeneration({
                         text: this.currentText,
                         tokens: this.currentTokens,
-                        startNewChecks: this.currentFunctionCallCurrentPartTokens.length === 1,
+                        startNewChecks: this.currentFunctionCallCurrentPartTokens.length === tokens.length,
                         triggerMustStartWithGeneration: true
                     });
                     if (prefixDetector.hasTriggeredStops)
-                        afterPrefixLeftoverTokens.push(token);
+                        pushAll(afterPrefixLeftoverTokens, tokens);
                     else {
                         prefixDetector.recordGeneration({
                             text: this.currentText,
                             tokens: this.currentTokens,
-                            startNewChecks: this.currentFunctionCallCurrentPartTokens.length === 1,
+                            startNewChecks: this.currentFunctionCallCurrentPartTokens.length === tokens.length,
                             triggerMustStartWithGeneration: true
                         });
                         pushAll(prefixDetectorRecordedTokens, this.currentTokens);
@@ -1109,8 +1204,8 @@ class GenerateResponseState {
                         }
                     }
                 }
-                for await (const token of this.evaluateWithContextShift(loadContextWindow)) {
-                    this.currentFunctionCallCurrentPartTokens.push(token);
+                for await (const tokens of this.evaluateWithContextShift(loadContextWindow)) {
+                    pushAll(this.currentFunctionCallCurrentPartTokens, tokens);
                     functionNameGenerationDoneDetector.recordGeneration({
                         text: this.currentText,
                         tokens: this.currentTokens
@@ -1154,8 +1249,8 @@ class GenerateResponseState {
                     });
                     StopGenerationDetector.resolveStopTriggers(this.functionsGrammar.stopGenerationTriggers, this.llamaChat.model.tokenizer)
                         .map((stopTrigger) => functionParamsGenerationDoneDetector.addStopTrigger(stopTrigger));
-                    for await (const token of this.evaluateWithContextShift(loadContextWindow)) {
-                        this.currentFunctionCallCurrentPartTokens.push(token);
+                    for await (const tokens of this.evaluateWithContextShift(loadContextWindow)) {
+                        pushAll(this.currentFunctionCallCurrentPartTokens, tokens);
                         functionParamsGenerationDoneDetector.recordGeneration({
                             text: this.currentText,
                             tokens: this.currentTokens
@@ -1213,8 +1308,8 @@ class GenerateResponseState {
                     LlamaText(new SpecialToken("EOT"))
                 ], this.llamaChat.model.tokenizer)
                     .map((stopTrigger) => sectionSuffixDetector.addStopTrigger(stopTrigger));
-                for await (const token of this.evaluateWithContextShift(loadContextWindow)) {
-                    this.currentFunctionCallCurrentPartTokens.push(token);
+                for await (const tokens of this.evaluateWithContextShift(loadContextWindow)) {
+                    pushAll(this.currentFunctionCallCurrentPartTokens, tokens);
                     sectionSuffixDetector.recordGeneration({
                         text: this.currentText,
                         tokens: this.currentTokens,
@@ -1258,17 +1353,17 @@ class GenerateResponseState {
     returnFunctionCallResults() {
         if (this.resFunctionCalls.length > 0) {
             this.releasePartiallyFreeTokensBeforeFunctionCallStart();
-            let modelResponse = this.llamaChat.model.detokenize(this.res);
-            let contextWindowModelResponse = this.llamaChat.model.detokenize(this.contextWindowsRes);
-            if (this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix) {
-                modelResponse = modelResponse.trimEnd();
-                contextWindowModelResponse = contextWindowModelResponse.trimEnd();
-            }
+            this.segmentHandler.onFinishedGeneration();
+            const trimWhitespaceSuffix = this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix;
+            const responseSegments = this.segmentHandler.getModelResponseSegments(trimWhitespaceSuffix);
             return {
-                response: modelResponse,
+                response: responseSegments
+                    .filter((segment) => typeof segment === "string")
+                    .join(""),
+                fullResponse: responseSegments,
                 lastEvaluation: {
-                    contextWindow: setLastTextInChatHistory("model", this.lastContextWindowHistory, this.contextWindowLastModelResponse + contextWindowModelResponse),
-                    cleanHistory: setLastTextInChatHistory("model", this.resolvedHistory, this.lastModelResponse + modelResponse),
+                    contextWindow: mergeGeneratedResultWithChatHistory("model", this.lastContextWindowHistory, this.segmentHandler.getContextWindowModelResponseSegments(trimWhitespaceSuffix)),
+                    cleanHistory: mergeGeneratedResultWithChatHistory("model", this.resolvedHistory, responseSegments),
                     contextShiftMetadata: this.lastHistoryCompressionMetadata
                 },
                 functionCalls: this.resFunctionCalls.map((functionCall) => {
@@ -1292,9 +1387,10 @@ class GenerateResponseState {
             await this.alignCurrentSequenceStateWithCurrentTokens();
             await this.createNewEvaluationIterator();
             while (await this.iterateEvaluation()) {
-                if (this.currentToken == null)
+                if (this.currentTokens.length === 0)
                     break;
-                yield this.currentToken;
+                if (!this.holdPartialTokensForNextEvaluation())
+                    yield this.currentTokens;
                 if (this.shouldAbort)
                     return;
                 if (this.updateShouldContextShift())
@@ -1367,9 +1463,14 @@ class GenerateResponseState {
         this.currentIterationReplacementToken = undefined;
         this.ensureNotAborted();
         this.generatedTokens++;
-        if (this.currentIteration != null && this.currentIteration?.done !== true) {
-            this.currentToken = this.currentIteration.value;
-            this.currentTokens = [this.currentToken];
+        if ((this.currentIteration != null && this.currentIteration?.done !== true) || this.pendingPartialTokens.length !== 0) {
+            this.currentToken = this.currentIteration?.value ?? undefined;
+            this.currentTokens = this.currentToken != null
+                ? this.pendingPartialTokens.length === 0
+                    ? [this.currentToken]
+                    : [...this.pendingPartialTokens, this.currentToken]
+                : [...this.pendingPartialTokens];
+            this.pendingPartialTokens.length = 0;
             this.currentText = this.llamaChat.model.detokenize(this.currentTokens, false, this.getLastTokens());
             if (this.functionEvaluationMode === false)
                 this.currentQueuedTokenRelease = this.streamRegulator.addChunk({
@@ -1382,6 +1483,19 @@ class GenerateResponseState {
         }
         return false;
     }
+    holdPartialTokensForNextEvaluation() {
+        if (this.pendingPartialTokens.length === 0 &&
+            this.currentText.endsWith(UNKNOWN_UNICODE_CHAR) &&
+            this.currentToken != null &&
+            !this.llamaChat.model.isSpecialToken(this.currentToken) &&
+            !this.llamaChat.model.isEogToken(this.currentToken)) {
+            this.pendingPartialTokens.length = 0;
+            pushAll(this.pendingPartialTokens, this.currentTokens);
+            this.streamRegulator.removeChunkIfLast(this.currentQueuedTokenRelease);
+            return true;
+        }
+        return false;
+    }
     waitOnPartialCharactersOrWhiteSpaceTokens() {
         if (this.currentText.endsWith(UNKNOWN_UNICODE_CHAR) || ((this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix) && this.currentText?.trim() === "") || (this.currentText === "" && this.locksToReleaseOnValidGeneration.length > 0 &&
             !this.llamaChat.model.isSpecialToken(this.currentToken))) {
@@ -1449,21 +1563,22 @@ class GenerateResponseState {
             const { firstRemainingGenerationAfterStop } = StopGenerationDetector.getFirstRemainingGenerationAfterStop(triggeredStops);
             this.removeFoundStartIgnoreTextsFromPendingTokens(true);
             this.pushPendingTokensAndCallOnToken();
-            let modelResponse = this.llamaChat.model.detokenize(this.res);
-            let contextWindowModelResponse = this.llamaChat.model.detokenize(this.contextWindowsRes);
-            if (this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix) {
-                modelResponse = modelResponse.trimEnd();
-                contextWindowModelResponse = contextWindowModelResponse.trimEnd();
-            }
+            this.segmentHandler.onFinishedGeneration();
+            const trimWhitespaceSuffix = this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix;
+            const responseSegments = this.segmentHandler.getModelResponseSegments(trimWhitespaceSuffix);
+            const response = responseSegments
+                .filter((segment) => typeof segment === "string")
+                .join("");
             const lastEvaluation = {
-                contextWindow: setLastTextInChatHistory(lastHistoryItemType, this.lastContextWindowHistory, this.contextWindowLastModelResponse + contextWindowModelResponse),
-                cleanHistory: setLastTextInChatHistory(lastHistoryItemType, this.resolvedHistory, this.lastModelResponse + modelResponse),
+                contextWindow: mergeGeneratedResultWithChatHistory(lastHistoryItemType, this.lastContextWindowHistory, this.segmentHandler.getContextWindowModelResponseSegments(trimWhitespaceSuffix)),
+                cleanHistory: mergeGeneratedResultWithChatHistory(lastHistoryItemType, this.resolvedHistory, responseSegments),
                 contextShiftMetadata: this.lastHistoryCompressionMetadata
             };
             const isEogToken = this.llamaChat.model.isEogToken(this.currentToken);
             if (isEogToken || this.stopGenerationDetector.hasTriggeredStops) {
                 return {
-                    response: modelResponse,
+                    response,
+                    fullResponse: responseSegments,
                     lastEvaluation,
                     metadata: {
                         remainingGenerationAfterStop: firstRemainingGenerationAfterStop,
@@ -1474,7 +1589,8 @@ class GenerateResponseState {
                 };
             }
             return {
-                response: modelResponse,
+                response,
+                fullResponse: responseSegments,
                 lastEvaluation,
                 metadata: {
                     remainingGenerationAfterStop: firstRemainingGenerationAfterStop,
@@ -1511,17 +1627,17 @@ class GenerateResponseState {
     }
     handleMaxTokensTrigger(lastHistoryItemType) {
         if (this.isMaxTokensTriggered()) {
-            let modelResponse = this.llamaChat.model.detokenize(this.res);
-            let contextWindowModelResponse = this.llamaChat.model.detokenize(this.contextWindowsRes);
-            if (this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix) {
-                modelResponse = modelResponse.trimEnd();
-                contextWindowModelResponse = contextWindowModelResponse.trimEnd();
-            }
+            this.segmentHandler.onFinishedGeneration();
+            const trimWhitespaceSuffix = this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix;
+            const responseSegments = this.segmentHandler.getModelResponseSegments(trimWhitespaceSuffix);
             return {
-                response: modelResponse,
+                response: responseSegments
+                    .filter((segment) => typeof segment === "string")
+                    .join(""),
+                fullResponse: responseSegments,
                 lastEvaluation: {
-                    contextWindow: setLastTextInChatHistory(lastHistoryItemType, this.lastContextWindowHistory, this.contextWindowLastModelResponse + contextWindowModelResponse),
-                    cleanHistory: setLastTextInChatHistory(lastHistoryItemType, this.resolvedHistory, this.lastModelResponse + modelResponse),
+                    contextWindow: mergeGeneratedResultWithChatHistory(lastHistoryItemType, this.lastContextWindowHistory, this.segmentHandler.getContextWindowModelResponseSegments(trimWhitespaceSuffix)),
+                    cleanHistory: mergeGeneratedResultWithChatHistory(lastHistoryItemType, this.resolvedHistory, responseSegments),
                     contextShiftMetadata: this.lastHistoryCompressionMetadata
                 },
                 metadata: {
@@ -1542,17 +1658,17 @@ class GenerateResponseState {
         if (this.shouldAbort && this.signal?.aborted && this.stopOnAbortSignal) {
             if (this.res.length === 0)
                 throw this.signal.reason;
-            let modelResponse = this.llamaChat.model.detokenize(this.res);
-            let contextWindowModelResponse = this.llamaChat.model.detokenize(this.contextWindowsRes);
-            if (this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix) {
-                modelResponse = modelResponse.trimEnd();
-                contextWindowModelResponse = contextWindowModelResponse.trimEnd();
-            }
+            this.segmentHandler.onFinishedGeneration();
+            const trimWhitespaceSuffix = this.grammar?.trimWhitespaceSuffix || this.trimWhitespaceSuffix;
+            const responseSegments = this.segmentHandler.getModelResponseSegments(trimWhitespaceSuffix);
             return {
-                response: modelResponse,
+                response: responseSegments
+                    .filter((segment) => typeof segment === "string")
+                    .join(""),
+                fullResponse: responseSegments,
                 lastEvaluation: {
-                    contextWindow: setLastTextInChatHistory(lastHistoryItemType, this.lastContextWindowHistory, this.contextWindowLastModelResponse + contextWindowModelResponse),
-                    cleanHistory: setLastTextInChatHistory(lastHistoryItemType, this.resolvedHistory, this.lastModelResponse + modelResponse),
+                    contextWindow: mergeGeneratedResultWithChatHistory(lastHistoryItemType, this.lastContextWindowHistory, this.segmentHandler.getContextWindowModelResponseSegments(trimWhitespaceSuffix)),
+                    cleanHistory: mergeGeneratedResultWithChatHistory(lastHistoryItemType, this.resolvedHistory, responseSegments),
                     contextShiftMetadata: this.lastHistoryCompressionMetadata
                 },
                 metadata: {
@@ -1565,10 +1681,8 @@ class GenerateResponseState {
     pushPendingTokensAndCallOnToken() {
         if (this.pendingTokens.length === 0)
             return;
-        this.onToken?.(this.pendingTokens.slice());
-        this.onTextChunk?.(this.llamaChat.model.detokenize(this.pendingTokens, false, this.res));
+        this.segmentHandler.processTokens(this.pendingTokens);
         pushAll(this.res, this.pendingTokens);
-        pushAll(this.contextWindowsRes, this.pendingTokens);
         this.pendingTokens.length = 0;
     }
     getLastTokens(maxTokens = maxRecentDetokenizerTokens) {
@@ -1577,8 +1691,470 @@ class GenerateResponseState {
             this.ignoredStartTextTokens,
             this.pendingTokens,
             this.streamRegulator.getLastQueuedChunkTokens(maxTokens),
-            this.getContextWindowFunctionCallsTokens()
+            this.getContextWindowFunctionCallsTokens(),
+            this.pendingPartialTokens
         ], maxTokens);
     }
 }
+class SegmentHandler {
+    model;
+    onToken;
+    onTextChunk;
+    onResponseChunk;
+    _closeAllSegmentsDetector;
+    _segmentDetectors;
+    _segmentsStack = [];
+    _segmentsStackSet = new Set();
+    _ownedSegmentsStackLength = 0;
+    _segments = [];
+    _segmentsStartTokenTrail = [];
+    _contextWindowSegments = [];
+    _contextWindowStartTokenTrail = [];
+    _initialTokensTrail;
+    _tokensTrail;
+    _streamRegulator = new TokenStreamRegulator();
+    _segmentDefinitions;
+    constructor({ model, onTextChunk, onToken, onResponseChunk, segmentDefinitions, closeAllSegments, initialSegmentStack, previousTokens }) {
+        this.model = model;
+        this.onTextChunk = onTextChunk;
+        this.onToken = onToken;
+        this.onResponseChunk = onResponseChunk;
+        this._initialTokensTrail = previousTokens.slice(-maxRecentDetokenizerTokens);
+        this._segmentsStartTokenTrail = previousTokens.slice(-maxRecentDetokenizerTokens);
+        this._tokensTrail = previousTokens.slice(-maxRecentDetokenizerTokens);
+        this._closeAllSegmentsDetector = closeAllSegments != null
+            ? new StopGenerationDetector()
+                .addStopTrigger(StopGenerationDetector.resolveLlamaTextTrigger(LlamaText(closeAllSegments), this.model.tokenizer))
+            : undefined;
+        this._segmentDetectors = new Map();
+        this._segmentsStack = initialSegmentStack;
+        this._segmentsStackSet = new Set(initialSegmentStack);
+        this._ownedSegmentsStackLength = initialSegmentStack.length;
+        this._segmentDefinitions = segmentDefinitions;
+        for (const [segment, { prefix, suffix }] of segmentDefinitions.entries()) {
+            this._segmentDetectors.set(segment, {
+                prefix: new StopGenerationDetector()
+                    .addStopTrigger(StopGenerationDetector.resolveLlamaTextTrigger(LlamaText(prefix), this.model.tokenizer)),
+                suffix: suffix != null
+                    ? new StopGenerationDetector()
+                        .addStopTrigger(StopGenerationDetector.resolveLlamaTextTrigger(LlamaText(suffix), this.model.tokenizer))
+                    : undefined
+            });
+        }
+    }
+    processTokens(tokens) {
+        if (tokens.length === 0)
+            return;
+        let pendingTokens = [];
+        for (const token of tokens) {
+            pendingTokens.push(token);
+            const currentText = this.model.detokenize(pendingTokens, false, this._tokensTrail);
+            if (currentText.endsWith(UNKNOWN_UNICODE_CHAR))
+                continue;
+            pushAll(this._tokensTrail, pendingTokens);
+            this._processTokens(pendingTokens, currentText);
+            pendingTokens = [];
+        }
+    }
+    onFinishedGeneration() {
+        this._clearDetectors();
+        this._pushCurrentTokens(this._streamRegulator.popFreeChunkTokens());
+    }
+    resetContextWindow() {
+        this._contextWindowSegments.length = 0;
+        this._contextWindowStartTokenTrail.length = 0;
+        pushAll(this._contextWindowStartTokenTrail, this._getTokenTrailFromResult());
+    }
+    openSegment(type) {
+        const now = Date.now();
+        this._segmentsStack.push(type);
+        this._segmentsStackSet.add(type);
+        this._segments.push({ type, tokens: [], ended: false, start: true, startTime: now });
+        this._contextWindowSegments.push({ type, tokens: [], ended: false, start: true, startTime: now });
+        this.onResponseChunk?.({
+            type: "segment",
+            segmentType: type,
+            tokens: [],
+            text: "",
+            segmentStartTime: new Date(now)
+        });
+    }
+    _processTokens(tokens, text) {
+        const queuedTokenRelease = this._streamRegulator.addChunk({
+            tokens,
+            text
+        });
+        const currentType = this._segmentsStack.at(-1);
+        const handleDetector = (stopDetector, action, type) => {
+            if (stopDetector == null)
+                return false;
+            stopDetector.recordGeneration({
+                text,
+                tokens,
+                queuedTokenRelease
+            });
+            if (stopDetector.hasTriggeredStops) {
+                const [leftTokens, leftText] = this._handleTriggeredStopDetector(stopDetector);
+                if (action === "pop")
+                    this._closeSegment(type);
+                else if (action === "push") {
+                    this.openSegment(type);
+                }
+                else if (action === "reset") {
+                    const now = Date.now();
+                    while (this._segmentsStack.length > 0) {
+                        const segmentType = this._segmentsStack.pop();
+                        this._segmentsStackSet.delete(segmentType);
+                        const lastSegment = this._segments.at(-1);
+                        if (lastSegment != null && !(lastSegment instanceof Array) && lastSegment.type === segmentType) {
+                            lastSegment.ended = true;
+                            lastSegment.endTime = now;
+                            this.onResponseChunk?.({
+                                type: "segment",
+                                segmentType: segmentType,
+                                tokens: [],
+                                text: "",
+                                segmentStartTime: undefined,
+                                segmentEndTime: new Date(now)
+                            });
+                        }
+                        else {
+                            this._segments.push({ type: segmentType, tokens: [], ended: true, start: false, endTime: now });
+                            this.onResponseChunk?.({
+                                type: "segment",
+                                segmentType: segmentType,
+                                tokens: [],
+                                text: "",
+                                segmentStartTime: undefined,
+                                segmentEndTime: new Date(now)
+                            });
+                        }
+                        const lastContextWindowSegment = this._contextWindowSegments.at(-1);
+                        if (lastContextWindowSegment != null && !(lastContextWindowSegment instanceof Array) &&
+                            lastContextWindowSegment.type === segmentType)
+                            lastContextWindowSegment.ended = true;
+                        else
+                            this._contextWindowSegments.push({ type: segmentType, tokens: [], ended: true, start: false, endTime: now });
+                    }
+                    this._ownedSegmentsStackLength = 0;
+                }
+                if (leftTokens.length > 0)
+                    this._processTokens(leftTokens, leftText);
+                return true;
+            }
+            return false;
+        };
+        if (currentType != null) {
+            if (handleDetector(this._closeAllSegmentsDetector, "reset", currentType))
+                return;
+            if (handleDetector(this._segmentDetectors.get(currentType)?.suffix, "pop", currentType))
+                return;
+        }
+        else
+            this._closeAllSegmentsDetector?.clearInProgressStops();
+        for (const [type, { prefix, suffix }] of this._segmentDetectors.entries()) {
+            if (!this._segmentsStackSet.has(type)) {
+                if (handleDetector(prefix, "push", type))
+                    return;
+            }
+            else
+                prefix.clearInProgressStops();
+            if (this._segmentsStackSet.has(type)) {
+                // `currentType` suffix is already handled above
+                if (type === currentType && handleDetector(suffix, "pop", type))
+                    return;
+            }
+            else
+                suffix?.clearInProgressStops();
+        }
+        this._pushCurrentTokens(this._streamRegulator.popFreeChunkTokens());
+    }
+    _handleTriggeredStopDetector(stopDetector) {
+        this._clearDetectors(stopDetector);
+        stopDetector.clearInProgressStops();
+        const triggeredStops = stopDetector.getTriggeredStops();
+        const freeTokens = this._streamRegulator.popFreeChunkTokens();
+        const partiallyFreeTokens = this._streamRegulator.getPartiallyFreeChunk(this.model.tokenizer);
+        const queuedTokensBeforeStopTrigger = getQueuedTokensBeforeStopTrigger(triggeredStops, partiallyFreeTokens, this.model.tokenizer);
+        const { firstRemainingGenerationAfterStop } = StopGenerationDetector.getFirstRemainingGenerationAfterStop(triggeredStops);
+        const remainingTokens = typeof firstRemainingGenerationAfterStop === "string"
+            ? firstRemainingGenerationAfterStop === ""
+                ? []
+                : this.model.tokenize(firstRemainingGenerationAfterStop, false)
+            : (firstRemainingGenerationAfterStop ?? []);
+        const remainingText = typeof firstRemainingGenerationAfterStop === "string"
+            ? firstRemainingGenerationAfterStop
+            : this.model.detokenize(remainingTokens, false, queuedTokensBeforeStopTrigger.length === 0
+                ? this._getTokenTrailFromResult()
+                : queuedTokensBeforeStopTrigger);
+        this._pushCurrentTokens([...freeTokens, ...queuedTokensBeforeStopTrigger]);
+        stopDetector.clearTriggeredStops();
+        this._streamRegulator.reset();
+        return [remainingTokens, remainingText];
+    }
+    _closeSegment(type) {
+        if (type == null)
+            return;
+        const lastSegment = this._segments.at(-1);
+        const now = Date.now();
+        if (lastSegment != null && !(lastSegment instanceof Array) && lastSegment.type === type && this._segmentsStack.at(-1) === type) {
+            if (lastSegment.ended !== true) {
+                lastSegment.ended = true;
+                lastSegment.endTime = now;
+                this.onResponseChunk?.({
+                    type: "segment",
+                    segmentType: type,
+                    tokens: [],
+                    text: "",
+                    segmentStartTime: undefined,
+                    segmentEndTime: new Date(now)
+                });
+            }
+            this._segmentsStackSet.delete(this._segmentsStack.pop());
+            if (this._segmentsStack.length < this._ownedSegmentsStackLength)
+                this._ownedSegmentsStackLength = this._segmentsStack.length;
+            const lastContextWindowSegment = this._contextWindowSegments.at(-1);
+            if (lastContextWindowSegment != null && !(lastContextWindowSegment instanceof Array) &&
+                lastContextWindowSegment.type === type && this._segmentsStack.at(-1) === type) {
+                if (lastContextWindowSegment.ended !== true) {
+                    lastContextWindowSegment.ended = true;
+                    lastContextWindowSegment.endTime = now;
+                }
+            }
+            else
+                this._contextWindowSegments.push({ type, tokens: [], ended: true, start: false, endTime: now });
+            return;
+        }
+        const typeIndex = this._segmentsStack.lastIndexOf(type);
+        if (typeIndex < 0)
+            return;
+        for (let i = this._segmentsStack.length - 1; i >= typeIndex; i--) {
+            const segmentType = this._segmentsStack.pop();
+            this._segmentsStackSet.delete(segmentType);
+            if (this._segmentsStack.length < this._ownedSegmentsStackLength)
+                this._ownedSegmentsStackLength = this._segmentsStack.length;
+            this._segments.push({ type: segmentType, tokens: [], ended: true, start: false, endTime: now });
+            this._contextWindowSegments.push({ type: segmentType, tokens: [], ended: true, start: false, endTime: now });
+            this.onResponseChunk?.({
+                type: "segment",
+                segmentType: segmentType,
+                tokens: [],
+                text: "",
+                segmentStartTime: undefined,
+                segmentEndTime: new Date(now)
+            });
+        }
+    }
+    _clearDetectors(skipDetector) {
+        if (this._closeAllSegmentsDetector !== skipDetector) {
+            this._closeAllSegmentsDetector?.clearInProgressStops();
+            this._closeAllSegmentsDetector?.clearTriggeredStops();
+        }
+        for (const { prefix, suffix } of this._segmentDetectors.values()) {
+            if (prefix !== skipDetector) {
+                prefix.clearInProgressStops();
+                prefix.clearTriggeredStops();
+            }
+            if (suffix !== skipDetector) {
+                suffix?.clearInProgressStops();
+                suffix?.clearTriggeredStops();
+            }
+        }
+    }
+    _pushCurrentTokens(tokens) {
+        const lastSegment = this._segments.at(-1);
+        const lastContextWindowSegment = this._contextWindowSegments.at(-1);
+        const type = this._segmentsStack.at(-1);
+        if (type == null) {
+            if (lastSegment == null) {
+                const text = (this.onResponseChunk != null || this.onTextChunk != null)
+                    ? this.model.detokenize(tokens, false, this._getTokenTrailFromResult())
+                    : "";
+                this._segments.push(tokens);
+                this.onToken?.(tokens.slice());
+                this.onTextChunk?.(text);
+                this.onResponseChunk?.({ type: undefined, segmentType: undefined, tokens: tokens.slice(), text });
+            }
+            else {
+                if (lastSegment instanceof Array) {
+                    const text = (this.onResponseChunk != null || this.onTextChunk != null)
+                        ? this.model.detokenize(tokens, false, this._getTokenTrailFromResult())
+                        : "";
+                    pushAll(lastSegment, tokens);
+                    this.onToken?.(tokens);
+                    this.onTextChunk?.(text);
+                    this.onResponseChunk?.({ type: undefined, segmentType: undefined, tokens, text });
+                }
+                else
+                    this._segments.push(tokens);
+            }
+            if (lastContextWindowSegment == null)
+                this._contextWindowSegments.push(tokens.slice());
+            else {
+                if (lastContextWindowSegment instanceof Array)
+                    pushAll(lastContextWindowSegment, tokens);
+                else
+                    this._contextWindowSegments.push(tokens.slice());
+            }
+        }
+        else {
+            const now = Date.now();
+            if (lastSegment == null) {
+                const text = this.onResponseChunk != null
+                    ? this.model.detokenize(tokens, false, this._getTokenTrailFromResult())
+                    : "";
+                this._segments.push({
+                    type,
+                    tokens,
+                    ended: false,
+                    start: this._segmentsStack.length > this._ownedSegmentsStackLength,
+                    startTime: now
+                });
+                this.onResponseChunk?.({
+                    type: "segment",
+                    segmentType: type,
+                    tokens: tokens.slice(),
+                    text,
+                    segmentStartTime: new Date(now)
+                });
+            }
+            else {
+                const text = this.onResponseChunk != null
+                    ? this.model.detokenize(tokens, false, this._getTokenTrailFromResult())
+                    : "";
+                if (lastSegment instanceof Array || lastSegment.type !== type) {
+                    this._segments.push({
+                        type,
+                        tokens,
+                        ended: false,
+                        start: this._segmentsStack.length > this._ownedSegmentsStackLength,
+                        startTime: now
+                    });
+                    this.onResponseChunk?.({
+                        type: "segment",
+                        segmentType: type,
+                        tokens: tokens.slice(),
+                        text,
+                        segmentStartTime: new Date(now)
+                    });
+                }
+                else {
+                    pushAll(lastSegment.tokens, tokens);
+                    this.onResponseChunk?.({
+                        type: "segment",
+                        segmentType: type,
+                        tokens: tokens.slice(),
+                        text,
+                        segmentStartTime: undefined
+                    });
+                }
+            }
+            if (lastContextWindowSegment == null)
+                this._contextWindowSegments.push({
+                    type,
+                    tokens: tokens.slice(),
+                    ended: false,
+                    start: this._segmentsStack.length > this._ownedSegmentsStackLength,
+                    startTime: now
+                });
+            else {
+                if (lastContextWindowSegment instanceof Array || lastContextWindowSegment.type !== type)
+                    this._contextWindowSegments.push({
+                        type,
+                        tokens: tokens.slice(),
+                        ended: false,
+                        start: this._segmentsStack.length > this._ownedSegmentsStackLength,
+                        startTime: now
+                    });
+                else
+                    pushAll(lastContextWindowSegment.tokens, tokens);
+            }
+        }
+    }
+    _getTokenTrailFromResult() {
+        const res = [];
+        for (let i = this._segments.length - 1; i >= 0; i--) {
+            const segment = this._segments[i];
+            const segmentTokens = segment instanceof Array
+                ? segment
+                : segment.tokens;
+            for (let j = segmentTokens.length - 1; j >= 0; j--) {
+                res.unshift(segmentTokens[j]);
+                if (res.length >= maxRecentDetokenizerTokens)
+                    return res;
+            }
+        }
+        for (let i = this._initialTokensTrail.length - 1; i >= 0; i--) {
+            res.unshift(this._initialTokensTrail[i]);
+            if (res.length >= maxRecentDetokenizerTokens)
+                return res;
+        }
+        return res;
+    }
+    getModelResponseSegments(trimWhitespaceSuffix = false) {
+        return this._getModelResponseForSegments(this._segments, this._segmentsStartTokenTrail, trimWhitespaceSuffix);
+    }
+    getContextWindowModelResponseSegments(trimWhitespaceSuffix = false) {
+        return this._getModelResponseForSegments(this._contextWindowSegments, this._contextWindowStartTokenTrail, trimWhitespaceSuffix);
+    }
+    _getModelResponseForSegments(rawSegments, recentTokens, trimWhitespaceSuffix) {
+        let tokenTrail = resolveLastTokens([recentTokens]);
+        return rawSegments.map((rawSegment, index) => {
+            const isLast = index === rawSegments.length - 1;
+            if (rawSegment instanceof Array) {
+                let text = this.model.detokenize(rawSegment, false, tokenTrail);
+                if (isLast && trimWhitespaceSuffix)
+                    text = text.trimEnd();
+                tokenTrail = resolveLastTokens([tokenTrail, rawSegment]);
+                return text;
+            }
+            let text = this.model.detokenize(rawSegment.tokens, false, tokenTrail);
+            if (isLast && rawSegment.ended && trimWhitespaceSuffix)
+                text = text.trimEnd();
+            tokenTrail = resolveLastTokens([tokenTrail, rawSegment.tokens]);
+            const segmentDefinition = this._segmentDefinitions.get(rawSegment.type);
+            return {
+                type: "segment",
+                segmentType: rawSegment.type,
+                text,
+                ended: rawSegment.ended,
+                raw: segmentDefinition == null
+                    ? LlamaText([text]).toJSON()
+                    : LlamaText([
+                        rawSegment.start
+                            ? segmentDefinition.prefix
+                            : "",
+                        text,
+                        rawSegment.ended
+                            ? (segmentDefinition.suffix ?? "")
+                            : ""
+                    ]).toJSON(),
+                startTime: rawSegment.startTime != null
+                    ? new Date(rawSegment.startTime).toISOString()
+                    : undefined,
+                endTime: rawSegment.endTime != null
+                    ? new Date(rawSegment.endTime).toISOString()
+                    : undefined
+            };
+        });
+    }
+    static getStackFromModelResponse(modelResponse) {
+        const stack = [];
+        const stackSet = new Set();
+        for (const item of modelResponse) {
+            if (typeof item === "string" || isChatModelResponseFunctionCall(item))
+                continue;
+            void item.type;
+            if (item.ended && stack.at(-1) === item.segmentType) {
+                stack.pop();
+                stackSet.delete(item.segmentType);
+            }
+            else if (!item.ended && !stackSet.has(item.segmentType)) {
+                stack.push(item.segmentType);
+                stackSet.add(item.segmentType);
+            }
+        }
+        return stack;
+    }
+}
 //# sourceMappingURL=LlamaChat.js.map