Show warning message when last user input get pruned

core/index.d.ts

@@ -321,7 +321,8 @@ export type ChatMessageRole =
   | "assistant"
   | "thinking"
   | "system"
-  | "tool";
+  | "tool"
+  | "warning";
 
 export type TextMessagePart = {
   type: "text";
@@ -385,12 +386,18 @@ export interface SystemChatMessage {
   content: string;
 }
 
+export interface WarningChatMessage {
+  role: "warning";
+  content: string;
+}
+
 export type ChatMessage =
   | UserChatMessage
   | AssistantChatMessage
   | ThinkingChatMessage
   | SystemChatMessage
-  | ToolResultChatMessage;
+  | ToolResultChatMessage
+  | WarningChatMessage;
 
 export interface ContextItemId {
   providerTitle: string;

core/llm/countTokens.ts

@@ -27,9 +27,9 @@ class LlamaEncoding implements Encoding {
 }
 
 class NonWorkerAsyncEncoder implements AsyncEncoder {
-  constructor(private readonly encoding: Encoding) { }
+  constructor(private readonly encoding: Encoding) {}
 
-  async close(): Promise<void> { }
+  async close(): Promise<void> {}
 
   async encode(text: string): Promise<number[]> {
     return this.encoding.encode(text);
@@ -243,13 +243,15 @@ function pruneChatHistory(
   chatHistory: ChatMessage[],
   contextLength: number,
   tokensForCompletion: number,
-): ChatMessage[] {
+): [ChatMessage[], boolean] {
   let totalTokens =
     tokensForCompletion +
     chatHistory.reduce((acc, message) => {
       return acc + countChatMessageTokens(modelName, message);
     }, 0);
 
+  let shouldWarn = false;
+
   // 0. Prune any messages that take up more than 1/3 of the context length
   const longestMessages = [...chatHistory];
   longestMessages.sort((a, b) => b.content.length - a.content.length);
@@ -275,6 +277,7 @@ function pruneChatHistory(
       content,
     );
     totalTokens -= delta;
+    shouldWarn = true;
   }
 
   // 1. Replace beyond last 5 messages with summary
@@ -327,9 +330,10 @@ function pruneChatHistory(
       tokensForCompletion,
     );
     totalTokens = contextLength;
+    shouldWarn = true;
   }
 
-  return chatHistory;
+  return [chatHistory, shouldWarn];
 }
 
 function messageIsEmpty(message: ChatMessage): boolean {
@@ -367,6 +371,7 @@ function chatMessageIsEmpty(message: ChatMessage): boolean {
         !message.toolCalls
       );
     case "thinking":
+    case "warning":
     case "tool":
       return false;
   }
@@ -381,11 +386,16 @@ function compileChatMessages(
   prompt: string | undefined = undefined,
   functions: any[] | undefined = undefined,
   systemMessage: string | undefined = undefined,
-): ChatMessage[] {
+): [ChatMessage[], boolean] {
   let msgsCopy = msgs
     ? msgs
-      .map((msg) => ({ ...msg }))
-      .filter((msg) => !chatMessageIsEmpty(msg) && msg.role !== "system")
+        .map((msg) => ({ ...msg }))
+        .filter(
+          (msg) =>
+            !chatMessageIsEmpty(msg) &&
+            msg.role !== "system" &&
+            msg.role !== "warning",
+        )
     : [];
 
   msgsCopy = addSpaceToAnyEmptyMessages(msgsCopy);
@@ -445,7 +455,7 @@ function compileChatMessages(
     }
   }
 
-  const history = pruneChatHistory(
+  const [history, shouldWarn] = pruneChatHistory(
     modelName,
     msgsCopy,
     contextLength,
@@ -459,7 +469,7 @@ function compileChatMessages(
 
   const flattenedHistory = flattenMessages(history);
 
-  return flattenedHistory;
+  return [flattenedHistory, shouldWarn];
 }
 
 export {
@@ -470,6 +480,5 @@ export {
   pruneLinesFromTop,
   pruneRawPromptFromTop,
   pruneStringFromBottom,
-  pruneStringFromTop
+  pruneStringFromTop,
 };
-

core/llm/index.ts

@@ -195,11 +195,11 @@ export abstract class BaseLLM implements ILLM {
         options.completionOptions?.maxTokens ??
         (llmInfo?.maxCompletionTokens
           ? Math.min(
-            llmInfo.maxCompletionTokens,
-            // Even if the model has a large maxTokens, we don't want to use that every time,
-            // because it takes away from the context length
-            this.contextLength / 4,
-          )
+              llmInfo.maxCompletionTokens,
+              // Even if the model has a large maxTokens, we don't want to use that every time,
+              // because it takes away from the context length
+              this.contextLength / 4,
+            )
           : DEFAULT_MAX_TOKENS),
     };
     this.requestOptions = options.requestOptions;
@@ -766,7 +766,18 @@ export abstract class BaseLLM implements ILLM {
 
     completionOptions = this._modifyCompletionOptions(completionOptions);
 
-    const messages = this._compileChatMessages(completionOptions, _messages);
+    const [messages, shouldWarn] = this._compileChatMessages(
+      completionOptions,
+      _messages,
+    );
+
+    if (shouldWarn) {
+      yield {
+        role: "warning",
+        content:
+          "The context has reached its limit. This may lead to less accurate or incomplete answers.",
+      };
+    }
 
     const prompt = this.templateMessages
       ? this.templateMessages(messages)
@@ -840,7 +851,6 @@ export abstract class BaseLLM implements ILLM {
             signal,
             completionOptions,
           )) {
-
             if (chunk.role === "assistant") {
               completion += chunk.content;
               yield chunk;
@@ -948,15 +958,15 @@ export abstract class BaseLLM implements ILLM {
     );
   }
 
-  protected async * _streamComplete(
+  protected async *_streamComplete(
     prompt: string,
     signal: AbortSignal,
     options: CompletionOptions,
   ): AsyncGenerator<string> {
     throw new Error("Not implemented");
   }
 
-  protected async * _streamChat(
+  protected async *_streamChat(
     messages: ChatMessage[],
     signal: AbortSignal,
     options: CompletionOptions,