feat: add various strategies for handling context window

Author: mateuszlampert

packages/react-native-executorch/src/constants/llmDefaults.ts

@@ -1,4 +1,5 @@
 import { ChatConfig, Message } from '../types/llm';
+import { MessageCountContextStrategy } from '../utils/llms/context_strategy/MessageCountContextStrategy';
 
 /**
  * Default system prompt used to guide the behavior of Large Language Models (LLMs).
@@ -48,5 +49,7 @@ export const DEFAULT_CONTEXT_WINDOW_LENGTH = 5;
 export const DEFAULT_CHAT_CONFIG: ChatConfig = {
   systemPrompt: DEFAULT_SYSTEM_PROMPT,
   initialMessageHistory: DEFAULT_MESSAGE_HISTORY,
-  contextWindowLength: DEFAULT_CONTEXT_WINDOW_LENGTH,
+  contextStrategy: new MessageCountContextStrategy(
+    DEFAULT_CONTEXT_WINDOW_LENGTH
+  ),
 };

packages/react-native-executorch/src/controllers/LLMController.ts

@@ -305,15 +305,29 @@ export class LLMController {
   }
 
   public async sendMessage(message: string): Promise<string> {
-    this.messageHistoryCallback([
+    const updatedHistory = [
       ...this._messageHistory,
-      { content: message, role: 'user' },
-    ]);
-
-    const messageHistoryWithPrompt: Message[] = [
-      { content: this.chatConfig.systemPrompt, role: 'system' },
-      ...this._messageHistory.slice(-this.chatConfig.contextWindowLength),
+      { content: message, role: 'user' as const },
     ];
+    this.messageHistoryCallback(updatedHistory);
+
+    const countTokensCallback = (messages: Message[]) => {
+      const rendered = this.applyChatTemplate(
+        messages,
+        this.tokenizerConfig,
+        this.toolsConfig?.tools,
+        // eslint-disable-next-line camelcase
+        { tools_in_user_message: false, add_generation_prompt: true }
+      );
+      return this.nativeModule.getTokenCount(rendered);
+    };
+
+    const messageHistoryWithPrompt =
+      this.chatConfig.contextStrategy.buildContext(
+        this.chatConfig.systemPrompt,
+        updatedHistory,
+        countTokensCallback
+      );
 
     const response = await this.generate(
       messageHistoryWithPrompt,

packages/react-native-executorch/src/types/llm.ts

@@ -212,13 +212,13 @@ export type LLMTool = Object;
  *
  * @category Types
  * @property {Message[]} initialMessageHistory - An array of `Message` objects that represent the conversation history. This can be used to provide initial context to the model.
- * @property {number} contextWindowLength - The number of messages from the current conversation that the model will use to generate a response. The higher the number, the more context the model will have. Keep in mind that using larger context windows will result in longer inference time and higher memory usage.
  * @property {string} systemPrompt - Often used to tell the model what is its purpose, for example - "Be a helpful translator".
+ * @property {ContextStrategy} contextStrategy - Defines a strategy for managing the conversation context window and message history.
  */
 export interface ChatConfig {
   initialMessageHistory: Message[];
-  contextWindowLength: number;
   systemPrompt: string;
+  contextStrategy: ContextStrategy;
 }
 
 /**
@@ -251,6 +251,26 @@ export interface GenerationConfig {
   batchTimeInterval?: number;
 }
 
+/**
+ * Defines a strategy for managing the conversation context window and message history.
+ *
+ * @category Types
+ */
+export interface ContextStrategy {
+  /**
+   * Constructs the final array of messages to be sent to the model for the current inference step.
+   * * @param systemPrompt - The top-level instructions or persona assigned to the model.
+   * @param history - The complete conversation history up to the current point.
+   * @param getTokenCount - A callback function provided by the LLM controller that calculates the exact number of tokens a specific array of messages will consume once formatted.
+   * @returns The optimized array of messages, ready to be processed by the model.
+   */
+  buildContext(
+    systemPrompt: string,
+    history: Message[],
+    getTokenCount: (messages: Message[]) => number
+  ): Message[];
+}
+
 /**
  * Special tokens used in Large Language Models (LLMs).
  *

packages/react-native-executorch/src/utils/llms/context_strategy/MessageCountContextStrategy.ts

@@ -0,0 +1,37 @@
+import { DEFAULT_CONTEXT_WINDOW_LENGTH } from '../../../constants/llmDefaults';
+import { ContextStrategy, Message } from '../../../types/llm';
+
+/**
+ * A simple context strategy that retains a fixed number of the most recent messages.
+ * This strategy trims the conversation history based purely on the message count.
+ *
+ * @category Utils
+ */
+export class MessageCountContextStrategy implements ContextStrategy {
+  /**
+   * Initializes the MessageCountContextStrategy.
+   * * @param {number} windowLength - The maximum number of recent messages to retain in the context. Defaults to {@link DEFAULT_CONTEXT_WINDOW_LENGTH}.
+   */
+  constructor(
+    private readonly windowLength: number = DEFAULT_CONTEXT_WINDOW_LENGTH
+  ) {}
+
+  /**
+   * Builds the context by slicing the history to retain only the most recent `windowLength` messages.
+   *
+   * @param {string} systemPrompt - The top-level instructions for the model.
+   * @param {Message[]} history - The complete conversation history.
+   * @param {(messages: Message[]) => number} _getTokenCount - Unused in this strategy.
+   * @returns {Message[]} The truncated message history with the system prompt at the beginning.
+   */
+  buildContext(
+    systemPrompt: string,
+    history: Message[],
+    _getTokenCount: (messages: Message[]) => number
+  ): Message[] {
+    return [
+      { content: systemPrompt, role: 'system' as const },
+      ...history.slice(-this.windowLength),
+    ];
+  }
+}

