Merge pull request #3598 from continuedev/dallin/codeblock-symbols

Include Symbols from Toolbar Codeblocks

gui/src/components/markdown/StyledMarkdownPreview.tsx

@@ -28,6 +28,8 @@ import { fixDoubleDollarNewLineLatex } from "./utils/fixDoubleDollarLatex";
 import { selectUIConfig } from "../../redux/slices/configSlice";
 import { ToolTip } from "../gui/Tooltip";
 import { v4 as uuidv4 } from "uuid";
+import { ContextItemWithId, RangeInFileWithContents } from "core";
+import { getContextItemsFromHistory } from "../../redux/thunks/updateFileSymbols";
 
 const StyledMarkdown = styled.div<{
   fontSize?: number;
@@ -139,34 +141,34 @@ const StyledMarkdownPreview = memo(function StyledMarkdownPreview(
   // So they won't use the most up-to-date state values
   // So in this case we just put them in refs
 
-  // Grab context items that are this one or further back
+  // The logic here is to get file names from
+  // 1. Context items found in past messages
+  // 2. Toolbar Codeblocks found in past messages
   const history = useAppSelector((state) => state.session.history);
-  const previousFileContextItems = useMemo(() => {
+  const allSymbols = useAppSelector((state) => state.session.symbols);
+  const pastFileInfo = useMemo(() => {
     const index = props.itemIndex;
     if (index === undefined) {
-      return [];
+      return {
+        symbols: [],
+        rifs: [],
+      };
     }
-    const previousItems = history.flatMap((item, i) =>
-      i <= index ? item.contextItems : [],
-    );
-    return previousItems.filter(
-      (item) => item.uri?.type === "file" && item?.uri?.value,
-    );
-  }, [props.itemIndex, history]);
-  const previousFileContextItemsRef = useUpdatingRef(previousFileContextItems);
-
-  // Extract global symbols for files matching previous context items
-  const allSymbols = useAppSelector((state) => state.session.symbols);
-  const previousFileContextItemSymbols = useMemo(() => {
-    const uniqueUris = new Set(
-      previousFileContextItems.map((item) => item.uri!.value!),
+    const pastContextItems = getContextItemsFromHistory(history, index);
+    const rifs = pastContextItems.map((item) =>
+      ctxItemToRifWithContents(item, true),
     );
-    return Object.entries(allSymbols)
-      .filter((e) => uniqueUris.has(e[0]))
+    const symbols = Object.entries(allSymbols)
+      .filter((e) => pastContextItems.find((item) => item.uri!.value === e[0]))
       .map((f) => f[1])
       .flat();
-  }, [allSymbols, previousFileContextItems]);
-  const symbolsRef = useUpdatingRef(previousFileContextItemSymbols);
+
+    return {
+      symbols,
+      rifs,
+    };
+  }, [props.itemIndex, history, allSymbols]);
+  const pastFileInfoRef = useUpdatingRef(pastFileInfo);
 
   const [reactContent, setMarkdownSource] = useRemark({
     remarkPlugins: [
@@ -297,33 +299,28 @@ const StyledMarkdownPreview = memo(function StyledMarkdownPreview(
           const content = getCodeChildrenContent(codeProps.children);
 
           if (content) {
+            const { symbols, rifs } = pastFileInfoRef.current;
+
             // Insert file links for matching previous context items
             // With some reasonable limitations on what might be a filename
-            if (
-              previousFileContextItemsRef.current?.length &&
-              content.includes(".") &&
-              content.length > 2
-            ) {
-              const ctxItem = previousFileContextItemsRef.current.find(
-                (item) => item.uri!.value!.split("/").pop() === content, // Exact match for last segment of URI
+            if (rifs.length && content.includes(".") && content.length > 2) {
+              const match = rifs.find(
+                (rif) => rif.filepath.split("/").pop() === content, // Exact match for last segment of URI
               );
 
-              if (ctxItem) {
-                const rif = ctxItemToRifWithContents(ctxItem);
-                return <FilenameLink rif={rif} />;
+              if (match) {
+                return <FilenameLink rif={match} />;
               }
             }
 
             // Insert symbols for exact matches
-            const exactSymbol = symbolsRef.current.find(
-              (s) => s.name === content,
-            );
+            const exactSymbol = symbols.find((s) => s.name === content);
             if (exactSymbol) {
               return <SymbolLink content={content} symbol={exactSymbol} />;
             }
 
             // Partial matches - this is the case where the llm returns e.g. `subtract(number)` instead of `subtract`
-            const partialSymbol = symbolsRef.current.find((s) =>
+            const partialSymbol = symbols.find((s) =>
               content.startsWith(s.name),
             );
             if (partialSymbol) {

gui/src/redux/thunks/gatherContext.ts

@@ -6,14 +6,10 @@ import {
   MessageContent,
   RangeInFile,
 } from "core";
-import resolveEditorContent, {
-  hasSlashCommandOrContextProvider,
-} from "../../components/mainInput/resolveInput";
+import resolveEditorContent from "../../components/mainInput/resolveInput";
 import { ThunkApiType } from "../store";
 import { selectDefaultModel } from "../slices/configSlice";
-import { setIsGatheringContext } from "../slices/sessionSlice";
 import { findUriInDirs, getUriPathBasename } from "core/util/uri";
-import { updateFileSymbolsFromNewContextItems } from "./updateFileSymbols";
 
 export const gatherContext = createAsyncThunk<
   {
@@ -97,8 +93,6 @@ export const gatherContext = createAsyncThunk<
       }
     }
 
-    dispatch(updateFileSymbolsFromNewContextItems(selectedContextItems));
-
     if (promptPreamble) {
       if (typeof content === "string") {
         content = promptPreamble + content;

gui/src/redux/thunks/streamResponse.ts

@@ -16,6 +16,7 @@ import { resetStateForNewMessage } from "./resetStateForNewMessage";
 import { streamNormalInput } from "./streamNormalInput";
 import { streamSlashCommand } from "./streamSlashCommand";
 import { streamThunkWrapper } from "./streamThunkWrapper";
+import { updateFileSymbolsFromFiles } from "./updateFileSymbols";
 
 const getSlashCommandForInput = (
   input: MessageContent,
@@ -83,6 +84,15 @@ export const streamResponseThunk = createAsyncThunk<
         const unwrapped = unwrapResult(result);
         const { selectedContextItems, selectedCode, content } = unwrapped;
 
+        // symbols for both context items AND selected codeblocks
+        const filesForSymbols = [
+          ...selectedContextItems
+            .filter((item) => item.uri?.type === "file" && item?.uri?.value)
+            .map((item) => item.uri!.value),
+          ...selectedCode.map((rif) => rif.filepath),
+        ];
+        dispatch(updateFileSymbolsFromFiles(filesForSymbols));
+
         dispatch(
           updateHistoryItemAtIndex({
             index: inputIndex,

gui/src/redux/thunks/updateFileSymbols.ts

@@ -1,28 +1,53 @@
 import { createAsyncThunk } from "@reduxjs/toolkit";
 import { ThunkApiType } from "../store";
 import { updateFileSymbols } from "../slices/sessionSlice";
-import { ContextItemWithId } from "core";
+import { ChatHistoryItem, ContextItemWithId } from "core";
+
+export function getContextItemsFromHistory(
+  historyItems: ChatHistoryItem[],
+  priorToIndex?: number,
+) {
+  const pastHistoryItems = historyItems.filter(
+    (_, i) => i <= (priorToIndex ?? historyItems.length - 1),
+  );
+
+  const pastNormalContextItems = pastHistoryItems.flatMap((item) => {
+    return (
+      item.contextItems?.filter(
+        (item) => item.uri?.type === "file" && item.uri?.value,
+      ) ?? []
+    );
+  });
+  const pastToolbarContextItems: ContextItemWithId[] = pastHistoryItems
+    .filter(
+      (item) => item.editorState && Array.isArray(item.editorState.content),
+    )
+    .flatMap((item) => item.editorState.content)
+    .filter(
+      (content) =>
+        content?.type === "codeBlock" &&
+        content?.attrs?.item?.uri?.value &&
+        content.attrs.item.uri?.type === "file",
+    )
+    .map((content) => content.attrs!.item!);
+
+  return [...pastNormalContextItems, ...pastToolbarContextItems];
+}
 
 /*
     Get file symbols for given context items
     Overwrite symbols for existing file matches
 */
-export const updateFileSymbolsFromNewContextItems = createAsyncThunk<
+export const updateFileSymbolsFromFiles = createAsyncThunk<
   void,
-  ContextItemWithId[],
+  string[],
   ThunkApiType
 >(
   "symbols/updateFromContextItems",
-  async (contextItems, { dispatch, extra, getState }) => {
+  async (filepaths, { dispatch, extra, getState }) => {
     try {
       // Get unique file uris from context items
-      const uniqueUris = Array.from(
-        new Set(
-          contextItems
-            .filter((item) => item.uri?.type === "file" && item?.uri?.value)
-            .map((item) => item.uri!.value),
-        ),
-      );
+      const uniqueUris = Array.from(new Set(filepaths));
 
       // Update symbols for those files
       if (uniqueUris.length > 0) {
@@ -36,11 +61,7 @@ export const updateFileSymbolsFromNewContextItems = createAsyncThunk<
         dispatch(updateFileSymbols(result.content));
       }
     } catch (e) {
-      console.error(
-        "Error updating file symbols from context items",
-        e,
-        contextItems,
-      );
+      console.error("Error updating file symbols from filepaths", e, filepaths);
     }
   },
 );
@@ -58,9 +79,8 @@ export const updateFileSymbolsFromHistory = createAsyncThunk<
     const state = getState();
 
     // Get unique context item file uris from all history
-    const contextItems = state.session.history.flatMap(
-      (item) => item.contextItems,
-    );
+    const contextItems = getContextItemsFromHistory(state.session.history);
+
     const uniqueUris = new Set(
       contextItems
         .filter((item) => item.uri?.type === "file" && item?.uri?.value)