@devilsdev/rag-pipeline-utils 2.4.5 → 2.6.0

package/src/loader/csv-loader.js

@@ -4,7 +4,8 @@
 
 const fs = require('fs/promises');
 const path = require('path');
-const { parse } = require('csv-parse/sync');
+// csv-parse loads when a CSV file is read, not on import.
+const loadCsvParse = () => require('csv-parse/sync'); // eslint-disable-line global-require
 
 class CSVLoader {
   /**
@@ -20,7 +21,7 @@ class CSVLoader {
       const src = raw.replace(/^\uFEFF/, ''); // Remove BOM
       const normalizedSrc = src.replace(/\r\n|\r/g, '\n'); // Normalize line endings
 
-      const records = parse(normalizedSrc, { columns: true });
+      const records = loadCsvParse().parse(normalizedSrc, { columns: true });
 
       return [
         {

package/src/loader/document-reader.js

@@ -0,0 +1,91 @@
+'use strict';
+
+const fs = require('fs');
+const path = require('path');
+const fg = require('fast-glob');
+const { MarkdownLoader } = require('./markdown-loader.js');
+const { HTMLLoader } = require('./html-loader.js');
+const { CSVLoader } = require('./csv-loader.js');
+
+const textLoader = {
+  async load(filePath) {
+    const text = await fs.promises.readFile(filePath, 'utf8');
+    return [
+      {
+        chunk: () => {
+          const { ChunkingEngine } = require('../chunking/chunking-engine');
+          return new ChunkingEngine({
+            strategy: 'sentence',
+            maxLen: 500,
+          }).chunk(text);
+        },
+      },
+    ];
+  },
+};
+
+const LOADERS_BY_EXTENSION = {
+  '.md': new MarkdownLoader(),
+  '.markdown': new MarkdownLoader(),
+  '.html': new HTMLLoader(),
+  '.htm': new HTMLLoader(),
+  '.csv': new CSVLoader(),
+  '.txt': textLoader,
+};
+
+const SUPPORTED_EXTENSIONS = Object.keys(LOADERS_BY_EXTENSION);
+
+/**
+ * Read a file or every supported file under a directory and return its
+ * chunks together with the file they came from. Unlike DirectoryLoader this
+ * keeps the source path, which ingestion needs for stable chunk ids,
+ * re-ingestion and citations.
+ *
+ * @param {string} inputPath - File or directory
+ * @param {object} [options]
+ * @param {string} [options.baseDir=process.cwd()] - Sources are reported relative to this
+ * @returns {Promise<Array<{source: string, chunks: string[]}>>}
+ */
+async function readDocuments(inputPath, options = {}) {
+  const baseDir = path.resolve(options.baseDir || process.cwd());
+  const absolute = path.resolve(inputPath);
+  if (!fs.existsSync(absolute)) {
+    throw new Error(`Path not found: ${inputPath}`);
+  }
+
+  let files;
+  if (fs.statSync(absolute).isDirectory()) {
+    const pattern = `**/*{${SUPPORTED_EXTENSIONS.join(',')}}`;
+    files = (
+      await fg(pattern, {
+        cwd: absolute,
+        absolute: true,
+        caseSensitiveMatch: false,
+        ignore: ['**/node_modules/**', '**/.git/**'],
+      })
+    ).sort();
+  } else {
+    const ext = path.extname(absolute).toLowerCase();
+    if (!LOADERS_BY_EXTENSION[ext]) {
+      throw new Error(
+        `Unsupported file type "${ext}". Supported: ${SUPPORTED_EXTENSIONS.join(', ')}`,
+      );
+    }
+    files = [absolute];
+  }
+
+  const documents = [];
+  for (const file of files) {
+    const loader = LOADERS_BY_EXTENSION[path.extname(file).toLowerCase()];
+    const loaded = await loader.load(file);
+    const chunks = loaded
+      .flatMap((doc) => doc.chunk())
+      .map((chunk) => String(chunk).trim())
+      .filter(Boolean);
+    const source = path.relative(baseDir, file).split(path.sep).join('/');
+    documents.push({ source, chunks });
+  }
+  return documents;
+}
+
+module.exports = { readDocuments, SUPPORTED_EXTENSIONS };

package/src/loader/html-loader.js

@@ -4,7 +4,8 @@
 
 const fs = require('fs/promises');
 const path = require('path');
-const { JSDOM } = require('jsdom');
+// jsdom (about 100 packages) loads when an HTML file is read, not on import.
+const loadJsdom = () => require('jsdom'); // eslint-disable-line global-require
 
 /**
  * HTMLLoader reads HTML files and extracts visible text.
@@ -19,6 +20,7 @@ class HTMLLoader {
   async load(filePath) {
     const absPath = path.resolve(filePath);
     const raw = await fs.readFile(absPath, 'utf-8');
+    const { JSDOM } = loadJsdom();
     const dom = new JSDOM(raw);
     let text = dom.window.document.body.textContent || '';
 

package/src/loader/markdown-loader.js

@@ -10,7 +10,8 @@ const fs = require('fs/promises');
 // eslint-disable-line global-require
 const path = require('path');
 // eslint-disable-line global-require
-const { marked } = require('marked');
+// marked loads when a Markdown file is read, not on import.
+const loadMarked = () => require('marked'); // eslint-disable-line global-require
 // eslint-disable-line global-require
 
 /**
@@ -26,7 +27,7 @@ class MarkdownLoader {
   async load(_filePath) {
     const absPath = path.resolve(_filePath);
     const raw = await fs.readFile(absPath, 'utf-8');
-    const html = marked(raw);
+    const html = loadMarked().marked(raw);
     const text = html.replace(/<[^>]+>/g, ''); // strip tags
 
     return [

package/src/reranker/cascade-reranker.js

@@ -8,6 +8,7 @@
  */
 
 const { EventEmitter } = require('events');
+const { throwIfAborted } = require('../connectors/connector-utils');
 
 /** @type {object} */
 const DEFAULT_CONFIG = {};
@@ -37,6 +38,7 @@ class CascadeReranker extends EventEmitter {
    * @param {Array<object|string>} documents - Documents to rerank
    * @param {object} [options]
    * @param {number} [options.topK] - Maximum number of results to return
+   * @param {AbortSignal} [options.signal] - Forwarded to every stage; checked between stages
    * @returns {Promise<object[]>} Reranked documents from the final stage
    * @throws {Error} If no rerankers are configured
    */
@@ -52,6 +54,7 @@ class CascadeReranker extends EventEmitter {
     let currentDocs = documents;
 
     for (let i = 0; i < this.rerankers.length; i++) {
+      throwIfAborted(options.signal);
       const reranker = this.rerankers[i];
       currentDocs = await reranker.rerank(query, currentDocs, options);
 

package/src/reranker/embedding-reranker.js

@@ -1,28 +1,53 @@
 'use strict';
 
 /**
- * Version: 1.0.0
+ * Version: 1.1.0
  * Path: /src/reranker/embedding-reranker.js
  * Description: Cosine similarity reranker using any embedder
  * Author: Ali Kahwaji
  */
 
 const { EventEmitter } = require('events');
+const { createHash } = require('crypto');
+const {
+  isAbortError,
+  throwIfAborted,
+  raceAbort,
+} = require('../connectors/connector-utils');
+
+/** @type {{ cacheSize: number, concurrency: number }} */
+const DEFAULT_CONFIG = {
+  // Maximum number of document embeddings kept in the LRU cache (0 disables).
+  cacheSize: 1000,
+  // Maximum number of document embedding calls in flight at once.
+  concurrency: 4,
+};
 
-/** @type {object} */
-const DEFAULT_CONFIG = {};
+/**
+ * Hash document text so cache keys stay small regardless of document size.
+ * @param {string} text
+ * @returns {string}
+ */
+function contentHash(text) {
+  return createHash('sha256').update(text, 'utf8').digest('hex');
+}
 
 /**
  * Reranks documents by computing cosine similarity between query and
  * document embeddings produced by an external embedder.
  * Implements the reranker contract: rerank(query, documents, options).
  *
+ * Document embeddings are cached in a bounded LRU keyed by a SHA-256 of the
+ * document text, and embedding calls run with bounded concurrency.
+ *
  * @extends EventEmitter
  */
 class EmbeddingReranker extends EventEmitter {
   /**
    * @param {object} embedder - Must expose .embedQuery(text) or .embed(text)
    * @param {object} [options]
+   * @param {number} [options.cacheSize=1000] - Max cached document embeddings (0 disables caching)
+   * @param {number} [options.concurrency=4] - Max concurrent embedding calls
    */
   constructor(embedder, options = {}) {
     super();
@@ -31,6 +56,79 @@ class EmbeddingReranker extends EventEmitter {
     }
     this.embedder = embedder;
     this.config = { ...DEFAULT_CONFIG, ...options };
+    if (
+      !(
+        Number.isInteger(this.config.concurrency) && this.config.concurrency > 0
+      )
+    ) {
+      throw new Error(
+        'EmbeddingReranker: concurrency must be a positive integer',
+      );
+    }
+    /** @type {Map<string, Promise<number[]>>} LRU: oldest entry first */
+    this._cache = new Map();
+    this.cacheStats = { hits: 0, misses: 0 };
+  }
+
+  /**
+   * Call the embedder, forwarding the signal only when one is given.
+   * @private
+   */
+  _embed(text, signal) {
+    const fn = this.embedder.embedQuery ? 'embedQuery' : 'embed';
+    return signal
+      ? this.embedder[fn](text, { signal })
+      : this.embedder[fn](text);
+  }
+
+  /**
+   * Embed a document, using the LRU cache. In-flight promises are cached so
+   * duplicate documents in one batch are embedded once; failed embeddings
+   * are evicted so they can be retried.
+   * @private
+   */
+  _embedDocument(text, signal) {
+    const { cacheSize } = this.config;
+    if (!cacheSize || cacheSize <= 0) {
+      this.cacheStats.misses++;
+      return Promise.resolve().then(() => this._embed(text, signal));
+    }
+
+    const key = contentHash(text);
+    const cached = this._cache.get(key);
+    if (cached) {
+      this.cacheStats.hits++;
+      // Refresh recency.
+      this._cache.delete(key);
+      this._cache.set(key, cached);
+      // The shared promise may have been started by another request whose
+      // signal aborted; that abort must not fail this request.
+      return cached.catch((err) => {
+        if (isAbortError(err) && !(signal && signal.aborted)) {
+          return this._embedDocument(text, signal);
+        }
+        throw err;
+      });
+    }
+
+    this.cacheStats.misses++;
+    const promise = Promise.resolve().then(() => this._embed(text, signal));
+    this._cache.set(key, promise);
+    promise.catch(() => {
+      if (this._cache.get(key) === promise) this._cache.delete(key);
+    });
+    while (this._cache.size > cacheSize) {
+      this._cache.delete(this._cache.keys().next().value);
+    }
+    return promise;
+  }
+
+  /**
+   * Clear the document embedding cache.
+   */
+  clearCache() {
+    this._cache.clear();
+    this.cacheStats = { hits: 0, misses: 0 };
   }
 
   /**
@@ -40,6 +138,7 @@ class EmbeddingReranker extends EventEmitter {
    * @param {Array<object|string>} documents - Documents to rerank
    * @param {object} [options]
    * @param {number} [options.topK] - Maximum number of results to return
+   * @param {AbortSignal} [options.signal] - Cancels pending embedding calls
    * @returns {Promise<object[]>} Reranked documents with .relevanceScore
    */
   async rerank(query, documents, options = {}) {
@@ -47,27 +146,43 @@ class EmbeddingReranker extends EventEmitter {
       return [];
     }
 
-    const { topK = documents.length } = options;
+    const { topK = documents.length, signal } = options;
+    throwIfAborted(signal);
 
-    // Get query embedding
-    const queryEmbedding = this.embedder.embedQuery
-      ? await this.embedder.embedQuery(query)
-      : await this.embedder.embed(query);
+    const queryEmbedding = await raceAbort(
+      Promise.resolve().then(() => this._embed(query, signal)),
+      signal,
+    );
 
-    // Get document embeddings and compute similarities
-    const scored = await Promise.all(
-      documents.map(async (doc, originalIndex) => {
+    // Embed documents with at most `concurrency` calls in flight, keeping
+    // results in input order.
+    const embeddings = new Array(documents.length);
+    let next = 0;
+    const worker = async () => {
+      while (next < documents.length) {
+        throwIfAborted(signal);
+        const i = next++;
+        const doc = documents[i];
         const text = doc.content || doc.text || String(doc);
-
-        const docEmbedding = this.embedder.embedQuery
-          ? await this.embedder.embedQuery(text)
-          : await this.embedder.embed(text);
-
-        const similarity = this._cosineSimilarity(queryEmbedding, docEmbedding);
-
-        return { doc, similarity, originalIndex };
-      }),
-    );
+        embeddings[i] = await raceAbort(
+          this._embedDocument(text, signal),
+          signal,
+        );
+      }
+    };
+    const workers = [];
+    const poolSize = Math.min(this.config.concurrency, documents.length);
+    for (let w = 0; w < poolSize; w++) workers.push(worker());
+    await Promise.all(workers);
+
+    const scored = documents.map((doc, originalIndex) => ({
+      doc,
+      similarity: this._cosineSimilarity(
+        queryEmbedding,
+        embeddings[originalIndex],
+      ),
+      originalIndex,
+    }));
 
     // Sort by similarity descending, then by original index for stability
     scored.sort((a, b) => {

package/src/reranker/llm-reranker.js

@@ -1,28 +1,83 @@
 'use strict';
 
+const {
+  createAbortError,
+  isAbortError,
+  raceAbort,
+  throwIfAborted,
+} = require('../connectors/connector-utils');
+
+// Characters of each document shown to the LLM; keeps the ranking prompt small.
+const DEFAULT_MAX_DOC_CHARS = 500;
+
 class LLMReranker {
   constructor(options = {}) {
     this.llm = options.llm;
+    this.maxDocChars = options.maxDocChars || DEFAULT_MAX_DOC_CHARS;
+  }
+
+  /**
+   * Build the instruction sent to the LLM. Documents are numbered from 0 so
+   * the response can be validated as indices into `docs`.
+   */
+  _buildRankingPrompt(query, docs) {
+    const listing = docs
+      .map((doc, i) => {
+        const text = String(doc.text || doc.content || '').slice(
+          0,
+          this.maxDocChars,
+        );
+        return `[${i}] ${text.replace(/\s+/g, ' ')}`;
+      })
+      .join('\n');
+    return [
+      'Rank the documents below by how relevant they are to the query.',
+      'Respond with only a JSON array of document numbers, most relevant',
+      'first, for example [2, 0, 1]. Do not include any other text.',
+      '',
+      `Query: ${query}`,
+      '',
+      'Documents:',
+      listing,
+    ].join('\n');
+  }
+
+  /**
+   * Parse the first JSON array in the response; models often wrap it in
+   * prose or a code fence.
+   */
+  _parseRanking(response) {
+    const text = typeof response === 'string' ? response : String(response);
+    const match = text.match(/\[[\s\d,]*\]/);
+    return JSON.parse(match ? match[0] : text);
   }
 
   async rerank(query, docs, opts = {}) {
     if (!Array.isArray(docs)) {
       return [];
     }
+    const signal = opts.signal;
+    throwIfAborted(signal);
 
     // If no LLM provided, fall back to simple overlap scoring
     if (!this.llm) {
-      return this._fallbackRerank(query, docs);
+      return this._limit(this._fallbackRerank(query, docs), opts);
     }
 
     try {
-      // Use LLM to get reranking order
-      const response = await this.llm.generate(query, docs, opts);
+      // Ask the LLM for a ranking. Documents are inlined in the prompt, so no
+      // separate retrieval context is passed.
+      const response = await raceAbort(
+        Promise.resolve().then(() =>
+          this.llm.generate(this._buildRankingPrompt(query, docs), [], opts),
+        ),
+        signal,
+      );
 
       // Validate and sanitize LLM output before parsing (security fix)
       let rankingOrder;
       try {
-        const parsed = JSON.parse(response);
+        const parsed = this._parseRanking(response);
 
         // Security validation: ensure parsed result is an array
         if (!Array.isArray(parsed)) {
@@ -49,22 +104,34 @@ class LLMReranker {
         }
       } catch (parseErr) {
         // Fallback: use simple overlap scoring if parsing fails
-        return this._fallbackRerank(query, docs);
+        return this._limit(this._fallbackRerank(query, docs), opts);
       }
 
-      // Apply the ranking order
-      const reranked = [];
-      for (const index of rankingOrder) {
-        reranked.push(docs[index]);
-      }
+      // Apply the ranking order; documents the LLM left out keep their
+      // original relative order after the ranked ones instead of vanishing.
+      const ranked = new Set(rankingOrder);
+      const reranked = [
+        ...rankingOrder.map((index) => docs[index]),
+        ...docs.filter((_, index) => !ranked.has(index)),
+      ];
 
-      return reranked;
+      return this._limit(reranked, opts);
     } catch (error) {
+      // Cancellation is not an LLM failure: propagate it instead of falling back.
+      if (isAbortError(error) || (signal && signal.aborted)) {
+        throw createAbortError(signal);
+      }
       // Fall back to simple overlap scoring on LLM error
-      return this._fallbackRerank(query, docs);
+      return this._limit(this._fallbackRerank(query, docs), opts);
     }
   }
 
+  _limit(docs, opts) {
+    return Number.isInteger(opts.topK) && opts.topK > 0
+      ? docs.slice(0, opts.topK)
+      : docs;
+  }
+
   _fallbackRerank(query, docs) {
     // Simple bag-of-words overlap scoring
     const queryTokens = query

package/src/reranker/scoring-reranker.js

@@ -9,6 +9,7 @@
 
 const { EventEmitter } = require('events');
 const { tokenize } = require('../evaluate/scoring');
+const { throwIfAborted } = require('../connectors/connector-utils');
 
 /** @type {{ k1: number, b: number }} */
 const DEFAULT_CONFIG = { k1: 1.2, b: 0.75 };
@@ -37,9 +38,11 @@ class ScoringReranker extends EventEmitter {
    * @param {Array<object|string>} documents - Documents to rerank
    * @param {object} [options]
    * @param {number} [options.topK] - Maximum number of results to return
+   * @param {AbortSignal} [options.signal] - Rejects with AbortError when already aborted
    * @returns {Promise<object[]>} Reranked documents with .relevanceScore
    */
   async rerank(query, documents, options = {}) {
+    throwIfAborted(options.signal);
     if (!Array.isArray(documents) || documents.length === 0) {
       return [];
     }

package/src/retrieval/bm25-search.js

@@ -1,7 +1,8 @@
 'use strict';
 
 const { EventEmitter } = require('events');
-const { tokenize } = require('../evaluate/scoring');
+// Called through the module object so tests can observe tokenization.
+const scoring = require('../evaluate/scoring');
 
 const DEFAULT_CONFIG = { k1: 1.2, b: 0.75 };
 
@@ -11,6 +12,8 @@ class BM25Search extends EventEmitter {
     this.config = { ...DEFAULT_CONFIG, ...options };
     this.invertedIndex = new Map(); // token -> [{docIndex, tf}]
     this.documents = [];
+    /** @type {number[]} token count per document, computed at index time */
+    this.docLengths = [];
     this.avgDocLength = 0;
     this.docCount = 0;
   }
@@ -27,14 +30,16 @@ class BM25Search extends EventEmitter {
     this.documents = documents;
     this.docCount = documents.length;
     this.invertedIndex.clear();
+    this.docLengths = new Array(documents.length);
 
     let totalLength = 0;
 
     for (let docIndex = 0; docIndex < documents.length; docIndex++) {
       const doc = documents[docIndex];
       const text = doc.content || doc.text || '';
-      const tokens = tokenize(text);
+      const tokens = scoring.tokenize(text);
 
+      this.docLengths[docIndex] = tokens.length;
       totalLength += tokens.length;
 
       // Count term frequencies for this document
@@ -73,7 +78,7 @@ class BM25Search extends EventEmitter {
       return [];
     }
 
-    const queryTokens = tokenize(query);
+    const queryTokens = scoring.tokenize(query);
     if (queryTokens.length === 0) {
       return [];
     }
@@ -94,9 +99,9 @@ class BM25Search extends EventEmitter {
       const idf = Math.log((N - df + 0.5) / (df + 0.5) + 1);
 
       for (const { docIndex, tf } of postings) {
-        const doc = this.documents[docIndex];
-        const docText = doc.content || doc.text || '';
-        const docLen = tokenize(docText).length;
+        // Document lengths were computed at index time; scoring never
+        // re-tokenizes documents.
+        const docLen = this.docLengths[docIndex];
 
         const numerator = tf * (k1 + 1);
         const denominator =
@@ -137,6 +142,7 @@ class BM25Search extends EventEmitter {
   clear() {
     this.invertedIndex.clear();
     this.documents = [];
+    this.docLengths = [];
     this.avgDocLength = 0;
     this.docCount = 0;
   }