add support for docs via config

Author: priyashpatil

core/context/providers/DocsContextProvider.ts

@@ -5,6 +5,7 @@ import {
   ContextProviderExtras,
   ContextSubmenuItem,
   LoadSubmenuItemsArgs,
+  SiteIndexingConfig,
 } from "../../index.js";
 import { DocsService } from "../../indexing/docs/DocsService.js";
 import configs from "../../indexing/docs/preIndexedDocs.js";
@@ -108,40 +109,44 @@ class DocsContextProvider extends BaseContextProvider {
     ];
   }
 
-  async loadSubmenuItems(
-    args: LoadSubmenuItemsArgs,
-  ): Promise<ContextSubmenuItem[]> {
-    const docs = await this.docsService.list();
-    const submenuItems: ContextSubmenuItem[] = docs.map((doc) => ({
+  // Get combined site configs from preIndexedDocs and options.sites.
+  private _getDocsSitesConfig(): SiteIndexingConfig[] {
+    return [...configs, ...(this.options?.sites || [])];
+  }
+
+  // Get indexed docs as ContextSubmenuItems from database.
+  private async _getIndexedDocsContextSubmenuItems(): Promise<ContextSubmenuItem[]> {
+    return (await this.docsService.list()).map((doc) => ({
       title: doc.title,
       description: new URL(doc.baseUrl).hostname,
       id: doc.baseUrl,
-      metadata: {
-        preIndexed: !!configs.find((config) => config.title === doc.title),
-      },
     }));
+  }
 
-    submenuItems.push(
-      ...configs
-        // After it's actually downloaded, we don't want to show twice
-        .filter(
-          (config) => !submenuItems.some((item) => item.id === config.startUrl),
-        )
-        .map((config) => ({
-          title: config.title,
-          description: new URL(config.startUrl).hostname,
-          id: config.startUrl,
-          metadata: {
-            preIndexed: true,
-          },
-          // iconUrl: config.faviconUrl,
-        })),
-    );
+  async loadSubmenuItems(
+    args: LoadSubmenuItemsArgs,
+  ): Promise<ContextSubmenuItem[]> {
+    const submenuItemsMap = new Map<string, ContextSubmenuItem>();
+
+    for (const item of await this._getIndexedDocsContextSubmenuItems()) {
+      submenuItemsMap.set(item.id, item);
+    }
+
+    for (const config of this._getDocsSitesConfig()) {
+      submenuItemsMap.set(config.startUrl, {
+        id: config.startUrl,
+        title: config.title,
+        description: new URL(config.startUrl).hostname,
+        metadata: { preIndexed: !!configs.find((cnf) => cnf.title === config.title), },
+      });
+    }
+
+    const submenuItems = Array.from(submenuItemsMap.values());
 
     // Sort submenuItems such that the objects with titles which don't occur in configs occur first, and alphabetized
     submenuItems.sort((a, b) => {
-      const aTitleInConfigs = a.metadata?.preIndexed;
-      const bTitleInConfigs = b.metadata?.preIndexed;
+      const aTitleInConfigs = a.metadata?.preIndexed ?? false;
+      const bTitleInConfigs = b.metadata?.preIndexed ?? false;
 
       // Primary criterion: Items not in configs come first
       if (!aTitleInConfigs && bTitleInConfigs) {

core/core.ts

@@ -242,6 +242,19 @@ export class Core {
       await this.docsService.delete(baseUrl);
       this.messenger.send("refreshSubmenuItems", undefined);
     });
+    on("context/indexDocs", async (msg) => {
+      const config = await this.config();
+      const provider: any = config.contextProviders?.find(
+        (provider) => provider.description.title === "docs",
+      );
+
+      const siteIndexingOptions: SiteIndexingConfig[] = provider ?
+        (mProvider => mProvider?.options?.sites || [])({ ...provider }) :
+        [];
+
+      await this.indexDocs(siteIndexingOptions, msg.data.reIndex);
+      this.ide.infoPopup("Docs indexing completed");
+    });
     on("context/loadSubmenuItems", async (msg) => {
       const config = await this.config();
       const items = config.contextProviders
@@ -623,4 +636,16 @@ export class Core {
       this.indexingState = update;
     }
   }
+
+  private async indexDocs(sites: SiteIndexingConfig[], reIndex: boolean): Promise<void> {
+    for (const site of sites) {
+      for await (const update of this.docsService.indexAndAdd(site, new TransformersJsEmbeddingsProvider(), reIndex)) {
+        // Temporary disabled posting progress updates to the UI due to
+        // possible collision with code indexing progress updates.
+
+        // this.messenger.request("indexProgress", update);
+        // this.indexingState = update;
+      }
+    }
+  }
 }

core/indexing/docs/DocsService.ts

@@ -30,6 +30,7 @@ export class DocsService {
   private static instance: DocsService;
   private static DOCS_TABLE_NAME = "docs";
   private _sqliteTable: Database | undefined;
+  private docsIndexingQueue: Set<string> = new Set();
 
   public static getInstance(): DocsService {
     if (!DocsService.instance) {
@@ -204,10 +205,16 @@ export class DocsService {
   async *indexAndAdd(
     siteIndexingConfig: SiteIndexingConfig,
     embeddingsProvider: EmbeddingsProvider,
+    reIndex: boolean = false,
   ): AsyncGenerator<IndexingProgressUpdate> {
-    const startUrl = new URL(siteIndexingConfig.startUrl);
+    const startUrl = new URL(siteIndexingConfig.startUrl.toString());
 
-    if (await this.has(siteIndexingConfig.startUrl.toString())) {
+    if (this.docsIndexingQueue.has(startUrl.toString())) {
+      console.log("Already in queue");
+      return;
+    }
+
+    if (!reIndex && await this.has(startUrl.toString())) {
       yield {
         progress: 1,
         desc: "Already indexed",
@@ -216,38 +223,53 @@ export class DocsService {
       return;
     }
 
+    // Mark the site as currently being indexed
+    this.docsIndexingQueue.add(startUrl.toString());
+
     yield {
       progress: 0,
       desc: "Finding subpages",
       status: "indexing",
     };
 
     const articles: Article[] = [];
+    let processedPages = 0;
+    let maxKnownPages = 1;
 
     // Crawl pages and retrieve info as articles
     for await (const page of crawlPage(startUrl, siteIndexingConfig.maxDepth)) {
+      processedPages++;
       const article = pageToArticle(page);
       if (!article) {
         continue;
       }
       articles.push(article);
 
+      // Use a heuristic approach for progress calculation
+      const progress = Math.min(processedPages / maxKnownPages, 1);
+
       yield {
-        progress: 0,
+        progress, // Yield the heuristic progress
         desc: `Finding subpages (${page.path})`,
         status: "indexing",
       };
+
+      // Increase maxKnownPages to delay progress reaching 100% too soon
+      if (processedPages === maxKnownPages) {
+        maxKnownPages *= 2;
+      }
     }
 
     const chunks: Chunk[] = [];
     const embeddings: number[][] = [];
 
     // Create embeddings of retrieved articles
     console.log("Creating Embeddings for ", articles.length, " articles");
-    for (const article of articles) {
+    for (let i = 0; i < articles.length; i++) {
+      const article = articles[i];
       yield {
-        progress: Math.max(1, Math.floor(100 / (articles.length + 1))),
-        desc: `${article.subpath}`,
+        progress: i / articles.length,
+        desc: `Creating Embeddings: ${article.subpath}`,
         status: "indexing",
       };
 
@@ -268,7 +290,14 @@ export class DocsService {
 
     // Add docs to databases
     console.log("Adding ", embeddings.length, " embeddings to db");
+    yield {
+      progress: 0.5,
+      desc: `Adding ${embeddings.length} embeddings to db`,
+      status: "indexing",
+    };
+
     await this.add(siteIndexingConfig.title, startUrl, chunks, embeddings);
+    this.docsIndexingQueue.delete(startUrl.toString());
 
     yield {
       progress: 1,

core/protocol/core.ts

@@ -68,6 +68,7 @@ export type ToCoreFromIdeOrWebviewProtocol = {
   "autocomplete/complete": [AutocompleteInput, string[]];
   "context/addDocs": [SiteIndexingConfig, void];
   "context/removeDocs": [{ baseUrl: string }, void];
+  "context/indexDocs": [{ reIndex: boolean }, void];
   "autocomplete/cancel": [undefined, void];
   "autocomplete/accept": [{ completionId: string }, void];
   "command/run": [

core/protocol/passThrough.ts

@@ -25,6 +25,7 @@ export const WEBVIEW_TO_CORE_PASS_THROUGH: (keyof ToCoreFromWebviewProtocol)[] =
     "context/loadSubmenuItems",
     "context/addDocs",
     "context/removeDocs",
+    "context/indexDocs",
     "autocomplete/complete",
     "autocomplete/cancel",
     "autocomplete/accept",

docs/docs/customization/context-providers.md

@@ -50,6 +50,28 @@ Type `@docs` to index and retrieve snippets from any documentation site. You can
 
 Continue also pre-indexes a number of common sites, listed [here](https://github.com/continuedev/continue/blob/main/core/indexing/docs/preIndexedDocs.ts). The embeddings for these sites are hosted by us, but downloaded for local use after the first time. All other indexing occurs entirely locally.
 
+#### Adding a Documentation Site via Configuration
+
+To add a documentation site via configuration, update the `config.json` file as follows:
+
+```json
+{
+  "name": "docs",
+  "params": {
+    "sites": [
+      {
+        "title": "ExampleDocs",
+        "startUrl": "https://exampledocs.com/docs",
+        "rootUrl": "https://exampledocs.com",
+        "maxDepth": 3 // Default
+      }
+    ]
+  }
+}
+```
+
+Then run the docs indexing command: VS Code: `cmd/ctrl + shift + p`, then select `Continue: Index Docs`
+
 ### Open Files
 
 Type '@open' to reference the contents of all of your open files. Set `onlyPinned` to `true` to only reference pinned files.

extensions/vscode/package.json

@@ -229,6 +229,12 @@
         "category": "Continue",
         "title": "Fix Grammar / Spelling",
         "group": "Continue"
+      },
+      {
+        "command": "continue.docsIndex",
+        "category": "Continue",
+        "title": "Index Docs",
+        "group": "Continue"
       }
     ],
     "keybindings": [

extensions/vscode/src/commands.ts

@@ -25,6 +25,7 @@ import { VerticalPerLineDiffManager } from "./diff/verticalPerLine/manager";
 import { QuickEdit } from "./quickPicks/QuickEdit";
 import { Battery } from "./util/battery";
 import type { VsCodeWebviewProtocol } from "./webviewProtocol";
+import { Core } from "core/core";
 
 let fullScreenPanel: vscode.WebviewPanel | undefined;
 
@@ -169,6 +170,7 @@ const commandsMap: (
   verticalDiffManager: VerticalPerLineDiffManager,
   continueServerClientPromise: Promise<ContinueServerClient>,
   battery: Battery,
+  core: Core,
 ) => { [command: string]: (...args: any) => any } = (
   ide,
   extensionContext,
@@ -178,6 +180,7 @@ const commandsMap: (
   verticalDiffManager,
   continueServerClientPromise,
   battery,
+  core
 ) => {
   /**
    * Streams an inline edit to the vertical diff manager.
@@ -307,6 +310,9 @@ const commandsMap: (
     "continue.toggleAuxiliaryBar": () => {
       vscode.commands.executeCommand("workbench.action.toggleAuxiliaryBar");
     },
+    "continue.docsIndex": async () => {
+      core.invoke("context/indexDocs", {reIndex: false});
+    },
     "continue.focusContinueInput": async () => {
       const fullScreenTab = getFullScreenTab();
       if (!fullScreenTab) {
@@ -690,6 +696,7 @@ export function registerAllCommands(
   verticalDiffManager: VerticalPerLineDiffManager,
   continueServerClientPromise: Promise<ContinueServerClient>,
   battery: Battery,
+  core: Core,
 ) {
   for (const [command, callback] of Object.entries(
     commandsMap(
@@ -701,6 +708,7 @@ export function registerAllCommands(
       verticalDiffManager,
       continueServerClientPromise,
       battery,
+      core
     ),
   )) {
     context.subscriptions.push(

extensions/vscode/src/extension/VsCodeExtension.ts

@@ -190,6 +190,7 @@ export class VsCodeExtension {
       this.verticalDiffManager,
       this.core.continueServerClientPromise,
       this.battery,
+      this.core
     );
 
     registerDebugTracker(this.sidebar.webviewProtocol, this.ide);

gui/package-lock.json

@@ -66,6 +66,7 @@
   "devDependencies": {
     "@swc/cli": "^0.3.12",
     "@swc/core": "^1.5.25",
+    "@types/lodash": "^4.17.6",
     "@types/node": "^20.5.6",
     "@types/node-fetch": "^2.6.4",
     "@types/react": "^18.0.27",

gui/package.json

@@ -16,6 +16,7 @@ import { isJetBrains } from "../util";
 import { setLocalStorage } from "../util/localStorage";
 import useChatHandler from "./useChatHandler";
 import { useWebviewListener } from "./useWebviewListener";
+import { debounce } from "lodash";
 
 function useSetup(dispatch: Dispatch<any>) {
   const [configLoaded, setConfigLoaded] = useState<boolean>(false);
@@ -86,8 +87,13 @@ function useSetup(dispatch: Dispatch<any>) {
     });
   });
 
+  const debouncedIndexDocs = debounce(() => {
+    ideMessenger.request("context/indexDocs", {reIndex: false});
+  }, 1000);
+
   useWebviewListener("configUpdate", async () => {
     loadConfig();
+    debouncedIndexDocs();
   });
 
   useWebviewListener("submitMessage", async (data) => {