@devilsdev/rag-pipeline-utils 2.1.11 → 2.2.0

package/src/plugins/sample-embedder.js

@@ -0,0 +1,53 @@
+/**
+ * Sample Embedder Plugin
+ * Minimal fixture for test compatibility with configurable batch processing
+ */
+
+class SampleEmbedder {
+  constructor(options = {}) {
+    this.name = 'sample-embedder';
+    this.version = '1.0.0';
+    this.options = options;
+
+    // Configurable batch size from environment or options
+    this.batchSize =
+      parseInt(process.env.RAG_EMBEDDER_BATCH_SIZE) || options.batchSize || 50; // Default batch size
+  }
+
+  async embed(texts) {
+    if (!Array.isArray(texts)) {
+      return new Array(384).fill(0).map(() => Math.random());
+    }
+
+    // Process in batches for better memory management
+    const results = [];
+    for (let i = 0; i < texts.length; i += this.batchSize) {
+      const batch = texts.slice(i, i + this.batchSize);
+      const batchResults = await this._processBatch(batch);
+      results.push(...batchResults);
+
+      // Optional progress callback
+      if (this.options.onProgress) {
+        this.options.onProgress({
+          processed: Math.min(i + this.batchSize, texts.length),
+          total: texts.length,
+          stage: 'embedding',
+        });
+      }
+    }
+
+    return results;
+  }
+
+  async _processBatch(batch) {
+    // Simulate batch processing with slight delay for realism
+    await new Promise((resolve) => setTimeout(resolve, 10));
+    return batch.map(() => new Array(384).fill(0).map(() => Math.random()));
+  }
+
+  async embedQuery(query) {
+    return new Array(384).fill(0).map(() => Math.random());
+  }
+}
+
+module.exports = { SampleEmbedder };

package/src/plugins/sample-llm.js

@@ -0,0 +1,30 @@
+/**
+ * Sample LLM Plugin
+ * Minimal fixture for test compatibility
+ */
+
+class SampleLLM {
+  constructor(options = {}) {
+    this.name = 'sample-llm';
+    this.version = '1.0.0';
+    this.options = options;
+  }
+
+  async generate(prompt, context = []) {
+    // Mock response for testing
+    return `Generated response for: ${prompt.substring(0, 50)}...`;
+  }
+
+  async *generateStream(prompt, context = []) {
+    const response = `Generated response for: ${prompt.substring(0, 50)}...`;
+    const tokens = response.split(' ');
+
+    for (const token of tokens) {
+      yield token + ' ';
+      // Small delay for realistic streaming
+      await new Promise((resolve) => setTimeout(resolve, 10));
+    }
+  }
+}
+
+module.exports = { SampleLLM };

package/src/plugins/sample-loader.js

@@ -0,0 +1,24 @@
+/**
+ * Sample Loader Plugin
+ * Minimal fixture for test compatibility
+ */
+
+class SampleLoader {
+  constructor(options = {}) {
+    this.name = 'sample-loader';
+    this.version = '1.0.0';
+    this.options = options;
+  }
+
+  async load(filePath) {
+    return [
+      {
+        content: 'Sample document content',
+        metadata: { source: filePath },
+        chunk: () => ['Sample chunk 1', 'Sample chunk 2'],
+      },
+    ];
+  }
+}
+
+module.exports = { SampleLoader };

package/src/plugins/sample-retriever.js

