Forbid file mentioning for large files

core/core.ts

@@ -22,6 +22,7 @@ import { DataLogger } from "./data/log";
 import { streamDiffLines } from "./edit/streamDiffLines";
 import { CodebaseIndexer, PauseToken } from "./indexing/CodebaseIndexer";
 import DocsService from "./indexing/docs/DocsService";
+import { countTokens } from "./llm/countTokens";
 import Ollama from "./llm/llms/Ollama";
 import { createNewPromptFileV2 } from "./promptFiles/v2/createNewPromptFile";
 import { callTool } from "./tools/callTool";
@@ -888,6 +889,26 @@ export class Core {
 
       return { contextItems };
     });
+
+    on("isItemTooBig", async ({ data: { item, selectedModelTitle } }) => {
+      const { config } = await this.configHandler.loadConfig();
+
+      if (!config) {
+        return false;
+      }
+
+      const llm = await this.configHandler.llmFromTitle(selectedModelTitle);
+
+      // Count the size of the file tokenwise
+      const tokens = countTokens(item.content);
+
+      // File exceeds context length of the model
+      if (tokens > llm.contextLength - llm.completionOptions!.maxTokens!) {
+        return true;
+      }
+
+      return false;
+    });
   }
 
   private indexingCancellationController: AbortController | undefined;

core/protocol/core.ts

@@ -196,4 +196,8 @@ export type ToCoreFromIdeOrWebviewProtocol = {
   ];
   "clipboardCache/add": [{ content: string }, void];
   "controlPlane/openUrl": [{ path: string; orgSlug: string | undefined }, void];
+  isItemTooBig: [
+    { item: ContextItemWithId; selectedModelTitle: string | undefined },
+    boolean,
+  ];
 };

core/protocol/passThrough.ts

@@ -59,6 +59,7 @@ export const WEBVIEW_TO_CORE_PASS_THROUGH: (keyof ToCoreFromWebviewProtocol)[] =
     "didChangeSelectedOrg",
     "tools/call",
     "controlPlane/openUrl",
+    "isItemTooBig",
   ];
 
 // Message types to pass through from core to webview

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

@@ -126,6 +126,7 @@ class MessageTypes {
             "didChangeSelectedOrg",
             "tools/call",
             "controlPlane/openUrl",
+            "isItemTooBig",
         )
     }
 }

gui/src/components/mainInput/AtMentionDropdown/index.tsx

@@ -5,6 +5,7 @@ import {
   PlusIcon,
 } from "@heroicons/react/24/outline";
 import { Editor } from "@tiptap/react";
+import { RangeInFile } from "core";
 import {
   forwardRef,
   useContext,
@@ -24,6 +25,7 @@ import {
   vscQuickInputBackground,
 } from "../..";
 import { IdeMessengerContext } from "../../../context/IdeMessenger";
+import { useAppSelector } from "../../../redux/hooks";
 import { setDialogMessage, setShowDialog } from "../../../redux/slices/uiSlice";
 import { fontSize } from "../../../util";
 import FileIcon from "../../FileIcon";
@@ -138,11 +140,31 @@ interface AtMentionDropdownProps {
   onClose: () => void;
 }
 
+const formatFileSize = (fileSize: number) => {
+  const KB = 1000;
+  const MB = 1000_000;
+  const GB = 1000_000_000;
+
+  if (fileSize > GB) {
+    return `${(fileSize / GB).toFixed(1)} GB`;
+  } else if (fileSize > MB) {
+    return `${(fileSize / MB).toFixed(1)} MB`;
+  } else if (fileSize > KB) {
+    return `${(fileSize / KB).toFixed(1)} KB`;
+  }
+
+  return `${fileSize} byte${fileSize > 1 ? "s" : ""}`;
+};
+
 const AtMentionDropdown = forwardRef((props: AtMentionDropdownProps, ref) => {
   const dispatch = useDispatch();
 
   const ideMessenger = useContext(IdeMessengerContext);
 
+  const selectedModelTitle = useAppSelector(
+    (store) => store.config.defaultModelTitle,
+  );
+
   const [selectedIndex, setSelectedIndex] = useState(0);
 
   const [subMenuTitle, setSubMenuTitle] = useState<string | undefined>(
@@ -157,6 +179,88 @@ const AtMentionDropdown = forwardRef((props: AtMentionDropdownProps, ref) => {
 
   const [allItems, setAllItems] = useState<ComboBoxItem[]>([]);
 
+  async function isItemTooBig(
+    name: string,
+    query: string,
+  ): Promise<[boolean, number]> {
+    const selectedCode: RangeInFile[] = [];
+    // Get context item from core
+    const contextResult = await ideMessenger.request(
+      "context/getContextItems",
+      {
+        name,
+        query,
+        fullInput: "",
+        selectedCode,
+        selectedModelTitle: selectedModelTitle ?? "",
+      },
+    );
+
+    if (contextResult.status === "error") {
+      return [false, -1];
+    }
+
+    const item = contextResult.content[0];
+
+    // Check if the context item exceeds the context length of the selected model
+    const result = await ideMessenger.request("isItemTooBig", {
+      item,
+      selectedModelTitle: selectedModelTitle,
+    });
+
+    if (result.status === "error") {
+      return [false, -1];
+    }
+
+    const size = new Blob([item.content]).size;
+
+    return [result.content, size];
+  }
+
+  function handleItemTooBig(
+    fileExceeds: boolean,
+    fileSize: number,
+    item: ComboBoxItem,
+  ) {
+    if (fileExceeds) {
+      props.editor
+        .chain()
+        .focus()
+        .command(({ tr, state }) => {
+          const text = state.doc.textBetween(
+            0,
+            state.selection.from,
+            "\n",
+            "\n",
+          ); // Get the text before the cursor
+          const lastAtIndex = text.lastIndexOf("@");
+
+          if (lastAtIndex !== -1) {
+            // Delete text after the last "@"
+            tr.delete(lastAtIndex + 1, state.selection.from);
+            return true;
+          }
+          return false;
+        })
+        .run();
+
+      // Trigger warning message
+      ideMessenger.ide.showToast(
+        "warning",
+        fileSize > 0 ? "File exceeds context length" : "Can't load the file",
+        {
+          modal: true,
+          detail:
+            fileSize > 0
+              ? `'${item.title}' is ${formatFileSize(fileSize)} which exceeds the allowed context length and connot be processed by the model`
+              : `'${item.title}' could not be loaded. Please check if the file exists and has the correct permissions.`,
+        },
+      );
+    } else {
+      props.command({ ...item, itemType: item.type });
+    }
+  }
+
   useEffect(() => {
     const items = [...props.items];
     if (subMenuTitle === "Type to search docs") {
@@ -243,7 +347,13 @@ const AtMentionDropdown = forwardRef((props: AtMentionDropdownProps, ref) => {
     }
 
     if (item) {
-      props.command({ ...item, itemType: item.type });
+      if (item.type === "file" && item.query) {
+        isItemTooBig(item.type, item.query).then(([fileExceeds, fileSize]) =>
+          handleItemTooBig(fileExceeds, fileSize, item),
+        );
+      } else {
+        props.command({ ...item, itemType: item.type });
+      }
     }
   };