[RAG] Embed workspace files

extensions/void/package.json

@@ -146,5 +146,16 @@
     "typescript": "5.5.4",
     "typescript-eslint": "^8.3.0",
     "uuid": "^10.0.0"
+  },
+  "dependencies": {
+    "@anthropic-ai/sdk": "^0.27.1",
+    "@langchain/community": "^0.3.4",
+    "@langchain/core": "^0.3.8",
+    "@langchain/openai": "^0.3.5",
+    "@opensearch-project/opensearch": "^2.12.0",
+    "langchain": "^0.3.2"
+  },
+  "overrides": {
+    "@langchain/core": "^0.3.8"
   }
 }

extensions/void/src/ai/constants.ts

@@ -0,0 +1,11 @@
+export enum ApiProvider {
+	ANTHROPIC = 'anthropic',
+	OPENAI = 'openAI',
+	GREPTILE = 'greptile',
+	OLLAMA = 'ollama',
+	OPENAI_COMPATIBLE = 'openAICompatible'
+}
+
+export enum VectorStore {
+	OPENSEARCH = 'openSearch'
+}

extensions/void/src/ai/embed.ts

@@ -0,0 +1,172 @@
+import * as vscode from "vscode";
+import * as path from "path";
+import * as fs from "fs";
+import { OpenAIEmbeddings } from "@langchain/openai";
+import {
+	RecursiveCharacterTextSplitter,
+	SupportedTextSplitterLanguage,
+} from "@langchain/textsplitters";
+import { JSONLoader } from "langchain/document_loaders/fs/json";
+import { TextLoader } from "langchain/document_loaders/fs/text";
+import { Embeddings } from "@langchain/core/embeddings";
+import { getVectorStoreClient } from "./vectorStore/index";
+import { VoidConfig } from "../sidebar/contextForConfig";
+import { ApiProvider } from "./constants";
+
+enum FileType {
+	UNKNOWN = "unknown",
+	CODE = "code",
+	TEXT = "text",
+	JSON = "json",
+}
+
+const detectFileType = (
+	file: vscode.Uri
+): { type: FileType; language?: SupportedTextSplitterLanguage } => {
+	switch (path.extname(file.fsPath)) {
+		case ".js":
+		case ".ts":
+		case ".jsx":
+		case ".tsx":
+			return { type: FileType.CODE, language: "js" };
+		case ".py":
+			return { type: FileType.CODE, language: "python" };
+		case ".md":
+			return { type: FileType.CODE, language: "markdown" };
+		case ".tex":
+			return { type: FileType.CODE, language: "latex" };
+		case ".html":
+			return { type: FileType.CODE, language: "html" };
+		case ".php":
+			return { type: FileType.CODE, language: "php" };
+		case ".txt":
+		case ".csv":
+			return { type: FileType.TEXT };
+		case ".json":
+			return { type: FileType.JSON };
+		default:
+			return { type: FileType.UNKNOWN };
+	}
+};
+
+const getSplitter = (
+	type: FileType,
+	language?: SupportedTextSplitterLanguage
+) => {
+	switch (type) {
+		case FileType.CODE:
+			return (
+				language &&
+				RecursiveCharacterTextSplitter.fromLanguage(language, {
+					chunkSize: 300,
+					chunkOverlap: 50,
+					separators: ["\n\n", "\n", " ", ""],
+				})
+			);
+
+		case FileType.TEXT:
+			return new RecursiveCharacterTextSplitter({
+				chunkSize: 1000,
+				chunkOverlap: 100,
+			});
+
+		default:
+			return null;
+	}
+};
+
+const getLoader = (type: FileType, file: vscode.Uri) => {
+	switch (type) {
+		case FileType.JSON:
+			return new JSONLoader(file.fsPath);
+
+		case FileType.TEXT:
+		case FileType.CODE:
+			return new TextLoader(file.fsPath);
+
+		default:
+			return null;
+	}
+};
+
+const getEmbeddingClient = (voidConfig: VoidConfig): Embeddings | null => {
+	switch (voidConfig.default.whichApi) {
+		case ApiProvider.OPENAI:
+			return new OpenAIEmbeddings({
+				model: voidConfig.openAI.embedding,
+				apiKey: voidConfig.openAI.apikey,
+			});
+		default:
+			return null;
+	}
+};
+
+export const embedWorkspaceFiles = async (voidConfig: VoidConfig) => {
+	const embeddingClient = getEmbeddingClient(voidConfig);
+	const vectorStore =
+		embeddingClient && getVectorStoreClient(voidConfig, embeddingClient);
+
+	// if embedding and vector store keys are configured, proceed
+	if (embeddingClient && vectorStore) {
+		const excludePatterns = Object.keys(
+			vscode.workspace.getConfiguration("files").get("exclude") || {}
+		).join(",");
+
+		const files = await vscode.workspace.findFiles(
+			"**",
+			`{${excludePatterns}}`
+		);
+
+		files?.forEach(async (file) => {
+			console.debug(`Embedding file: ${file.fsPath}`);
+
+			// check if file has been modified since last embedding
+			const stat = fs.statSync(file.fsPath);
+			const mtime = stat.mtime.getTime();
+			const storedMtime = await vectorStore.getStoredMtime(file.fsPath);
+
+			// either file is new or has been modified since last embedding
+			if (!storedMtime || mtime > storedMtime) {
+				const { type, language } = detectFileType(file);
+				const textSplitter = getSplitter(type, language);
+				const fileLoader = getLoader(type, file);
+
+				// for already embedded files, delete the old embeddings so they don't show up in search results
+				if (storedMtime) {
+					console.debug(`File ${file.fsPath} modified since last embedding`);
+					await vectorStore.deleteDocuments(file.fsPath);
+				} else {
+					console.debug(`File ${file.fsPath} is new`);
+				}
+
+				// if we handle this file type, embed it and save to vector store
+				if (textSplitter && fileLoader) {
+					const docs = await fileLoader.load();
+
+					const docsWithMetadata = docs.map((doc) => ({
+						...doc,
+						metadata: {
+							...doc.metadata,
+							mtime,
+						},
+					}));
+
+					const chunks = await textSplitter.splitDocuments(docsWithMetadata);
+
+					await vectorStore.uploadDocuments(chunks);
+
+					console.debug(`File ${file.fsPath} embedded`);
+				} else {
+					console.debug(`File ${file.fsPath} is of an unsupported type`);
+				}
+			} else {
+				console.debug(`File ${file.fsPath} is up to date, skipping embedding`);
+			}
+		});
+	} else {
+		console.error("Embedding client or vector store client not configured", {
+			embeddingClient,
+			vectorStore,
+		});
+	}
+};

