Merge pull request #2336 from continuedev/pe/repo-map-enhancements

feat: enhance repo map logic

.vscode/launch.json

@@ -12,7 +12,7 @@
       "cwd": "${workspaceFolder}/extensions/vscode",
       "args": [
         // Pass a directory to manually test in
-        "${workspaceFolder}/manual-testing-sandbox",
+        "${workspaceFolder}/gui",
         "${workspaceFolder}/manual-testing-sandbox/test.js",
         "--extensionDevelopmentPath=${workspaceFolder}/extensions/vscode"
       ],

core/config/util.ts

@@ -1,4 +1,10 @@
-import { ContextProviderWithParams, ModelDescription } from "../";
+import {
+  ContextProviderWithParams,
+  ContinueConfig,
+  ILLM,
+  ModelDescription,
+  ModelRoles,
+} from "../";
 import { editConfigJson } from "../util/paths";
 
 function stringify(obj: any, indentation?: number): string {
@@ -61,3 +67,20 @@ export function deleteModel(title: string) {
     return config;
   });
 }
+
+export function getModelByRole<T extends keyof ModelRoles>(
+  config: ContinueConfig,
+  role: T,
+): ILLM | undefined {
+  const roleTitle = config.experimental?.modelRoles?.[role];
+
+  if (!roleTitle) {
+    return undefined;
+  }
+
+  const matchingModel = config.models.find(
+    (model) => model.title === roleTitle,
+  );
+
+  return matchingModel;
+}

core/context/providers/CodebaseContextProvider.ts

