VS Code extensions

import * as path from "node:path";
import * as lancedb from "@lancedb/lancedb";
import "@lancedb/lancedb/embedding/transformers";
import { Int32, Utf8 } from "apache-arrow";

// Pass an absolute directory, such as `context.storageUri.fsPath`. A relative
// path resolves against the extension host's working directory, which is not
// your workspace and may be read-only (on macOS it is usually `/`).
export async function openDatabase(storageDir: string) {
  return lancedb.connect(path.join(storageDir, "lancedb"));
}

export type Chunk = { path: string; line: number; text: string };

// Split a Markdown file into paragraphs, keeping the line each one starts on
// so a search result can jump straight to it.
export function chunkMarkdown(filePath: string, content: string): Chunk[] {
  const lines = content.split(/\r?\n/);
  const chunks: Chunk[] = [];
  let start = 0;
  for (let i = 0; i <= lines.length; i++) {
    if (i === lines.length || lines[i].trim() === "") {
      const text = lines.slice(start, i).join("\n").trim();
      if (text) {
        chunks.push({ path: filePath, line: start, text });
      }
      start = i + 1;
    }
  }
  return chunks;
}

export async function indexChunks(db: lancedb.Connection, chunks: Chunk[]) {
  // Runs all-MiniLM-L6-v2 locally through Transformers.js: no API key needed.
  const embedder = (await lancedb.embedding
    .getRegistry()
    .get("huggingface")
    ?.create()) as lancedb.embedding.EmbeddingFunction;

  const schema = lancedb.embedding.LanceSchema({
    path: new Utf8(),
    line: new Int32(),
    text: embedder.sourceField(new Utf8()),
    vector: embedder.vectorField(),
  });

  // LanceDB fills in the `vector` column from `text` as the rows are written.
  return db.createTable("chunks", chunks, { schema, mode: "overwrite" });
}

export async function searchChunks(
  db: lancedb.Connection,
  query: string,
  limit = 5,
) {
  // The table stores its embedding function in its metadata, so a text query
  // is embedded with the same model that embedded the rows.
  const table = await db.openTable("chunks");
  return table
    .search(query)
    .select(["path", "line", "text", "_distance"])
    .limit(limit)
    .toArray();
}

"contributes": {
  "commands": [
    {
      "command": "lancedbSearch.index",
      "title": "Index Markdown in Workspace",
      "category": "LanceDB"
    },
    {
      "command": "lancedbSearch.search",
      "title": "Search Markdown",
      "category": "LanceDB"
    }
  ]
}

import * as vscode from "vscode";
import {
  type Chunk,
  chunkMarkdown,
  indexChunks,
  openDatabase,
  searchChunks,
} from "./search";

export async function activate(context: vscode.ExtensionContext) {
  // One index per workspace. `storageUri` is undefined when no folder is open.
  const storage = context.storageUri ?? context.globalStorageUri;

  // Transformers.js downloads the embedding model on first use. Keep it in
  // global storage so every workspace shares one copy and it survives
  // extension updates. Use `import()`: LanceDB loads the ES module build, and
  // `require()` would return a separate copy whose settings LanceDB never sees.
  const { env } = await import("@huggingface/transformers");
  env.cacheDir = vscode.Uri.joinPath(context.globalStorageUri, "models").fsPath;

  context.subscriptions.push(
    vscode.commands.registerCommand("lancedbSearch.index", async () => {
      const files = await vscode.workspace.findFiles(
        "**/*.md",
        "**/node_modules/**",
      );
      const chunks: Chunk[] = [];
      for (const file of files) {
        const bytes = await vscode.workspace.fs.readFile(file);
        // Store the file's URI, not a path relative to its folder: in a
        // multi-root workspace `findFiles` searches every folder, and only the
        // URI says which one a result came from.
        chunks.push(
          ...chunkMarkdown(file.toString(), new TextDecoder().decode(bytes)),
        );
      }
      if (chunks.length === 0) {
        vscode.window.showWarningMessage("No Markdown files to index.");
        return 0;
      }

      await vscode.window.withProgress(
        {
          location: vscode.ProgressLocation.Notification,
          title: `Indexing ${files.length} Markdown files`,
        },
        async () => {
          await vscode.workspace.fs.createDirectory(storage);
          await indexChunks(await openDatabase(storage.fsPath), chunks);
        },
      );
      vscode.window.showInformationMessage(
        `Indexed ${chunks.length} paragraphs from ${files.length} files.`,
      );
      return chunks.length;
    }),

    // Called with a query (for example from a test or another extension), the
    // command returns the matching rows instead of showing a picker.
    vscode.commands.registerCommand(
      "lancedbSearch.search",
      async (query?: string) => {
        const text =
          query ??
          (await vscode.window.showInputBox({
            prompt: "Search Markdown by meaning",
          }));
        if (!text) {
          return [];
        }
        const results = await searchChunks(
          await openDatabase(storage.fsPath),
          text,
        );
        if (query !== undefined) {
          return results;
        }

        const picked = await vscode.window.showQuickPick(
          results.map((row) => {
            const uri = vscode.Uri.parse(row.path);
            return {
              label: `${vscode.workspace.asRelativePath(uri)}:${row.line + 1}`,
              description: `distance ${row._distance.toFixed(3)}`,
              detail: row.text,
              uri,
              line: row.line,
            };
          }),
          { placeHolder: text, matchOnDetail: true },
        );
        if (picked) {
          const editor = await vscode.window.showTextDocument(picked.uri);
          const start = new vscode.Position(picked.line, 0);
          editor.selection = new vscode.Selection(start, start);
          editor.revealRange(
            new vscode.Range(start, start),
            vscode.TextEditorRevealType.InCenter,
          );
        }
        return results;
      },
    ),
  );
}

export function deactivate() {}