extensions/void/src/ai/vectorStore/index.ts

@@ -0,0 +1,25 @@
+import { Document } from "@langchain/core/documents";
+import openSearchInstance from "./openSearch";
+import { Embeddings } from "@langchain/core/embeddings";
+import { VoidConfig } from "../../sidebar/contextForConfig";
+import { VectorStore } from "../constants";
+
+export const INDEX_NAME = "void";
+
+export interface VectorStoreAdapter {
+	uploadDocuments: (documents: Document[]) => Promise<void>;
+	deleteDocuments: (path: string) => Promise<void>;
+	getStoredMtime: (path: string) => Promise<number | null>;
+}
+
+export const getVectorStoreClient = (
+	voidConfig: VoidConfig,
+	embeddingApi: Embeddings
+) => {
+	switch (voidConfig.default.vectorStore) {
+		case VectorStore.OPENSEARCH:
+			return openSearchInstance(voidConfig, embeddingApi);
+		default:
+			return null;
+	}
+};

extensions/void/src/ai/vectorStore/openSearch.ts

@@ -0,0 +1,74 @@
+import { Client as OpenSearchClient } from "@opensearch-project/opensearch";
+import { Document } from "@langchain/core/documents";
+import { OpenSearchVectorStore } from "@langchain/community/vectorstores/opensearch";
+import { Embeddings, EmbeddingsInterface } from "@langchain/core/embeddings";
+import { INDEX_NAME, VectorStoreAdapter } from ".";
+import { VoidConfig } from "../../sidebar/contextForConfig";
+
+const uploadDocuments = async (
+	client: OpenSearchClient,
+	embeddingApi: EmbeddingsInterface,
+	documents: Document[]
+) => {
+	await OpenSearchVectorStore.fromDocuments(documents, embeddingApi, {
+		client,
+		indexName: INDEX_NAME,
+	});
+};
+
+const deleteDocuments = async (client: OpenSearchClient, path: string) => {
+	const hasIndexResponse = await client.indices.exists({ index: INDEX_NAME });
+
+	if (hasIndexResponse.body) {
+		await client.deleteByQuery({
+			index: INDEX_NAME,
+			body: {
+				query: {
+					match: {
+						source: path,
+					},
+				},
+			},
+		});
+	}
+};
+
+const getStoredMtime = async (
+	client: OpenSearchClient,
+	embeddingApi: EmbeddingsInterface,
+	path: string
+) => {
+	const hasIndexResponse = await client.indices.exists({ index: INDEX_NAME });
+
+	if (hasIndexResponse.body) {
+		const vectorStore = await OpenSearchVectorStore.fromExistingIndex(
+			embeddingApi,
+			{ client, indexName: INDEX_NAME }
+		);
+
+		const results = await vectorStore.similaritySearch("", 1, {
+			source: path,
+		});
+
+		return results?.[0]?.metadata?.mtime;
+	} else {
+		return null;
+	}
+};
+
+export default (
+	voidConfig: VoidConfig,
+	embeddingApi: Embeddings
+): VectorStoreAdapter => {
+	const client = new OpenSearchClient({
+		nodes: [voidConfig.openSearch.endpoint],
+	});
+
+	return {
+		uploadDocuments: (documents: Document[]) =>
+			uploadDocuments(client, embeddingApi, documents),
+		deleteDocuments: (path: string) => deleteDocuments(client, path),
+		getStoredMtime: (path: string) =>
+			getStoredMtime(client, embeddingApi, path),
+	};
+};