@@ -0,0 +1,84 @@
+/**
+ * Sample Retriever Plugin
+ * Minimal fixture for test compatibility with configurable batch processing
+ */
+
+class SampleRetriever {
+  constructor(options = {}) {
+    this.name = 'sample-retriever';
+    this.version = '1.0.0';
+    this.options = options;
+    this.storage = [];
+
+    // Configurable batch size from environment or options
+    this.batchSize =
+      parseInt(process.env.RAG_RETRIEVER_BATCH_SIZE) ||
+      options.batchSize ||
+      100; // Default batch size for retrieval
+  }
+
+  async store(vectors) {
+    // Process storage in batches for better memory management
+    if (Array.isArray(vectors) && vectors.length > this.batchSize) {
+      for (let i = 0; i < vectors.length; i += this.batchSize) {
+        const batch = vectors.slice(i, i + this.batchSize);
+        this.storage.push(...batch);
+
+        // Optional progress callback
+        if (this.options.onProgress) {
+          this.options.onProgress({
+            processed: Math.min(i + this.batchSize, vectors.length),
+            total: vectors.length,
+            stage: 'storing',
+          });
+        }
+
+        // Small delay to prevent overwhelming memory
+        await new Promise((resolve) => setTimeout(resolve, 5));
+      }
+    } else {
+      this.storage.push(...vectors);
+    }
+    return true;
+  }
+
+  async retrieve(queryVector, k = 5) {
+    // For large result sets, process in batches
+    const totalResults = Math.min(k, this.storage.length);
+    const results = [];
+
+    for (let i = 0; i < totalResults; i += this.batchSize) {
+      const batchSize = Math.min(this.batchSize, totalResults - i);
+      const batchResults = this._retrieveBatch(i, batchSize);
+      results.push(...batchResults);
+
+      // Optional progress callback
+      if (this.options.onProgress && totalResults > this.batchSize) {
+        this.options.onProgress({
+          processed: Math.min(i + this.batchSize, totalResults),
+          total: totalResults,
+          stage: 'retrieving',
+        });
+      }
+
+      // Small delay for large retrievals
+      if (totalResults > this.batchSize) {
+        await new Promise((resolve) => setTimeout(resolve, 5));
+      }
+    }
+
+    return results;
+  }
+
+  _retrieveBatch(startIndex, batchSize) {
+    return this.storage
+      .slice(startIndex, startIndex + batchSize)
+      .map((vector, index) => ({
+        content: `Retrieved document ${startIndex + index + 1}`,
+        score: Math.random(),
+        metadata: { index: startIndex + index },
+      }));
+  }
+}
+
+module.exports = { SampleRetriever };

package/src/query.js

@@ -0,0 +1,117 @@
+/**
+ * Version: 2.0.0
+ * Path: /src/query.js
+ * Description: Handles querying the RAG pipeline with streaming support
+ * Author: Ali Kahwaji
+ */
+
+const { createRagPipeline } = require('./core/create-pipeline.js');
+// eslint-disable-line global-require
+const { loadPluginsFromJson } = require('./config/load-plugin-config.js');
+// eslint-disable-line global-require
+const { logger } = require('./utils/logger.js');
+// eslint-disable-line global-require
+
+/**
+ * Query the RAG pipeline with a standard response
+ * @param {string} prompt - Query prompt
+ * @param {object} config - Pipeline configuration
+ * @returns {Promise<string>} Generated response
+ */
+async function queryPipeline(_prompt, config) {
+  if (!prompt) {
+    throw new Error('No prompt provided for query.');
+  }
+
+  if (!config) {
+    throw new Error('No configuration provided for query.');
+  }
+
+  try {
+    // Load plugins from configuration
+    await loadPluginsFromJson(process.cwd());
+
+    // Extract plugin names from config
+    const loaderName = Object.keys(config.loader)[0];
+    const embedderName = Object.keys(config.embedder)[0];
+    const retrieverName = Object.keys(config.retriever)[0];
+    const llmName = Object.keys(config.llm)[0];
+
+    // Create pipeline instance
+    const pipeline = createRagPipeline(
+      {
+        loader: loaderName,
+        embedder: embedderName,
+        retriever: retrieverName,
+        llm: llmName,
+      },
+      {
+        useRetry: true,
+        useLogging: true,
+        useReranker: config.useReranker || false,
+      },
+    );
+
+    // Execute query
+    return await pipeline.query(prompt);
+  } catch (error) {
+    logger.error('Query pipeline failed', { error: error.message, prompt });
+    throw error;
+  }
+}
+
+/**
+ * Query the RAG pipeline with streaming response
+ * @param {string} prompt - Query prompt
+ * @param {object} config - Pipeline configuration
+ * @returns {AsyncIterable<string>} Stream of response tokens
+ */
+async function* queryPipelineStream(prompt, config) {
+  if (!prompt) {
+    throw new Error('No prompt provided for streaming query.');
+  }
+
+  if (!config) {
+    throw new Error('No configuration provided for streaming query.');
+  }
+
+  try {
+    // Load plugins from configuration
+    await loadPluginsFromJson(process.cwd());
+
+    // Extract plugin names from config
+    const loaderName = Object.keys(config.loader)[0];
+    const embedderName = Object.keys(config.embedder)[0];
+    const retrieverName = Object.keys(config.retriever)[0];
+    const llmName = Object.keys(config.llm)[0];
+
+    // Create pipeline instance
+    const pipeline = createRagPipeline(
+      {
+        loader: loaderName,
+        embedder: embedderName,
+        retriever: retrieverName,
+        llm: llmName,
+      },
+      {
+        useRetry: true,
+        useLogging: true,
+        useReranker: config.useReranker || false,
+      },
+    );
+
+    // Execute streaming query
+    yield* pipeline.queryStream(prompt);
+  } catch (error) {
+    logger.error('Streaming query pipeline failed', {
+      error: error.message,
+      prompt,
+    });
+    throw error;
+  }
+}
+
+module.exports = {
+  queryPipeline,
+  queryPipelineStream,
+};

