Show warning message when the most recent user input gets pruned

First, I integrated llm/compileChat to precompile chat messages before
invoking llm/streamChat. The llm/compileChat function returns both the
compiled messages and a boolean flag indicating whether the most
recent user input has been pruned. This allows us to trigger a warning
to notify users when pruning occurs at the last message.

The messageOptions is introduced to streamChat as an optional parameter.
If we pass messageOptions and the precompiled attribute is set to true,
streamChat will skip the process of compiling chat messages, ensuring
that the messages won't go through the pruning process twice.

Author: Jazzcort

core/core.ts

@@ -474,6 +474,18 @@ export class Core {
       }
     });
 
+    on("llm/compileChat", async (msg) => {
+      const { title: modelName, messages, options } = msg.data;
+      const model = await this.configHandler.llmFromTitle(modelName);
+
+      const { compiledChatMessages, lastMessageTruncated } =
+        model.compileChatMessages(options, messages);
+      return {
+        compiledChatMessages,
+        lastMessageTruncated,
+      };
+    });
+
     on("llm/streamChat", (msg) =>
       llmStreamChat(
         this.configHandler,

core/index.d.ts

@@ -132,6 +132,7 @@ export interface ILLM extends LLMOptions {
     messages: ChatMessage[],
     signal: AbortSignal,
     options?: LLMFullCompletionOptions,
+    messageOptions?: MessageOptions,
   ): AsyncGenerator<ChatMessage, PromptLog>;
 
   chat(
@@ -140,6 +141,11 @@ export interface ILLM extends LLMOptions {
     options?: LLMFullCompletionOptions,
   ): Promise<ChatMessage>;
 
+  compileChatMessages(
+    options: LLMFullCompletionOptions,
+    messages: ChatMessage[],
+  ): { compiledChatMessages: ChatMEssage[]; lastMessageTruncated: boolean };
+
   embed(chunks: string[]): Promise<number[][]>;
 
   rerank(query: string, chunks: Chunk[]): Promise<number[]>;
@@ -1388,3 +1394,7 @@ export interface TerminalOptions {
   reuseTerminal?: boolean;
   terminalName?: string;
 }
+
+export interface MessageOptions {
+  precompiled: boolean;
+}

core/llm/countTokens.ts

@@ -251,13 +251,15 @@ function pruneChatHistory(
   chatHistory: ChatMessage[],
   contextLength: number,
   tokensForCompletion: number,
-): ChatMessage[] {
+): { prunedChatHistory: ChatMessage[]; lastMessageTruncated: boolean } {
   let totalTokens =
     tokensForCompletion +
     chatHistory.reduce((acc, message) => {
       return acc + countChatMessageTokens(modelName, message);
     }, 0);
 
+  let lastMessageTruncated = false;
+
   // 0. Prune any messages that take up more than 1/3 of the context length
   const zippedMessagesAndTokens: [ChatMessage, number][] = [];
 
@@ -363,8 +365,10 @@ function pruneChatHistory(
       tokensForCompletion,
     );
     totalTokens = contextLength;
+    lastMessageTruncated = true;
   }
-  return chatHistory;
+
+  return { prunedChatHistory: chatHistory, lastMessageTruncated };
 }
 
 function messageIsEmpty(message: ChatMessage): boolean {
@@ -527,7 +531,7 @@ function compileChatMessages({
   functions: any[] | undefined;
   systemMessage: string | undefined;
   rules: Rule[];
-}): ChatMessage[] {
+}): { compiledChatMessages: ChatMessage[]; lastMessageTruncated: boolean } {
   let msgsCopy = msgs
     ? msgs
         .map((msg) => ({ ...msg }))
@@ -581,7 +585,7 @@ function compileChatMessages({
     }
   }
 
-  const history = pruneChatHistory(
+  const { prunedChatHistory: history, lastMessageTruncated } = pruneChatHistory(
     modelName,
     msgsCopy,
     contextLength,
@@ -595,7 +599,7 @@ function compileChatMessages({
 
   const flattenedHistory = flattenMessages(history);
 
-  return flattenedHistory;
+  return { compiledChatMessages: flattenedHistory, lastMessageTruncated };
 }
 
 export {

core/llm/index.ts

@@ -18,6 +18,7 @@ import {
   ILLM,
   LLMFullCompletionOptions,
   LLMOptions,
+  MessageOptions,
   ModelCapability,
   ModelInstaller,
   PromptLog,
@@ -738,6 +739,15 @@ export abstract class BaseLLM implements ILLM {
     return { role: "assistant" as const, content: completion };
   }
 
+  compileChatMessages(
+    options: LLMFullCompletionOptions,
+    messages: ChatMessage[],
+  ) {
+    let { completionOptions } = this._parseCompletionOptions(options);
+    completionOptions = this._modifyCompletionOptions(completionOptions);
+    return this._compileChatMessages(completionOptions, messages);
+  }
+
   protected modifyChatBody(
     body: ChatCompletionCreateParams,
   ): ChatCompletionCreateParams {
@@ -762,13 +772,23 @@ export abstract class BaseLLM implements ILLM {
     _messages: ChatMessage[],
     signal: AbortSignal,
     options: LLMFullCompletionOptions = {},
+    messageOptions?: MessageOptions,
   ): AsyncGenerator<ChatMessage, PromptLog> {
     let { completionOptions, logEnabled } =
       this._parseCompletionOptions(options);
 
     completionOptions = this._modifyCompletionOptions(completionOptions);
 
-    const messages = this._compileChatMessages(completionOptions, _messages);
+    const { precompiled } = messageOptions ?? {};
+
+    let messages = _messages;
+    if (!precompiled) {
+      const { compiledChatMessages } = this._compileChatMessages(
+        completionOptions,
+        _messages,
+      );
+      messages = compiledChatMessages;
+    }
 
     const prompt = this.templateMessages
       ? this.templateMessages(messages)

core/llm/streamChat.ts

@@ -23,8 +23,13 @@ export async function* llmStreamChat(
     void TTS.kill();
   }
 
-  const { title, legacySlashCommandData, completionOptions, messages } =
-    msg.data;
+  const {
+    title,
+    legacySlashCommandData,
+    completionOptions,
+    messages,
+    messageOptions,
+  } = msg.data;
 
   const model = await configHandler.llmFromTitle(title);
 
@@ -113,6 +118,7 @@ export async function* llmStreamChat(
       messages,
       new AbortController().signal,
       completionOptions,
+      messageOptions,
     );
     let next = await gen.next();
     while (!next.done) {

core/protocol/core.ts

@@ -20,6 +20,7 @@ import type {
   FileSymbolMap,
   IdeSettings,
   LLMFullCompletionOptions,
+  MessageOptions,
   ModelDescription,
   PromptLog,
   RangeInFile,
@@ -115,12 +116,21 @@ export type ToCoreFromIdeOrWebviewProtocol = {
     },
     string,
   ];
+  "llm/compileChat": [
+    {
+      title: string;
+      messages: ChatMessage[];
+      options: LLMFullCompletionOptions;
+    },
+    { compiledChatMessages: ChatMessage[]; lastMessageTruncated: boolean },
+  ];
   "llm/listModels": [{ title: string }, string[] | undefined];
   "llm/streamChat": [
     {
       messages: ChatMessage[];
       completionOptions: LLMFullCompletionOptions;
       title: string;
+      messageOptions: MessageOptions;
       legacySlashCommandData?: {
         command: SlashCommandDescription;
         input: string;

core/protocol/passThrough.ts

@@ -35,6 +35,7 @@ export const WEBVIEW_TO_CORE_PASS_THROUGH: (keyof ToCoreFromWebviewProtocol)[] =
     "autocomplete/cancel",
     "autocomplete/accept",
     "tts/kill",
+    "llm/compileChat",
     "llm/complete",
     "llm/streamChat",
     "llm/listModels",

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/constants/MessageTypes.kt

@@ -101,6 +101,7 @@ class MessageTypes {
             "autocomplete/cancel",
             "autocomplete/accept",
             "tts/kill",
+            "llm/compileChat",
             "llm/complete",
             "llm/streamChat",
             "llm/listModels",

gui/src/pages/gui/Chat.tsx

@@ -132,6 +132,9 @@ export function Chat() {
   const hasDismissedExploreDialog = useAppSelector(
     (state) => state.ui.hasDismissedExploreDialog,
   );
+  const warningMessage = useAppSelector(
+    (state) => state.session.warningMessage,
+  );
 
   useEffect(() => {
     // Cmd + Backspace to delete current step
@@ -394,6 +397,14 @@ export function Chat() {
             </ErrorBoundary>
           </div>
         ))}
+        {warningMessage.length > 0 && (
+          <div className="relative m-2 flex justify-center rounded-md border border-solid border-[#ffc107] bg-transparent p-4">
+            <p className="thread-message text-[#ffc107]">
+              <b>Warning: </b>
+              {`${warningMessage}`}
+            </p>
+          </div>
+        )}
       </StepsDiv>
       <div className={"relative"}>
         {toolCallState?.status === "generated" && <ToolCallButtons />}

gui/src/redux/slices/sessionSlice.ts

@@ -56,6 +56,7 @@ type SessionState = {
   };
   activeToolStreamId?: [string, string];
   newestToolbarPreviewForInput: Record<string, string>;
+  warningMessage: string;
 };
 
 function isCodeToEditEqual(a: CodeToEdit, b: CodeToEdit) {
@@ -94,6 +95,7 @@ const initialState: SessionState = {
   },
   lastSessionId: undefined,
   newestToolbarPreviewForInput: {},
+  warningMessage: "",
 };
 
 export const sessionSlice = createSlice({
@@ -234,6 +236,7 @@ export const sessionSlice = createSlice({
     deleteMessage: (state, action: PayloadAction<number>) => {
       // Deletes the current assistant message and the previous user message
       state.history.splice(action.payload - 1, 2);
+      state.warningMessage = "";
     },
     updateHistoryItemAtIndex: (
       state,
@@ -679,6 +682,9 @@ export const sessionSlice = createSlice({
       state.newestToolbarPreviewForInput[payload.inputId] =
         payload.contextItemId;
     },
+    setWarningMessage: (state, action: PayloadAction<string>) => {
+      state.warningMessage = action.payload;
+    },
   },
   selectors: {
     selectIsGatheringContext: (state) => {
@@ -780,6 +786,7 @@ export const {
   deleteSessionMetadata,
   setNewestToolbarPreviewForInput,
   cycleMode,
+  setWarningMessage,
 } = sessionSlice.actions;
 
 export const {

gui/src/redux/thunks/streamNormalInput.ts

@@ -9,7 +9,8 @@ import {
   addPromptCompletionPair,
   selectUseTools,
   setToolGenerated,
-  streamUpdate
+  setWarningMessage,
+  streamUpdate,
 } from "../slices/sessionSlice";
 import { ThunkApiType } from "../store";
 import { callTool } from "./callTool";
@@ -39,6 +40,34 @@ export const streamNormalInput = createAsyncThunk<
     }
     const includeTools = useTools && modelSupportsTools(defaultModel);
 
+    const res = await extra.ideMessenger.request("llm/compileChat", {
+      title: defaultModel.title,
+      messages,
+      options: includeTools
+        ? {
+            tools: state.config.config.tools.filter(
+              (tool) =>
+                toolSettings[tool.function.name] !== "disabled" &&
+                toolGroupSettings[tool.group] !== "exclude",
+            ),
+          }
+        : {},
+    });
+
+    if (res.status === "error") {
+      throw new Error(res.error);
+    }
+
+    const { compiledChatMessages, lastMessageTruncated } = res.content;
+
+    if (lastMessageTruncated) {
+      dispatch(
+        setWarningMessage(
+          "The provided context items are too large. They have been truncated to fit within the model's context length.",
+        ),
+      );
+    }
+
     // Send request
     const gen = extra.ideMessenger.llmStreamChat(
       {
@@ -52,8 +81,9 @@ export const streamNormalInput = createAsyncThunk<
             }
           : {},
         title: defaultModel.title,
-        messages,
+        messages: compiledChatMessages,
         legacySlashCommandData,
+        messageOptions: { precompiled: true },
       },
       streamAborter.signal,
     );