@devilsdev/rag-pipeline-utils 2.4.5 → 2.6.0

package/src/core/create-pipeline.js

@@ -1,23 +1,161 @@
 'use strict';
-const { PluginRegistry } = require('./plugin-registry.js');
+// Share the exported singleton so plugins registered via
+// `pluginRegistry.register()` are resolvable by name here.
+const defaultRegistry = require('./plugin-registry.js');
+const { embedQuery, canEmbed } = require('../utils/embed-query');
+const {
+  normalizeCacheOption,
+  PipelineCache,
+} = require('../cache/pipeline-cache');
+const { createRunAbort } = require('../utils/run-abort');
 
-const defaultRegistry = new PluginRegistry();
+/** Metrics that only need the query, the retrieved docs and a reference. */
+const RETRIEVAL_METRICS = ['contextPrecision', 'contextRecall'];
 
 function resolve(registry, type, value) {
   const reg = registry || defaultRegistry;
   return typeof value === 'string' ? reg.get(type, value) : value;
 }
 
-function withTimeout(promise, ms) {
-  if (!ms) return promise;
-  return Promise.race([
-    promise,
-    new Promise((_, rej) =>
-      setTimeout(() => rej(new Error('Execution timeout')), ms),
-    ),
-  ]);
+/**
+ * Citations and evaluation for one run. `answer` is null when no llm is
+ * configured: citations are then skipped and evaluation is limited to the
+ * retrieval metrics (context precision, and context recall when a reference
+ * answer is given).
+ *
+ * @param {object} params
+ * @param {string} params.query
+ * @param {string|null} params.answer
+ * @param {Array<object>} params.results
+ * @param {object} params.options - Run options
+ * @returns {{ citations?: object, evaluation?: object }}
+ */
+function computeQuality({ query, answer, results, options }) {
+  const out = {};
+  const docs = Array.isArray(results) ? results : [];
+
+  if (options.citations && answer !== null) {
+    const { CitationTracker } = require('../citation/citation-tracker');
+    const tracker = new CitationTracker(options.citationOptions || {});
+    out.citations = tracker.track(answer, docs);
+  }
+
+  if (options.evaluate) {
+    const {
+      PipelineEvaluator,
+      DEFAULT_CONFIG,
+    } = require('../evaluate/pipeline-evaluator');
+    const evaluateOptions = options.evaluateOptions || {};
+    let config = evaluateOptions;
+    if (answer === null) {
+      const metrics = evaluateOptions.metrics || DEFAULT_CONFIG.metrics;
+      config = {
+        ...evaluateOptions,
+        metrics: metrics.filter((m) => RETRIEVAL_METRICS.includes(m)),
+      };
+    }
+    const evaluator = new PipelineEvaluator(config);
+    out.evaluation = evaluator.evaluate(
+      { query, answer: answer ?? '', results: docs },
+      {
+        citationResult: out.citations,
+        referenceAnswer:
+          options.referenceAnswer ?? evaluateOptions.referenceAnswer,
+      },
+    );
+  }
+
+  return out;
+}
+
+/**
+ * Pass a token stream through, stopping when the run's signal aborts. When
+ * citations or evaluation were requested, the streamed answer is accumulated
+ * and the stream ends with one final chunk
+ * `{ token: '', done: true, citations?, evaluation? }` (the source's own done
+ * chunk is replaced by it; a token it carried is emitted first).
+ *
+ * @param {AsyncIterable<object|string>} source
+ * @param {object} params
+ * @param {AbortSignal} params.signal
+ * @param {((text: string) => object)|null} params.finalize
+ * @param {() => void} params.onEnd
+ */
+async function* guardStream(source, { signal, finalize, onEnd }) {
+  try {
+    let text = '';
+    for await (const chunk of source) {
+      if (signal.aborted) throw signal.reason;
+      if (!finalize) {
+        yield chunk;
+        continue;
+      }
+      const token =
+        typeof chunk === 'string' ? chunk : String(chunk?.token ?? '');
+      text += token;
+      if (chunk && chunk.done) {
+        if (token) yield { ...chunk, token, done: false };
+        break;
+      }
+      yield chunk;
+    }
+    if (signal.aborted) throw signal.reason;
+    if (finalize) yield { token: '', done: true, ...finalize(text) };
+  } finally {
+    onEnd();
+  }
+}
+
+/**
+ * Pass a stream through and report its text when it completed: it reached a
+ * done chunk or its end, and the run was not aborted. A stream the consumer
+ * stops reading early is not reported, so a partial answer is never cached.
+ */
+async function* reportWhenComplete(source, signal, onComplete) {
+  let text = '';
+  let completed = false;
+  try {
+    for await (const chunk of source) {
+      text += typeof chunk === 'string' ? chunk : String(chunk?.token ?? '');
+      if (chunk && chunk.done) completed = true;
+      yield chunk;
+    }
+    completed = true;
+  } finally {
+    if (completed && !signal.aborted) await onComplete(text);
+  }
 }
 
+/**
+ * Create a RAG pipeline: (embed) -> retrieve -> (filter) -> (rerank) ->
+ * generate -> (citations, evaluation).
+ *
+ * `run({ query, queryVector, options })` options:
+ * - `topK` (default 3)
+ * - `timeout` (ms): aborts the run; resolves `{ success: false, error: 'Execution timeout' }`
+ * - `signal` (AbortSignal): caller cancellation; resolves
+ *   `{ success: false, error: 'Execution aborted' }` when it aborts
+ * - `stream`: return an async iterable of `{ token, done }` chunks
+ * - `citations`, `citationOptions`, `evaluate`, `evaluateOptions`,
+ *   `referenceAnswer` (ground truth for context recall)
+ * - `filterResults(docs)`: hook applied after retrieval
+ *
+ * `cache` (factory option): `true`, a CacheBackend, or
+ * `{ embeddings, retrieval, generation, ttlSeconds }` - see
+ * src/cache/pipeline-cache.js. Results then carry `cache: { embeddings,
+ * retrieval, generation }` ('hit', 'miss' or 'skip'), and the pipeline has
+ * `cacheStats()`.
+ *
+ * Every component call receives the run's AbortSignal: `retriever.retrieve({
+ * ..., signal })`, `embedder.embedQuery(text, { signal })` /
+ * `embedder.embed(text, { signal })`, `reranker.rerank(query, docs, { topK,
+ * signal })` and `llm.generate(query, docs, { ...options, signal })`.
+ * Components that ignore it still work; the run settles on abort anyway and
+ * no later stage starts.
+ *
+ * Without an llm the pipeline is retrieval-only: `answer` is null, citations
+ * are skipped and evaluation covers the retrieval metrics only.
+ */
 module.exports.createRagPipeline = function createRagPipeline({
   registry,
   loader,
@@ -25,6 +163,7 @@ module.exports.createRagPipeline = function createRagPipeline({
   retriever,
   llm,
   reranker,
+  cache,
 } = {}) {
   const reg = registry || defaultRegistry;
   const loaderI = loader ? resolve(reg, 'loader', loader) : null;
@@ -36,84 +175,191 @@ module.exports.createRagPipeline = function createRagPipeline({
           return [];
         },
       };
-  const llmI = llm
-    ? resolve(reg, 'llm', llm)
-    : {
-        async generate() {
-          return 'ok';
-        },
-      };
+  const llmI = llm ? resolve(reg, 'llm', llm) : null;
   const rerankerI = reranker ? resolve(reg, 'reranker', reranker) : null;
 
-  async function runOnce({ query, queryVector, options = {} }) {
-    const { timeout, stream } = options || {};
+  const cacheLayers = normalizeCacheOption(cache);
+  const pc = cacheLayers
+    ? new PipelineCache(cacheLayers, {
+        embedder: embedderI,
+        // The default stand-in retriever returns nothing; never cache it.
+        retriever: retriever ? retrieverI : null,
+        llm: llmI,
+      })
+    : null;
+
+  async function runOnce({ query, queryVector, options = {} } = {}) {
+    const opts = options || {};
+    const { timeout, stream } = opts;
+    const run = createRunAbort({ timeout, signal: opts.signal });
+    const { signal } = run;
+    let streaming = false;
+
+    // Per-run cache outcome per layer: 'hit', 'miss' or 'skip'.
+    const cacheStatus = {};
+    const withCache = (result) =>
+      pc ? { ...result, cache: { ...cacheStatus } } : result;
+
     const doWork = (async () => {
-      if (loaderI?.load) await loaderI.load();
-      if (embedderI?.embed && query) await embedderI.embed(query);
+      run.throwIfAborted();
+      if (loaderI?.load) {
+        await loaderI.load();
+        run.throwIfAborted();
+      }
+      let vector = queryVector;
+      if (!vector && query && canEmbed(embedderI)) {
+        const key = pc && pc.embeddingKey(query);
+        vector = pc && (await pc.get('embeddings', key, cacheStatus));
+        if (!vector) {
+          vector = await embedQuery(embedderI, query, { signal });
+          run.throwIfAborted();
+          if (pc) await pc.set('embeddings', key, vector);
+        }
+        run.throwIfAborted();
+      }
 
-      const docs = await retrieverI.retrieve({
-        query,
-        queryVector,
-        topK: options.topK ?? 3,
-      });
+      // Raw retriever output is cached before filterResults (guardrail ACLs)
+      // and the reranker, so both still run on every request.
+      const topK = opts.topK ?? 3;
+      const retrievalKey =
+        pc && pc.retrievalKey({ query, queryVector: vector, topK });
+      let docs = pc && (await pc.get('retrieval', retrievalKey, cacheStatus));
+      if (!docs) {
+        docs = await retrieverI.retrieve({
+          query,
+          queryVector: vector,
+          topK,
+          signal,
+        });
+        run.throwIfAborted();
+        if (pc && Array.isArray(docs)) {
+          await pc.set('retrieval', retrievalKey, docs);
+        }
+      }
+      run.throwIfAborted();
+      // Hook for filtering (e.g. ACL guardrails) before anything downstream,
+      // including the LLM, sees the retrieved documents.
+      if (typeof opts.filterResults === 'function' && Array.isArray(docs)) {
+        docs = await opts.filterResults(docs);
+        run.throwIfAborted();
+      }
       let results = docs;
       if (rerankerI?.rerank && Array.isArray(docs) && docs.length) {
         results = await rerankerI.rerank(query || '', docs, {
-          topK: options.topK ?? docs.length,
+          topK: opts.topK ?? docs.length,
+          signal,
         });
+        run.throwIfAborted();
       }
 
+      const genOptions = { ...opts, signal };
+      const wantsQuality = Boolean(opts.citations || opts.evaluate);
+      const generationKey =
+        pc && llmI
+          ? pc.generationKey({ query, docs: results, options: opts })
+          : null;
+      const cachedAnswer =
+        pc && llmI
+          ? await pc.get('generation', generationKey, cacheStatus)
+          : undefined;
+      run.throwIfAborted();
+
       if (stream) {
-        // Use streaming generate if available
-        if (llmI.generateStream) {
-          const streamResult = llmI.generateStream(query, results, options);
-          return streamResult;
+        let source;
+        let finalize = wantsQuality
+          ? (text) =>
+              computeQuality({ query, answer: text, results, options: opts })
+          : null;
+        if (!llmI) {
+          // Retrieval only: no tokens, one done chunk carrying the results.
+          const quality = computeQuality({
+            query,
+            answer: null,
+            results,
+            options: opts,
+          });
+          source = (async function* noAnswer() {})();
+          finalize = () => ({ results, ...quality });
+        } else if (typeof cachedAnswer === 'string') {
+          source = (async function* cached() {
+            yield { token: cachedAnswer, done: false };
+            yield { token: '', done: true };
+          })();
+        } else if (llmI.generateStream) {
+          source = await llmI.generateStream(query, results, genOptions);
         } else if (llmI.generate) {
           // Fallback: simulate streaming from regular generate
-          const answer = (await llmI.generate(query, results, options)) ?? '';
-          async function* gen() {
-            const chunks = String(answer || '').length
-              ? [String(answer)]
-              : [''];
-            for (const c of chunks) yield { token: c, done: false };
+          const answer =
+            (await llmI.generate(query, results, genOptions)) ?? '';
+          source = (async function* gen() {
+            yield { token: String(answer || ''), done: false };
             yield { token: '', done: true };
+          })();
+        }
+        if (source) {
+          run.throwIfAborted();
+          streaming = true;
+          if (pc && llmI && generationKey && typeof cachedAnswer !== 'string') {
+            source = reportWhenComplete(source, signal, (text) =>
+              pc.set('generation', generationKey, text),
+            );
           }
-          return gen();
+          return guardStream(source, {
+            signal,
+            finalize,
+            onEnd: run.detach,
+          });
         }
       }
 
-      const answer = (await llmI.generate?.(query, results, options)) ?? '';
-      // tiny delay so durations > 0 in timing tests
-      await new Promise((r) => setTimeout(r, 5));
-
-      const result = { success: true, query, results };
-
-      if (options.citations) {
-        const { CitationTracker } = require('../citation/citation-tracker');
-        const tracker = new CitationTracker(options.citationOptions || {});
-        result.answer = String(answer);
-        result.citations = tracker.track(String(answer), results);
+      if (!llmI) {
+        return withCache({
+          success: true,
+          query,
+          results,
+          answer: null,
+          ...computeQuality({ query, answer: null, results, options: opts }),
+        });
       }
 
-      if (options.evaluate) {
-        const { PipelineEvaluator } = require('../evaluate/pipeline-evaluator');
-        const evaluator = new PipelineEvaluator(options.evaluateOptions || {});
-        result.answer = result.answer || String(answer);
-        result.evaluation = evaluator.evaluate(
-          { query, answer: String(answer), results },
-          { citationResult: result.citations },
-        );
+      let answer;
+      if (typeof cachedAnswer === 'string') {
+        answer = cachedAnswer;
+      } else {
+        const raw = (await llmI.generate?.(query, results, genOptions)) ?? '';
+        run.throwIfAborted();
+        answer = String(raw);
+        if (pc) await pc.set('generation', generationKey, answer);
+        run.throwIfAborted();
       }
 
-      return result;
+      return withCache({
+        success: true,
+        query,
+        results,
+        answer,
+        ...computeQuality({ query, answer, results, options: opts }),
+      });
     })();
 
     try {
-      return await withTimeout(doWork, timeout);
+      return await run.race(doWork);
     } catch (e) {
-      return { success: false, error: String(e.message || e) };
+      return { success: false, error: String(e?.message || e) };
+    } finally {
+      // The timeout covers the run up to the returned stream; a caller
+      // signal keeps cancelling a stream until it finishes.
+      run.clearTimer();
+      if (!streaming) run.detach();
     }
   }
 
-  return { run: runOnce, cleanup() {} };
+  return {
+    run: runOnce,
+    cleanup() {},
+    /** Cache hits, misses and skips per layer; null when caching is off. */
+    cacheStats() {
+      return pc ? pc.stats() : null;
+    },
+  };
 };

package/src/core/performance/parallel-processor.js

@@ -36,16 +36,77 @@ class Semaphore {
   }
 }
 
+/** Built-in embedders whose embed() takes one string, not an array. */
+function isSingleTextEmbedder(embedder) {
+  /* eslint-disable global-require */
+  const singleText = [
+    require('../../connectors/local-embedder').LocalEmbedder,
+    require('../../connectors/openai-connector').OpenAIConnector,
+    require('../../connectors/ollama-connector').OllamaConnector,
+    require('../../connectors/cohere-connector').CohereConnector,
+  ];
+  /* eslint-enable global-require */
+  return singleText.some((Class) => Class && embedder instanceof Class);
+}
+
+const EMBED_MODES = ['auto', 'batch', 'single'];
+
 /**
  * Parallel processor for embedding operations with configurable concurrency
  */
 class ParallelEmbedder {
+  /**
+   * @param {object} embedder - Embedder with embed() and/or embedBatch()
+   * @param {object} [_options]
+   * @param {number} [_options.batchSize=10] - Texts per batch
+   * @param {number} [_options.maxConcurrency=3] - Batches in flight at once
+   * @param {number} [_options.retryAttempts=2] - Retries per failed batch
+   * @param {number} [_options.retryDelay=1000] - Base retry delay in ms
+   * @param {'auto'|'batch'|'single'} [_options.mode='auto'] - How a batch is
+   *   embedded: 'batch' calls embed(texts) and expects one vector per text,
+   *   'single' calls embed(text) for each text. 'auto' uses embedBatch(texts)
+   *   when the embedder has it, 'single' for the built-in connectors (which
+   *   embed one string at a time) and 'batch' otherwise.
+   */
   constructor(embedder, _options = {}) {
     this.embedder = embedder;
     this.batchSize = _options.batchSize || 10;
     this.maxConcurrency = _options.maxConcurrency || 3;
     this.retryAttempts = _options.retryAttempts || 2;
     this.retryDelay = _options.retryDelay || 1000;
+    const mode = _options.mode || 'auto';
+    if (!EMBED_MODES.includes(mode)) {
+      throw new Error(
+        `Invalid mode "${mode}". Expected one of: ${EMBED_MODES.join(', ')}`,
+      );
+    }
+    this.mode = mode;
+  }
+
+  /**
+   * Embed one batch according to the mode.
+   * @param {string[]} batch
+   * @returns {Promise<number[][]>}
+   * @private
+   */
+  async _embed(batch) {
+    const embedder = this.embedder;
+    if (this.mode === 'auto' && typeof embedder.embedBatch === 'function') {
+      return embedder.embedBatch(batch);
+    }
+    const single =
+      this.mode === 'single' ||
+      (this.mode === 'auto' && isSingleTextEmbedder(embedder));
+    if (!single) {
+      return embedder.embed(batch);
+    }
+    // One text at a time within a batch, so maxConcurrency still bounds the
+    // number of requests in flight.
+    const vectors = [];
+    for (const text of batch) {
+      vectors.push(await embedder.embed(text));
+    }
+    return vectors;
   }
 
   /**
@@ -133,7 +194,7 @@ class ParallelEmbedder {
 
     for (let attempt = 0; attempt <= this.retryAttempts; attempt++) {
       try {
-        const result = await this.embedder.embed(batch);
+        const result = await this._embed(batch);
 
         // Validate result
         if (!Array.isArray(result) || result.length !== batch.length) {