@@ -2,9 +2,9 @@ import {
   ContextItem,
   ContextProviderDescription,
   ContextProviderExtras,
-} from "../../index.js";
-import { BaseContextProvider } from "../index.js";
-import { retrieveContextItemsFromEmbeddings } from "../retrieval/retrieval.js";
+} from "../../";
+import { BaseContextProvider } from "../";
+import { retrieveContextItemsFromEmbeddings } from "../retrieval/retrieval";
 
 class CodebaseContextProvider extends BaseContextProvider {
   static description: ContextProviderDescription = {

core/context/providers/DocsContextProvider.ts

@@ -1,3 +1,4 @@
+import { INSTRUCTIONS_BASE_ITEM } from "./utils";
 import { BaseContextProvider } from "../";
 import {
   Chunk,
@@ -149,8 +150,7 @@ class DocsContextProvider extends BaseContextProvider {
         }))
         .reverse(),
       {
-        name: "Instructions",
-        description: "Instructions",
+        ...INSTRUCTIONS_BASE_ITEM,
         content:
           "Use the above documentation to answer the following question. You should not reference " +
           "anything outside of what is shown, unless it is a commonly known concept. Reference URLs " +

core/context/providers/OSContextProvider.ts

@@ -23,8 +23,10 @@ class OSContextProvider extends BaseContextProvider {
     return [
       {
         description: "Your operating system and CPU",
-        content: `I am running ${platform === "win32" ? "Windows" : platform} on ${cpu}.`,
-        name: "os",
+        content: `I am running ${
+          platform === "win32" ? "Windows" : platform
+        } on ${cpu}.`,
+        name: "Operating System",
       },
     ];
   }

core/context/providers/RepoMapContextProvider.ts

@@ -3,15 +3,28 @@ import {
   ContextItem,
   ContextProviderDescription,
   ContextProviderExtras,
+  ContextSubmenuItem,
+  LoadSubmenuItemsArgs,
 } from "../../";
-import generateRepoMap from "../../util/repoMap";
+import {
+  groupByLastNPathParts,
+  getBasename,
+  getUniqueFilePath,
+} from "../../util";
+import generateRepoMap from "../../util/generateRepoMap";
+
+const ENTIRE_PROJECT_ITEM: ContextSubmenuItem = {
+  id: "entire-codebase",
+  title: "Entire codebase",
+  description: "Search the entire codebase",
+};
 
 class RepoMapContextProvider extends BaseContextProvider {
   static description: ContextProviderDescription = {
     title: "repo-map",
     displayTitle: "Repository Map",
-    description: "List of files and signatures",
-    type: "normal",
+    description: "Select a folder",
+    type: "submenu",
     dependsOnIndexing: true,
   };
 
@@ -23,10 +36,28 @@ class RepoMapContextProvider extends BaseContextProvider {
       {
         name: "Repository Map",
         description: "Overview of the repository structure",
-        content: await generateRepoMap(extras.llm, extras.ide),
+        content: await generateRepoMap(extras.llm, extras.ide, {
+          dirs: query === ENTIRE_PROJECT_ITEM.id ? undefined : [query],
+        }),
       },
     ];
   }
+
+  async loadSubmenuItems(
+    args: LoadSubmenuItemsArgs,
+  ): Promise<ContextSubmenuItem[]> {
+    const folders = await args.ide.listFolders();
+    const folderGroups = groupByLastNPathParts(folders, 2);
+
+    return [
+      ENTIRE_PROJECT_ITEM,
+      ...folders.map((folder) => ({
+        id: folder,
+        title: getBasename(folder),
+        description: getUniqueFilePath(folder, folderGroups),
+      })),
+    ];
+  }
 }
 
 export default RepoMapContextProvider;

core/context/providers/index.ts

@@ -1,4 +1,4 @@
-import { ContextProviderName } from "../../";
+import { ContextItem, ContextProviderName } from "../../";
 import { BaseContextProvider } from "../";
 import CodeContextProvider from "./CodeContextProvider";
 import ContinueProxyContextProvider from "./ContinueProxyContextProvider";
@@ -30,7 +30,7 @@ import RepoMapContextProvider from "./RepoMapContextProvider";
  *
  * See this issue for details: https://github.com/continuedev/continue/issues/1365
  */
-const Providers: (typeof BaseContextProvider)[] = [
+export const Providers: (typeof BaseContextProvider)[] = [
   DiffContextProvider,
   FileTreeContextProvider,
   GitHubIssuesContextProvider,

core/context/providers/utils.ts

@@ -0,0 +1,7 @@
+import { ContextItem } from "../..";
+
+export const INSTRUCTIONS_BASE_ITEM: ContextItem = {
+  name: "Instructions",
+  description: "Instructions",
+  content: "",
+};

core/context/retrieval/repoMapRequest.ts

@@ -1,6 +1,33 @@
 import { Chunk, ContinueConfig, IDE, ILLM } from "../..";
+import { getModelByRole } from "../../config/util";
 import { stripImages } from "../../llm/images";
-import generateRepoMap from "../../util/repoMap";
+import generateRepoMap from "../../util/generateRepoMap";
+
+const SUPPORTED_MODEL_TITLE_FAMILIES = [
+  "claude-3",
+  "llama3.1",
+  "gemini-1.5",
+  "gpt-4",
+];
+
+function isSupportedModel(
+  config: ContinueConfig,
+  modelTitle?: string,
+): boolean {
+  if (config.experimental?.modelRoles?.applyCodeBlock) {
+    return true;
+  }
+
+  if (!modelTitle) {
+    return false;
+  }
+
+  const lowercaseModelTitle = modelTitle.toLowerCase();
+
+  return SUPPORTED_MODEL_TITLE_FAMILIES.some((title) =>
+    lowercaseModelTitle.includes(title),
+  );
+}
 
 export async function requestFilesFromRepoMap(
   defaultLlm: ILLM,
@@ -9,20 +36,16 @@ export async function requestFilesFromRepoMap(
   input: string,
   filterDirectory?: string,
 ): Promise<Chunk[]> {
-  const llm =
-    config.models.find(
-      (model) =>
-        model.title === config.experimental?.modelRoles?.repoMapFileSelection,
-    ) ?? defaultLlm;
+  const llm = getModelByRole(config, "repoMapFileSelection") ?? defaultLlm;
 
   // Only supported for Claude models right now
-  if (!llm.model.toLowerCase().includes("claude")) {
+  if (!isSupportedModel(config, llm.title)) {
     return [];
   }
 
   try {
     const repoMap = await generateRepoMap(llm, ide, {
-      signatures: false,
+      includeSignatures: false,
       dirs: filterDirectory ? [filterDirectory] : undefined,
     });
 

core/context/retrieval/retrieval.ts

@@ -1,14 +1,11 @@
-import {
-  BranchAndDir,
-  ContextItem,
-  ContextProviderExtras,
-} from "../../index.js";
-import TransformersJsEmbeddingsProvider from "../../indexing/embeddings/TransformersJsEmbeddingsProvider.js";
-import { resolveRelativePathInWorkspace } from "../../util/ideUtils.js";
-import { getRelativePath } from "../../util/index.js";
-import { RetrievalPipelineOptions } from "./pipelines/BaseRetrievalPipeline.js";
-import NoRerankerRetrievalPipeline from "./pipelines/NoRerankerRetrievalPipeline.js";
-import RerankerRetrievalPipeline from "./pipelines/RerankerRetrievalPipeline.js";
+import { BranchAndDir, ContextItem, ContextProviderExtras } from "../../";
+import TransformersJsEmbeddingsProvider from "../../indexing/embeddings/TransformersJsEmbeddingsProvider";
+import { resolveRelativePathInWorkspace } from "../../util/ideUtils";
+import { getRelativePath } from "../../util/";
+import { INSTRUCTIONS_BASE_ITEM } from "../providers/utils";
+import { RetrievalPipelineOptions } from "./pipelines/BaseRetrievalPipeline";
+import NoRerankerRetrievalPipeline from "./pipelines/NoRerankerRetrievalPipeline";
+import RerankerRetrievalPipeline from "./pipelines/RerankerRetrievalPipeline";
 
 export async function retrieveContextItemsFromEmbeddings(
   extras: ContextProviderExtras,
@@ -109,8 +106,7 @@ export async function retrieveContextItemsFromEmbeddings(
       };
     }),
     {
-      name: "Instructions",
-      description: "Instructions",
+      ...INSTRUCTIONS_BASE_ITEM,
       content:
         "Use the above code to answer the following question. You should not reference any files outside of what is shown, unless they are commonly known files, like a .gitignore or package.json. Reference the filenames whenever possible. If there isn't enough information to answer the question, suggest where the user might look to learn more.",
     },

core/indexing/CodeSnippetsIndex.ts

@@ -383,20 +383,18 @@ export class CodeSnippetsCodebaseIndex implements CodebaseIndex {
     const placeholders = likePatterns.map(() => "?").join(" OR path LIKE ");
 
     const query = `
-  SELECT path, signature
-  FROM code_snippets
-  WHERE path LIKE ${placeholders}
-  ORDER BY path
-  LIMIT ? OFFSET ?
-`;
+    SELECT DISTINCT path, signature
+    FROM code_snippets
+    WHERE path LIKE ${placeholders}
+    ORDER BY path, signature
+    LIMIT ? OFFSET ?
+  `;
 
     const rows = await db.all(query, [...likePatterns, batchSize, offset]);
 
-    const validRows = rows.filter((row) => row.path && row.signature !== null);
-
     const groupedByPath: { [path: string]: string[] } = {};
 
-    for (const { path, signature } of validRows) {
+    for (const { path, signature } of rows) {
       if (!groupedByPath[path]) {
         groupedByPath[path] = [];
       }

core/indexing/docs/DocsCrawler.ts

@@ -411,7 +411,7 @@ class ChromiumCrawler {
 
       for (const link of linkGroup) {
         enqueuedLinkCount++;
-        console.log({ enqueuedLinkCount, url: this.startUrl });
+        console.log({ enqueuedLinkCount, url: this.startUrl.toString() });
         if (enqueuedLinkCount <= this.maxRequestsPerCrawl) {
           yield* this.crawlSitePages(page, new URL(link), visitedLinks);
         }