Migrate to langchain

setup/ingest.ts

@@ -1,37 +1,29 @@
 import fs from "node:fs";
 import path from "node:path";
 import readline from "node:readline";
-import { OllamaEmbedding } from "@llamaindex/ollama";
+import * as lancedbSdk from "@lancedb/lancedb";
+import { Index } from "@lancedb/lancedb";
+import { LanceDB } from "@langchain/community/vectorstores/lancedb";
+import { Document } from "@langchain/core/documents";
+import { OllamaEmbeddings } from "@langchain/ollama";
+import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
 import * as cheerio from "cheerio";
 import { glob } from "glob";
-import {
-  Document,
-  SentenceSplitter,
-  Settings,
-  storageContextFromDefaults,
-  VectorStoreIndex,
-} from "llamaindex";
 import { SPEC_DIR, STORAGE_DIR } from "../constants";
 
-// Configure LlamaIndex to use local Ollama embeddings
-// This creates vector embeddings for semantic search without external APIs
-Settings.embedModel = new OllamaEmbedding({
+const embeddings = new OllamaEmbeddings({
   model: "nomic-embed-text-v2-moe",
 });
 
-// Text chunking configuration for larger chunks with more context preservation
-// Larger chunks reduce total number of nodes while still fitting within
-// the embedding model's 8192 token limit (~2048 chars ≈ 512 tokens)
-const sentenceSplitter = new SentenceSplitter({
-  chunkSize: 2048,
-  chunkOverlap: 50,
+const textSplitter = new RecursiveCharacterTextSplitter({
+  chunkSize: 4096,
+  chunkOverlap: 100,
+  separators: ["\n\n", "\n", ". ", " ", ""],
 });
 
-/**
- * Prompts the user for confirmation via stdin.
- * @param question - The question to display to the user
- * @returns Promise that resolves to true if user confirms (yes/y), false otherwise
- */
+const BREAKDOWN_TAGS = ["emu-table", "emu-grammar"] as const;
+const LARGE_DOC_THRESHOLD = 5000;
+
 function askUser(question: string): Promise<boolean> {
   const rl = readline.createInterface({
     input: process.stdin,
@@ -47,24 +39,7 @@ function askUser(question: string): Promise<boolean> {
   });
 }
 
-// Tags to extract from large sections for finer-grained chunking
-// Extend this array to add more tag types for breakdown
-const BREAKDOWN_TAGS = ["emu-table", "emu-grammar"] as const;
-const LARGE_DOC_THRESHOLD = 5000;
-
-/**
- * Extracts ECMAScript specification sections from HTML files and converts them
- * to Documents for vector indexing. Each section (emu-clause) becomes a separate
- * document with metadata for tracking.
- *
- * For large sections (> 5000 chars), attempts to break them down by extracting
- * content from specific structural tags (emu-table, emu-grammar, etc.) to create
- * more focused chunks. Falls back to the full section text if no breakdown tags
- * are found.
- *
- * @returns Array of Documents ready for indexing
- */
-async function ingestSpec() {
+async function ingestSpec(): Promise<Document[]> {
   const htmlFiles = await glob(path.join(SPEC_DIR, "*.html"));
   const documents: Document[] = [];
 
@@ -75,7 +50,6 @@ async function ingestSpec() {
     $("emu-clause").each((_i, elem) => {
       const id = $(elem).attr("id");
       const title = $(elem).find("h1").first().text().trim();
-      // Only extract immediate text to avoid excessive chunking of child sections
       const text = $(elem)
         .clone()
         .children("emu-clause")
@@ -88,7 +62,7 @@ async function ingestSpec() {
         return;
       }
 
-      // For large documents, attempt to break down by structural tags
+      // For large documents, break down by structural tags
       if (text.length > LARGE_DOC_THRESHOLD) {
         let subDocsCreated = false;
         const $section = $(elem).clone();
@@ -105,14 +79,14 @@ async function ingestSpec() {
             if (subText) {
               documents.push(
                 new Document({
-                  text: subText,
+                  pageContent: subText,
                   metadata: {
                     source: file,
-                    sectionId: subId,
-                    sectionTitle: `${title} [${tagName}]`,
+                    sectionid: subId,
+                    sectiontitle: `${title} [${tagName}]`,
                     type: "specification",
-                    parentSectionId: id,
-                    breakdownTag: tagName,
+                    parentsectionid: id,
+                    breakdowntag: tagName,
                   },
                 }),
               );
@@ -131,14 +105,14 @@ async function ingestSpec() {
         if (remainingText) {
           documents.push(
             new Document({
-              text: remainingText,
+              pageContent: remainingText,
               metadata: {
                 source: file,
-                sectionId: `${id}-prose-part-1`,
-                sectionTitle: `${title} [prose]`,
+                sectionid: `${id}-prose-part-1`,
+                sectiontitle: `${title} [prose]`,
                 type: "specification",
-                parentSectionId: id,
-                breakdownTag: "prose",
+                parentsectionid: id,
+                breakdowntag: "prose",
               },
             }),
           );
@@ -148,18 +122,19 @@ async function ingestSpec() {
         if (subDocsCreated || remainingText) {
           return;
         }
-        // Otherwise, fall through to add the full section document
       }
 
       // Add the full section document (for smaller sections or when no breakdown happened)
       documents.push(
         new Document({
-          text,
+          pageContent: text,
           metadata: {
             source: file,
-            sectionId: id,
-            sectionTitle: title,
+            sectionid: id,
+            sectiontitle: title,
             type: "specification",
+            parentsectionid: null,
+            breakdowntag: null,
           },
         }),
       );
@@ -168,99 +143,21 @@ async function ingestSpec() {
   return documents;
 }
 
-/**
- * Main execution pipeline:
- * 1. Ingest specification HTML files and convert to documents
- * 2. Split documents into smaller text chunks (nodes)
- * 3. Filter out oversized chunks that could exceed LLM context limits
- * 4. Build a vector index in batches to handle large document sets
- * 5. Persist the index to disk for later retrieval
- */
 async function main() {
   console.log("Ingesting specification...");
   const specDocs = await ingestSpec();
   console.log(`Ingested ${specDocs.length} specification sections.`);
 
-  console.log("Splitting documents into nodes...");
+  console.log("Splitting documents into chunks...");
+  const splitDocs = await textSplitter.splitDocuments(specDocs);
+  console.log(`Total chunks generated: ${splitDocs.length}`);
 
-  // Debug: Log largest documents (over 2000 chars) to diagnose oversized nodes
-  const largeDocs = specDocs
-    .map((doc, i) => ({
-      index: i,
-      length: doc.text.length,
-      sectionId: doc.metadata.sectionId,
-    }))
-    .filter((doc) => doc.length > 2000)
-    .sort((a, b) => b.length - a.length)
-    .slice(0, 50);
+  const db = await lancedbSdk.connect(STORAGE_DIR);
 
-  if (largeDocs.length > 0) {
-    console.log("\nDebug: Largest documents (> 2000 chars):");
-    largeDocs.forEach((doc) => {
-      console.log(
-        `  Doc ${doc.index}: ${doc.length} chars, section: ${doc.sectionId}`,
-      );
-    });
-    const remaining =
-      specDocs.filter((doc) => doc.text.length > 2000).length -
-      largeDocs.length;
-    if (remaining > 0) {
-      console.log(`  ... and ${remaining} more large documents`);
-    }
-  } else {
-    console.log("\nDebug: No documents over 2000 chars found");
-  }
-
-  const rawNodes = sentenceSplitter.getNodesFromDocuments(specDocs);
-  console.log(`Total raw nodes generated: ${rawNodes.length}`);
-
-  // Debug: Log node size distribution
-  const nodeSizes = rawNodes.map((n) => n.getContent().length);
-  const maxNodeSize = Math.max(...nodeSizes);
-  const avgNodeSize = nodeSizes.reduce((a, b) => a + b, 0) / nodeSizes.length;
-  console.log(
-    `\nDebug: Node size stats - Max: ${maxNodeSize}, Avg: ${Math.round(avgNodeSize)}`,
-  );
-
-  // Safety filter to ensure no node exceeds context limit
-  // Filter threshold set to chunkSize + buffer for metadata overhead
-  const MAX_NODE_LENGTH = 2500;
-  let skippedCount = 0;
-  const nodes = rawNodes.filter((node) => {
-    const contentLen = node.getContent().length;
-    if (contentLen > MAX_NODE_LENGTH) {
-      skippedCount++;
-      if (skippedCount <= 3) {
-        console.warn(
-          `Skipping node with length ${contentLen} from ${node.metadata.source || "unknown"} (section: ${node.metadata.sectionId})`,
-        );
-      }
-      return false;
-    }
-    return true;
-  });
-
-  if (skippedCount > 3) {
-    console.warn(`  ... and ${skippedCount - 3} more nodes skipped`);
-  }
-  console.log(`Total valid nodes for indexing: ${nodes.length}`);
-
-  console.log("Creating storage context...");
-  const storageContext = await storageContextFromDefaults({
-    persistDir: STORAGE_DIR,
-  });
-
-  console.log("Building index (this might take a while with local Ollama)...");
-
-  const BATCH_SIZE = 50;
-  let index: VectorStoreIndex | null = null;
-
-  // Try to load existing index if any
+  let table: lancedbSdk.Table;
   try {
-    index = await VectorStoreIndex.init({
-      storageContext,
-    });
-    console.log("Existing index found.");
+    table = await db.openTable("spec_vectors");
+    console.log("Existing table found.");
     const shouldOverwrite = await askUser(
       "Do you want to overwrite the existing vector store?",
     );
@@ -268,29 +165,21 @@ async function main() {
       console.log("Ingest cancelled by user.");
       process.exit(0);
     }
-    console.log("Overwriting existing index...");
-    // Reset index to null so we create a fresh one
-    index = null;
-  } catch (_e) {
-    console.log("No existing index found, starting fresh.");
+    console.log("Overwriting existing table...");
+    await db.dropTable("spec_vectors");
+    table = await db.createTable("spec_vectors", []);
+  } catch {
+    console.log("No existing table found, creating fresh...");
+    table = await db.createTable("spec_vectors", []);
   }
 
-  // Process nodes in batches to avoid overwhelming the embedding service
-  for (let i = 0; i < nodes.length; i += BATCH_SIZE) {
-    const batch = nodes.slice(i, i + BATCH_SIZE);
-    console.log(
-      `Processing batch ${i / BATCH_SIZE + 1} / ${Math.ceil(nodes.length / BATCH_SIZE)}...`,
-    );
+  console.log("Creating scalar indexes...");
+  await table.createIndex("sectionid", { config: Index.btree() });
+  await table.createIndex("type", { config: Index.btree() });
 
-    if (!index) {
-      index = await VectorStoreIndex.init({
-        storageContext,
-        nodes: batch,
-      });
-    } else {
-      await index.insertNodes(batch);
-    }
-  }
+  console.log("Storing documents with embeddings...");
+  const vectorStore = new LanceDB(embeddings, { table });
+  await vectorStore.addDocuments(splitDocs);
 
   console.log(`Index built and persisted to ${STORAGE_DIR}`);
 }