packages/react-native-executorch/src/utils/llms/context_strategy/NaiveContextStrategy.ts

@@ -0,0 +1,27 @@
+import { ContextStrategy, Message } from '../../../types/llm';
+
+/**
+ * A context strategy that performs no filtering or trimming of the message history.
+ * * This strategy is ideal when the developer wants to manually manage the conversation
+ * context (e.g., using a custom RAG pipeline or external database) and just needs the
+ * system prompt prepended to their pre-computed history.
+ *
+ * @category Utils
+ */
+export class NaiveContextStrategy implements ContextStrategy {
+  /**
+   * Builds the context by prepending the system prompt to the entire unfiltered history.
+   *
+   * @param {string} systemPrompt - The top-level instructions for the model.
+   * @param {Message[]} history - The complete conversation history.
+   * @param {(messages: Message[]) => number} _getTokenCount - Unused in this strategy.
+   * @returns {Message[]} The unedited message history with the system prompt at the beginning.
+   */
+  buildContext(
+    systemPrompt: string,
+    history: Message[],
+    _getTokenCount: (messages: Message[]) => number
+  ): Message[] {
+    return [{ content: systemPrompt, role: 'system' as const }, ...history];
+  }
+}

packages/react-native-executorch/src/utils/llms/context_strategy/SlidingWindowContextStrategy.ts

@@ -0,0 +1,67 @@
+import { ContextStrategy, Message } from '../../../types/llm';
+
+/**
+ * An advanced, token-aware context strategy that dynamically trims the message history
+ * to ensure it fits within the model's physical context limits.
+ * * This strategy calculates the exact token count of the formatted prompt. If the prompt
+ * exceeds the allowed token budget (`maxTokens` - `bufferTokens`), it recursively
+ * removes the oldest messages.
+ *
+ * @category Utils
+ */
+export class SlidingWindowContextStrategy implements ContextStrategy {
+  /**
+   * Initializes the SlidingWindowContextStrategy.
+   * @param {number} maxTokens - The absolute maximum number of tokens the model can process (e.g., 4096).
+   * @param {number} bufferTokens - The number of tokens to keep free for the model's generated response (e.g., 1000).
+   * @param {boolean} allowOrphanedAssistantMessages - Whether to allow orphaned assistant messages when trimming the history. If false, the strategy will ensure that an assistant message is not left without its preceding user message.
+   */
+  constructor(
+    private maxTokens: number,
+    private bufferTokens: number,
+    private allowOrphanedAssistantMessages: boolean = false
+  ) {}
+
+  /**
+   * Builds the context by recursively evicting the oldest messages until the total
+   * token count is safely within the defined budget.
+   *
+   * @param {string} systemPrompt - The top-level instructions for the model.
+   * @param {Message[]} history - The complete conversation history.
+   * @param {(messages: Message[]) => number} getTokenCount - Callback to calculate the exact token count of the rendered template.
+   * @returns {Message[]} The optimized message history guaranteed to fit the token budget.
+   */
+  buildContext(
+    systemPrompt: string,
+    history: Message[],
+    getTokenCount: (messages: Message[]) => number
+  ): Message[] {
+    let localHistory = [...history];
+    const tokenBudget = this.maxTokens - this.bufferTokens;
+
+    while (localHistory.length > 1) {
+      const candidateContext: Message[] = [
+        { content: systemPrompt, role: 'system' as const },
+        ...localHistory,
+      ];
+
+      if (getTokenCount(candidateContext) <= tokenBudget) {
+        return candidateContext;
+      }
+
+      localHistory.shift();
+
+      if (!this.allowOrphanedAssistantMessages) {
+        // Prevent leaving an orphaned "assistant" response
+        if (localHistory.length > 0 && localHistory[0]?.role === 'assistant') {
+          localHistory.shift();
+        }
+      }
+    }
+
+    return [
+      { content: systemPrompt, role: 'system' as const },
+      ...localHistory,
+    ];
+  }
+}