rework context strategies maContextLength

Author: mateuszlampert

apps/llm/app/llm/index.tsx

@@ -39,7 +39,7 @@ function LLMScreen() {
   useEffect(() => {
     llm.configure({
       chatConfig: {
-        contextStrategy: new SlidingWindowContextStrategy(2048, 512),
+        contextStrategy: new SlidingWindowContextStrategy(512),
       },
     });
     // eslint-disable-next-line react-hooks/exhaustive-deps

packages/react-native-executorch/common/rnexecutorch/host_objects/ModelHostObject.h

@@ -133,6 +133,11 @@ template <typename Model> class ModelHostObject : public JsiHostObject {
                                        synchronousHostFunction<&Model::setTopp>,
                                        "setTopp"));
 
+      addFunctions(JSI_EXPORT_FUNCTION(
+          ModelHostObject<Model>,
+          synchronousHostFunction<&Model::getMaxContextLength>,
+          "getMaxContextLength"));
+
       addFunctions(
           JSI_EXPORT_FUNCTION(ModelHostObject<Model>, unload, "unload"));
 

packages/react-native-executorch/common/rnexecutorch/models/llm/LLM.cpp

@@ -146,6 +146,16 @@ void LLM::setTopp(float topp) {
   }
   runner->set_topp(topp);
 }
+
+size_t LLM::getMaxContextLength() const {
+  if (!runner || !runner->is_loaded()) {
+    throw RnExecutorchError(
+        RnExecutorchErrorCode::ModuleNotLoaded,
+        "Can't get context length from a model that's not loaded");
+  }
+  return runner->get_max_context_length();
+}
+
 void LLM::unload() noexcept { runner.reset(nullptr); }
 
 } // namespace rnexecutorch::models::llm

packages/react-native-executorch/common/rnexecutorch/models/llm/LLM.h

@@ -31,6 +31,7 @@ class LLM : public BaseModel {
   void setTemperature(float temperature);
   void setTopp(float topp);
   void setTimeInterval(size_t timeInterval);
+  size_t getMaxContextLength() const;
 
 private:
   std::unique_ptr<example::Runner> runner;

packages/react-native-executorch/common/runner/runner.cpp

@@ -342,6 +342,26 @@ void Runner::set_topp(float topp) noexcept {
   }
 }
 
+int32_t Runner::get_max_context_length() const {
+  if (!is_loaded()) {
+    return static_cast<int32_t>(metadata_.at(kMaxContextLen));
+  }
+  return config_.max_context_length;
+}
+
+int32_t Runner::count_text_tokens(const std::string &text) const {
+  auto encodeResult =
+      tokenizer_->encode(text, numOfAddedBoSTokens, numOfAddedEoSTokens);
+
+  if (!encodeResult.ok()) {
+    throw rnexecutorch::RnExecutorchError(
+        rnexecutorch::RnExecutorchErrorCode::TokenizerError,
+        "Encoding failed during token count check.");
+  }
+
+  return static_cast<int32_t>(encodeResult.get().size());
+}
+
 int32_t Runner::resolve_max_new_tokens(int32_t num_prompt_tokens,
                                        int32_t max_seq_len,
                                        int32_t max_context_len,
@@ -368,17 +388,4 @@ int32_t Runner::resolve_max_new_tokens(int32_t num_prompt_tokens,
   return std::max(0, result);
 }
 
-int32_t Runner::count_text_tokens(const std::string &text) const {
-  auto encodeResult =
-      tokenizer_->encode(text, numOfAddedBoSTokens, numOfAddedEoSTokens);
-
-  if (!encodeResult.ok()) {
-    throw rnexecutorch::RnExecutorchError(
-        rnexecutorch::RnExecutorchErrorCode::TokenizerError,
-        "Encoding failed during token count check.");
-  }
-
-  return static_cast<int32_t>(encodeResult.get().size());
-}
-
 } // namespace example

packages/react-native-executorch/common/runner/runner.h

@@ -51,6 +51,7 @@ class Runner : public llm::IRunner {
   void set_temperature(float temperature) noexcept;
   void set_topp(float topp) noexcept;
   int32_t count_text_tokens(const std::string &text) const;
+  int32_t get_max_context_length() const;
 
   void stop() override;
   void reset() override;

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

@@ -321,11 +321,12 @@ export class LLMController {
       );
       return this.nativeModule.countTextTokens(rendered);
     };
-
+    const maxContextLength = this.nativeModule.getMaxContextLength();
     const messageHistoryWithPrompt =
       this.chatConfig.contextStrategy.buildContext(
         this.chatConfig.systemPrompt,
         updatedHistory,
+        maxContextLength,
         countTokensCallback
       );
 

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

@@ -261,12 +261,14 @@ 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 maxContextLength - The maximum number of tokens that the model can keep in the context.
    * @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[],
+    maxContextLength: number,
     getTokenCount: (messages: Message[]) => number
   ): Message[];
 }

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

@@ -22,11 +22,13 @@ export class MessageCountContextStrategy implements ContextStrategy {
    * @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.
+   * @param {number} _maxContextLength - Unused in this strategy.
    * @returns {Message[]} The truncated message history with the system prompt at the beginning.
    */
   buildContext(
     systemPrompt: string,
     history: Message[],
+    _maxContextLength: number,
     _getTokenCount: (messages: Message[]) => number
   ): Message[] {
     return [

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

@@ -3,8 +3,7 @@ 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.
+ * context.
  *
  * @category Utils
  */
@@ -14,12 +13,14 @@ export class NaiveContextStrategy implements ContextStrategy {
    *
    * @param {string} systemPrompt - The top-level instructions for the model.
    * @param {Message[]} history - The complete conversation history.
+   * @param {number} _maxContextLength - Unused in this strategy.
    * @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[],
+    _maxContextLength: number,
     _getTokenCount: (messages: Message[]) => number
   ): Message[] {
     return [{ content: systemPrompt, role: 'system' as const }, ...history];