extensions/void/src/extension.ts

@@ -3,6 +3,8 @@ import { DisplayChangesProvider } from './DisplayChangesProvider';
 import { BaseDiffArea, ChatThreads, MessageFromSidebar, MessageToSidebar } from './common/shared_types';
 import { SidebarWebviewProvider } from './SidebarWebviewProvider';
 import { v4 as uuidv4 } from 'uuid'
+import { embedWorkspaceFiles } from './ai/embed';
+import { getVoidConfig } from './sidebar/contextForConfig';
 
 const readFileContentOfUri = async (uri: vscode.Uri) => {
 	return Buffer.from(await vscode.workspace.fs.readFile(uri)).toString('utf8')
@@ -172,8 +174,10 @@ export function activate(context: vscode.ExtensionContext) {
 		}
 	)
 
-
-
+	// 6. Background jobs
+	const partialVoidConfig = context.globalState.get('partialVoidConfig') ?? {}
+	const voidConfig = getVoidConfig(partialVoidConfig)	
+	embedWorkspaceFiles(voidConfig)
 
 	// Gets called when user presses ctrl + k (mounts ctrl+k-style codelens)
 	// TODO need to build this

extensions/void/src/sidebar/SidebarChat.tsx

@@ -1,7 +1,6 @@
-import React, { FormEvent, useCallback, useEffect, useRef, useState } from "react";
+import React, { FormEvent, useCallback, useRef, useState } from "react";
 
 
-import { marked } from 'marked';
 import MarkdownRender from "./markdown/MarkdownRender";
 import BlockCode from "./markdown/BlockCode";
 import { File, ChatMessage, CodeSelection } from "../common/shared_types";