Migrate to langchain

This commit is contained in:
2026-03-31 13:06:59 +05:30
parent 23dcdd2f4f
commit 8bbd159224
4 changed files with 315 additions and 318 deletions
+63 -97
View File
@@ -1,25 +1,19 @@
import fs from "node:fs"; import fs from "node:fs";
import { OllamaEmbedding } from "@llamaindex/ollama"; import * as lancedbSdk from "@lancedb/lancedb";
import { OpenAI } from "@llamaindex/openai"; import { LanceDB } from "@langchain/community/vectorstores/lancedb";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { DynamicTool } from "@langchain/core/tools";
import { OllamaEmbeddings } from "@langchain/ollama";
import { ChatOpenAI } from "@langchain/openai";
import Graph from "graphology"; import Graph from "graphology";
import { import { AgentExecutor, createReactAgent } from "langchain/agents";
QueryEngineTool,
ReActAgent,
Settings,
storageContextFromDefaults,
VectorStoreIndex,
} from "llamaindex";
import { GRAPH_FILE, STORAGE_DIR } from "./constants"; import { GRAPH_FILE, STORAGE_DIR } from "./constants";
// Configure LlamaIndex to use local Ollama embeddings for semantic search const embeddings = new OllamaEmbeddings({
// This enables the query engine to perform similarity searches without external APIs
Settings.embedModel = new OllamaEmbedding({
model: "nomic-embed-text-v2-moe", model: "nomic-embed-text-v2-moe",
}); });
// Load API configuration from config.json
// Expects NVIDIA_API_KEY and NVIDIA_API_BASE for accessing NVIDIA's API endpoint
const config = JSON.parse(fs.readFileSync("./config.json", "utf-8")); const config = JSON.parse(fs.readFileSync("./config.json", "utf-8"));
const apiKey = config.NVIDIA_API_KEY; const apiKey = config.NVIDIA_API_KEY;
const baseURL = config.NVIDIA_API_BASE; const baseURL = config.NVIDIA_API_BASE;
@@ -28,81 +22,71 @@ if (!apiKey) {
console.warn("Please set NVIDIA_API_KEY in config.json."); console.warn("Please set NVIDIA_API_KEY in config.json.");
} }
// Initialize the LLM using NVIDIA's OpenAI-compatible API endpoint const llm = new ChatOpenAI({
// Model: openai/gpt-oss-120b with temperature 0 for deterministic responses modelName: "openai/gpt-oss-120b",
const llm = new OpenAI({ openAIApiKey: apiKey,
model: "openai/gpt-oss-120b", configuration: { baseURL },
apiKey: apiKey,
baseURL: baseURL,
temperature: 0, temperature: 0,
}); });
Settings.llm = llm;
/** const systemPrompt = `You are an expert in the ECMAScript specification and its implementation in engine262.
* Main function that initializes and runs the ECMAScript specification agent. Your goal is to explain how specific parts of the language work by combining information from the provided tools.
*
* The agent combines two information sources: CRITICAL INSTRUCTIONS:
* 1. Vector search index (spec_retriever) - for semantic text search across spec sections 1. ALWAYS prefer using the provided tools ('spec_retriever', 'fetch_section_chunks', and 'graph_explorer') to answer questions.
* 2. Graph knowledge base (graph_explorer) - for structural relationships between sections and code 2. Do NOT rely on your internal knowledge of JavaScript or the ECMAScript specification.
*/ 3. If the user asks about a function, you MUST first use 'graph_explorer' to find the associated specification section.
4. You MUST then use 'fetch_section_chunks' to read the actual text of that specification section before answering.
5. Base your explanations ONLY on the information retrieved from the tools.
6. If the tools do not provide enough information, state that clearly rather than guessing from your internal knowledge.`;
const prompt = ChatPromptTemplate.fromMessages([
["system", systemPrompt],
["human", "{input}"],
]);
async function main() { async function main() {
console.log("Loading indices and graph..."); console.log("Loading indices and graph...");
// Load the vector index from disk containing embedded spec sections const db = await lancedbSdk.connect(STORAGE_DIR);
const storageContext = await storageContextFromDefaults({ const table = await db.openTable("spec_vectors");
persistDir: STORAGE_DIR, const vectorStore = new LanceDB(embeddings, { table });
});
const index = await VectorStoreIndex.init({
storageContext,
});
// Load the knowledge graph mapping spec sections to implementation functions
const graphData = JSON.parse(fs.readFileSync(GRAPH_FILE, "utf-8")); const graphData = JSON.parse(fs.readFileSync(GRAPH_FILE, "utf-8"));
const graph = new Graph({ multi: true }); const graph = new Graph({ multi: true });
graph.import(graphData); graph.import(graphData);
// Create a query engine with top-3 similarity results for text retrieval const specRetrieverTool = new DynamicTool({
const queryEngine = index.asQueryEngine({ similarityTopK: 3 });
/**
* Tool for retrieving specification text via vector similarity search.
* Used to get detailed content of specific sections based on semantic queries.
*/
const queryEngineTool = new QueryEngineTool({
queryEngine,
metadata: {
name: "spec_retriever", name: "spec_retriever",
description: description:
"Queries the language specification for text content about specific sections or topics. Use this to get the detailed text of a section.", "Queries the language specification for text content about specific sections or topics. Use this to get the detailed text of a section.",
func: async (query: string) => {
const results = await vectorStore.similaritySearch(query, 3);
return results.map((r) => r.pageContent).join("\n\n");
}, },
}); });
/** const sectionRetrieverTool = new DynamicTool({
* Tool for exploring the knowledge graph connecting spec sections to implementation. name: "fetch_section_chunks",
* Enables structural navigation: finding which spec section a function implements description:
* or which functions implement a spec section. "Retrieves all text chunks from a specific specification section by sectionId. Use after finding a sectionId via spec_retriever.",
*/ func: async (sectionId: string) => {
const graphTool = { const results = await table
metadata: { .query()
.where(`sectionid = '${sectionId}'`)
.limit(100)
.toArray();
return results.map((r: { text: string }) => r.text).join("\n\n");
},
});
const graphTool = new DynamicTool({
name: "graph_explorer", name: "graph_explorer",
description: description:
"Explores structural relationships between specification sections and implementation code (functions). Use this to find which spec section a function implements. Input: section ID or function name.", "Explores structural relationships between specification sections and implementation code (functions). Use this to find which spec section a function implements. Input: section ID or function name.",
parameters: { func: async (query: string) => {
type: "object",
properties: {
query: {
type: "string",
description: "The section ID or function name to explore.",
},
},
required: ["query"],
},
},
call: async ({ query }: { query: string }) => {
console.log(`[Tool: graph_explorer] Querying for: ${query}`); console.log(`[Tool: graph_explorer] Querying for: ${query}`);
// Try exact node match, or prepend 'func-' prefix for function names
let nodeId = query; let nodeId = query;
if (!graph.hasNode(nodeId)) { if (!graph.hasNode(nodeId)) {
if (graph.hasNode(`func-${query}`)) { if (graph.hasNode(`func-${query}`)) {
@@ -111,7 +95,6 @@ async function main() {
} }
if (graph.hasNode(nodeId)) { if (graph.hasNode(nodeId)) {
// Collect node info and all connected nodes
const neighbors = graph.neighbors(nodeId); const neighbors = graph.neighbors(nodeId);
const nodeAttr = graph.getNodeAttributes(nodeId); const nodeAttr = graph.getNodeAttributes(nodeId);
@@ -120,7 +103,6 @@ async function main() {
if (nodeAttr.file) result += `- File: ${nodeAttr.file}\n`; if (nodeAttr.file) result += `- File: ${nodeAttr.file}\n`;
result += `\nConnected parts:\n`; result += `\nConnected parts:\n`;
// List all connected nodes with their relationship types
neighbors.forEach((neighbor) => { neighbors.forEach((neighbor) => {
const attr = graph.getNodeAttributes(neighbor); const attr = graph.getNodeAttributes(neighbor);
const edges = graph.edges(nodeId, neighbor); const edges = graph.edges(nodeId, neighbor);
@@ -133,48 +115,32 @@ async function main() {
return `No information found in graph for ${query}. Use spec_retriever to search text.`; return `No information found in graph for ${query}. Use spec_retriever to search text.`;
}, },
}; });
/** const agent = await createReactAgent({
* ReAct agent that reasons about ECMAScript specification. llm,
* tools: [specRetrieverTool, sectionRetrieverTool, graphTool],
* The agent follows this workflow: prompt,
* 1. For function queries: graph_explorer → spec_retriever → explanation });
* 2. For section queries: spec_retriever → explanation
*
* Critical constraints ensure tool-based answers rather than internal knowledge.
*/
const agent = new ReActAgent({
tools: [queryEngineTool, graphTool],
llm: llm,
verbose: true,
systemPrompt: `You are an expert in the ECMAScript specification and its implementation in engine262.
Your goal is to explain how specific parts of the language work by combining information from the provided tools.
CRITICAL INSTRUCTIONS: const agentExecutor = new AgentExecutor({
1. ALWAYS prefer using the provided tools ('spec_retriever' and 'graph_explorer') to answer questions. agent,
2. Do NOT rely on your internal knowledge of JavaScript or the ECMAScript specification. tools: [specRetrieverTool, sectionRetrieverTool, graphTool],
3. If the user asks about a function, you MUST first use 'graph_explorer' to find the associated specification section.
4. You MUST then use 'spec_retriever' to read the actual text of that specification section before answering.
5. Base your explanations ONLY on the information retrieved from the tools.
6. If the tools do not provide enough information, state that clearly rather than guessing from your internal knowledge.`,
}); });
console.log("Agent is ready!"); console.log("Agent is ready!");
// Accept user query from command line argument, or use default question
const message = const message =
process.argv[2] || process.argv[2] ||
"Which spec section does Evaluate_IfStatement implement? and what does that section say?"; "Which spec section does Evaluate_IfStatement implement? and what does that section say?";
console.log(`User: ${message}`); console.log(`User: ${message}`);
// Execute the agent with the user's query const response = await agentExecutor.invoke({
const response = await agent.chat({ input: message,
message: message,
}); });
console.log("\n--- Agent Response ---\n"); console.log("\n--- Agent Response ---\n");
console.log(response.toString()); console.log(response.output);
} }
main().catch(console.error); main().catch(console.error);
+191 -49
View File
File diff suppressed because one or more lines are too long
+6 -6
View File
@@ -19,16 +19,16 @@
"description": "", "description": "",
"type": "module", "type": "module",
"dependencies": { "dependencies": {
"@llamaindex/core": "^0.6.22", "@lancedb/lancedb": "^0.5.0",
"@llamaindex/env": "^0.1.30", "@langchain/community": "0.2.0",
"@llamaindex/node-parser": "^2.0.22", "@langchain/core": "^0.2.0",
"@llamaindex/ollama": "^0.1.23", "@langchain/ollama": "^0.1.0",
"@llamaindex/openai": "^0.4.22", "@langchain/openai": "^0.1.0",
"acorn": "^8.16.0", "acorn": "^8.16.0",
"cheerio": "^1.2.0", "cheerio": "^1.2.0",
"glob": "^13.0.6", "glob": "^13.0.6",
"graphology": "^0.26.0", "graphology": "^0.26.0",
"llamaindex": "^0.12.1" "langchain": "^0.2.0"
}, },
"devDependencies": { "devDependencies": {
"typescript": "^5.5.2", "typescript": "^5.5.2",
+50 -161
View File
@@ -1,37 +1,29 @@
import fs from "node:fs"; import fs from "node:fs";
import path from "node:path"; import path from "node:path";
import readline from "node:readline"; 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 * as cheerio from "cheerio";
import { glob } from "glob"; import { glob } from "glob";
import {
Document,
SentenceSplitter,
Settings,
storageContextFromDefaults,
VectorStoreIndex,
} from "llamaindex";
import { SPEC_DIR, STORAGE_DIR } from "../constants"; import { SPEC_DIR, STORAGE_DIR } from "../constants";
// Configure LlamaIndex to use local Ollama embeddings const embeddings = new OllamaEmbeddings({
// This creates vector embeddings for semantic search without external APIs
Settings.embedModel = new OllamaEmbedding({
model: "nomic-embed-text-v2-moe", model: "nomic-embed-text-v2-moe",
}); });
// Text chunking configuration for larger chunks with more context preservation const textSplitter = new RecursiveCharacterTextSplitter({
// Larger chunks reduce total number of nodes while still fitting within chunkSize: 4096,
// the embedding model's 8192 token limit (~2048 chars ≈ 512 tokens) chunkOverlap: 100,
const sentenceSplitter = new SentenceSplitter({ separators: ["\n\n", "\n", ". ", " ", ""],
chunkSize: 2048,
chunkOverlap: 50,
}); });
/** const BREAKDOWN_TAGS = ["emu-table", "emu-grammar"] as const;
* Prompts the user for confirmation via stdin. const LARGE_DOC_THRESHOLD = 5000;
* @param question - The question to display to the user
* @returns Promise that resolves to true if user confirms (yes/y), false otherwise
*/
function askUser(question: string): Promise<boolean> { function askUser(question: string): Promise<boolean> {
const rl = readline.createInterface({ const rl = readline.createInterface({
input: process.stdin, input: process.stdin,
@@ -47,24 +39,7 @@ function askUser(question: string): Promise<boolean> {
}); });
} }
// Tags to extract from large sections for finer-grained chunking async function ingestSpec(): Promise<Document[]> {
// 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() {
const htmlFiles = await glob(path.join(SPEC_DIR, "*.html")); const htmlFiles = await glob(path.join(SPEC_DIR, "*.html"));
const documents: Document[] = []; const documents: Document[] = [];
@@ -75,7 +50,6 @@ async function ingestSpec() {
$("emu-clause").each((_i, elem) => { $("emu-clause").each((_i, elem) => {
const id = $(elem).attr("id"); const id = $(elem).attr("id");
const title = $(elem).find("h1").first().text().trim(); const title = $(elem).find("h1").first().text().trim();
// Only extract immediate text to avoid excessive chunking of child sections
const text = $(elem) const text = $(elem)
.clone() .clone()
.children("emu-clause") .children("emu-clause")
@@ -88,7 +62,7 @@ async function ingestSpec() {
return; 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) { if (text.length > LARGE_DOC_THRESHOLD) {
let subDocsCreated = false; let subDocsCreated = false;
const $section = $(elem).clone(); const $section = $(elem).clone();
@@ -105,14 +79,14 @@ async function ingestSpec() {
if (subText) { if (subText) {
documents.push( documents.push(
new Document({ new Document({
text: subText, pageContent: subText,
metadata: { metadata: {
source: file, source: file,
sectionId: subId, sectionid: subId,
sectionTitle: `${title} [${tagName}]`, sectiontitle: `${title} [${tagName}]`,
type: "specification", type: "specification",
parentSectionId: id, parentsectionid: id,
breakdownTag: tagName, breakdowntag: tagName,
}, },
}), }),
); );
@@ -131,14 +105,14 @@ async function ingestSpec() {
if (remainingText) { if (remainingText) {
documents.push( documents.push(
new Document({ new Document({
text: remainingText, pageContent: remainingText,
metadata: { metadata: {
source: file, source: file,
sectionId: `${id}-prose-part-1`, sectionid: `${id}-prose-part-1`,
sectionTitle: `${title} [prose]`, sectiontitle: `${title} [prose]`,
type: "specification", type: "specification",
parentSectionId: id, parentsectionid: id,
breakdownTag: "prose", breakdowntag: "prose",
}, },
}), }),
); );
@@ -148,18 +122,19 @@ async function ingestSpec() {
if (subDocsCreated || remainingText) { if (subDocsCreated || remainingText) {
return; return;
} }
// Otherwise, fall through to add the full section document
} }
// Add the full section document (for smaller sections or when no breakdown happened) // Add the full section document (for smaller sections or when no breakdown happened)
documents.push( documents.push(
new Document({ new Document({
text, pageContent: text,
metadata: { metadata: {
source: file, source: file,
sectionId: id, sectionid: id,
sectionTitle: title, sectiontitle: title,
type: "specification", type: "specification",
parentsectionid: null,
breakdowntag: null,
}, },
}), }),
); );
@@ -168,99 +143,21 @@ async function ingestSpec() {
return documents; 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() { async function main() {
console.log("Ingesting specification..."); console.log("Ingesting specification...");
const specDocs = await ingestSpec(); const specDocs = await ingestSpec();
console.log(`Ingested ${specDocs.length} specification sections.`); 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 db = await lancedbSdk.connect(STORAGE_DIR);
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);
if (largeDocs.length > 0) { let table: lancedbSdk.Table;
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
try { try {
index = await VectorStoreIndex.init({ table = await db.openTable("spec_vectors");
storageContext, console.log("Existing table found.");
});
console.log("Existing index found.");
const shouldOverwrite = await askUser( const shouldOverwrite = await askUser(
"Do you want to overwrite the existing vector store?", "Do you want to overwrite the existing vector store?",
); );
@@ -268,29 +165,21 @@ async function main() {
console.log("Ingest cancelled by user."); console.log("Ingest cancelled by user.");
process.exit(0); process.exit(0);
} }
console.log("Overwriting existing index..."); console.log("Overwriting existing table...");
// Reset index to null so we create a fresh one await db.dropTable("spec_vectors");
index = null; table = await db.createTable("spec_vectors", []);
} catch (_e) { } catch {
console.log("No existing index found, starting fresh."); console.log("No existing table found, creating fresh...");
table = await db.createTable("spec_vectors", []);
} }
// Process nodes in batches to avoid overwhelming the embedding service console.log("Creating scalar indexes...");
for (let i = 0; i < nodes.length; i += BATCH_SIZE) { await table.createIndex("sectionid", { config: Index.btree() });
const batch = nodes.slice(i, i + BATCH_SIZE); await table.createIndex("type", { config: Index.btree() });
console.log(
`Processing batch ${i / BATCH_SIZE + 1} / ${Math.ceil(nodes.length / BATCH_SIZE)}...`,
);
if (!index) { console.log("Storing documents with embeddings...");
index = await VectorStoreIndex.init({ const vectorStore = new LanceDB(embeddings, { table });
storageContext, await vectorStore.addDocuments(splitDocs);
nodes: batch,
});
} else {
await index.insertNodes(batch);
}
}
console.log(`Index built and persisted to ${STORAGE_DIR}`); console.log(`Index built and persisted to ${STORAGE_DIR}`);
} }