package/src/reranker/llm-reranker.js

@@ -0,0 +1,85 @@
+'use strict';
+
+class LLMReranker {
+  constructor(options = {}) {
+    this.llm = options.llm;
+  }
+
+  async rerank(query, docs, opts = {}) {
+    if (!Array.isArray(docs)) {
+      return [];
+    }
+
+    // If no LLM provided, fall back to simple overlap scoring
+    if (!this.llm) {
+      return this._fallbackRerank(query, docs);
+    }
+
+    try {
+      // Use LLM to get reranking order
+      const response = await this.llm.generate(query, docs, opts);
+      const rankingOrder = JSON.parse(response);
+
+      // Apply the ranking order
+      const reranked = [];
+      for (const index of rankingOrder) {
+        if (index >= 0 && index < docs.length) {
+          reranked.push(docs[index]);
+        }
+      }
+
+      return reranked;
+    } catch (error) {
+      // Fall back to simple overlap scoring on LLM error
+      return this._fallbackRerank(query, docs);
+    }
+  }
+
+  _fallbackRerank(query, docs) {
+    // Simple bag-of-words overlap scoring
+    const queryTokens = query
+      .toLowerCase()
+      .split(/\s+/)
+      .filter((t) => t.length > 0);
+
+    const scored = docs.map((doc, originalIndex) => {
+      const text = doc.text || doc.content || '';
+      const docTokens = text
+        .toLowerCase()
+        .split(/\s+/)
+        .filter((t) => t.length > 0);
+
+      // Count overlapping tokens
+      let overlap = 0;
+      for (const qToken of queryTokens) {
+        if (docTokens.includes(qToken)) {
+          overlap++;
+        }
+      }
+
+      // Add tiny length-based tie-break
+      const lengthTieBreak = text.length * 0.0001;
+      const score = overlap + lengthTieBreak;
+
+      return { doc, score, originalIndex };
+    });
+
+    // Sort by score (descending), then by original index (ascending) for stable sort
+    scored.sort((a, b) => {
+      if (b.score !== a.score) {
+        return b.score - a.score;
+      }
+      return a.originalIndex - b.originalIndex;
+    });
+
+    return scored.map((item) => item.doc);
+  }
+}
+
+// Create default instance
+const defaultInstance = new LLMReranker();
+
+// CJS+ESM interop pattern
+module.exports = defaultInstance;
+module.exports.LLMReranker = LLMReranker;
+module.exports.default = module.exports;

package/src/scripts/script-utilities.js

@@ -0,0 +1,36 @@
+// src/scripts/script-utilities.js
+
+// default sleeper uses real timers (prod) - unref'd to not block process exit
+const sleep = (ms) =>
+  new Promise((resolve) => {
+    const t = setTimeout(resolve, ms);
+    // critical: don't keep the process alive just because of a timer
+    if (typeof t.unref === "function") t.unref();
+  });
+
+async function retry(fn, opts = {}) {
+  const {
+    retries = 3,
+    baseDelay = 100,
+    factor = 2,
+    onDelay = () => {},
+    sleeper = sleep, // <-- inject for tests
+  } = opts;
+
+  let lastErr;
+  for (let attempt = 0; attempt <= retries; attempt++) {
+    try {
+      // pass attempt if useful for logging/logic
+      return await fn(attempt);
+    } catch (err) {
+      lastErr = err;
+      if (attempt === retries) break;
+      const delay = baseDelay * Math.pow(factor, attempt);
+      onDelay(delay, attempt);
+      await sleeper(delay); // <-- no-op in tests
+    }
+  }
+  throw lastErr;
+}
+
+module.exports = { retry, sleep };