@devilsdev/rag-pipeline-utils 2.4.5 → 2.6.0

package/src/connectors/cohere-connector.js

@@ -1,19 +1,35 @@
 'use strict';
 
-const axios = require('axios');
+// axios loads on first use, so importing the package does not pay for it
+// (about 125 ms) unless this connector is constructed.
+const loadAxios = () => require('axios'); // eslint-disable-line global-require
+const { componentCacheKey } = require('../cache/component-keys');
 const { BaseConnector } = require('./base-connector');
+const {
+  missingApiKeyError,
+  resolveApiKey,
+  throwIfAborted,
+} = require('./connector-utils');
+
+const API_KEY_ENV = 'COHERE_API_KEY';
 
 /**
  * Default configuration for the Cohere connector.
+ * `apiKey` is a getter so the environment is read when a connector is
+ * constructed, not when this module is first required.
  * @type {Object}
  */
 const DEFAULT_CONFIG = {
-  apiKey: process.env.COHERE_API_KEY || '',
+  get apiKey() {
+    return process.env[API_KEY_ENV] || '';
+  },
   baseURL: 'https://api.cohere.ai/v1',
   embeddingModel: 'embed-english-v3.0',
   rerankModel: 'rerank-english-v3.0',
   inputType: 'search_document',
   queryInputType: 'search_query',
+  maxRetries: 3,
+  retryDelay: 1000,
 };
 
 /**
@@ -24,6 +40,21 @@ const DEFAULT_CONFIG = {
  * @extends BaseConnector
  */
 class CohereConnector extends BaseConnector {
+  /**
+   * Cache identity for createRagPipeline({ cache }).
+   * @param {'embedder'} role
+   * @returns {string|null}
+   */
+  getCacheKey(role) {
+    if (role !== 'embedder') return null;
+    const c = this.config;
+    return componentCacheKey('cohere-embed', {
+      baseURL: c.baseURL,
+      model: c.embeddingModel,
+      queryInputType: c.queryInputType,
+    });
+  }
+
   /**
    * @param {Object} [options]
    * @param {string} [options.apiKey] - Cohere API key
@@ -35,8 +66,12 @@ class CohereConnector extends BaseConnector {
    */
   constructor(options = {}) {
     super({ name: 'cohere', ...options });
-    this.config = { ...DEFAULT_CONFIG, ...options };
-    this.client = axios.create({
+    this.config = {
+      ...DEFAULT_CONFIG,
+      ...options,
+      apiKey: resolveApiKey(options, API_KEY_ENV),
+    };
+    this.client = loadAxios().create({
       baseURL: this.config.baseURL,
       headers: {
         Authorization: `Bearer ${this.config.apiKey}`,
@@ -45,37 +80,72 @@ class CohereConnector extends BaseConnector {
     });
   }
 
+  /**
+   * @throws {Error} code MISSING_API_KEY when no key is configured
+   * @private
+   */
+  _requireApiKey() {
+    if (!this.config.apiKey) {
+      throw missingApiKeyError('Cohere', API_KEY_ENV);
+    }
+  }
+
   /**
    * Connect to Cohere. Verifies the API key is present.
    * @returns {Promise<void>}
    * @throws {Error} If no API key is configured
    */
   async connect() {
-    if (!this.config.apiKey) {
-      throw new Error(
-        'Cohere API key required. Set COHERE_API_KEY or pass apiKey option.',
-      );
-    }
+    this._requireApiKey();
     await super.connect();
   }
 
+  /**
+   * POST with the connector's retry policy and cancellation.
+   * @param {string} path
+   * @param {object} body
+   * @param {AbortSignal} [signal]
+   * @returns {Promise<object>} Response data
+   * @private
+   */
+  async _post(path, body, signal) {
+    throwIfAborted(signal);
+    this._requireApiKey();
+    return this.withRetry(
+      async () => {
+        const response = await this.client.post(
+          path,
+          body,
+          signal ? { signal } : undefined,
+        );
+        return response.data;
+      },
+      this.config.maxRetries,
+      this.config.retryDelay,
+      { signal },
+    );
+  }
+
   /**
    * Generate an embedding vector for document text.
    * Uses the document input type by default.
    *
    * @param {string} text - Input text to embed
+   * @param {Object} [options={}]
+   * @param {AbortSignal} [options.signal] - Cancels the request
    * @returns {Promise<number[]>} Embedding vector
    */
-  async embed(text) {
-    return this.withRetry(async () => {
-      const response = await this.client.post('/embed', {
+  async embed(text, options = {}) {
+    const data = await this._post(
+      '/embed',
+      {
         texts: [text],
         model: this.config.embeddingModel,
         input_type: this.config.inputType,
-      });
-
-      return response.data.embeddings[0];
-    });
+      },
+      options.signal,
+    );
+    return data.embeddings[0];
   }
 
   /**
@@ -83,18 +153,21 @@ class CohereConnector extends BaseConnector {
    * Uses the query input type for better retrieval accuracy.
    *
    * @param {string} query - Query text to embed
+   * @param {Object} [options={}]
+   * @param {AbortSignal} [options.signal] - Cancels the request
    * @returns {Promise<number[]>} Embedding vector
    */
-  async embedQuery(query) {
-    return this.withRetry(async () => {
-      const response = await this.client.post('/embed', {
+  async embedQuery(query, options = {}) {
+    const data = await this._post(
+      '/embed',
+      {
         texts: [query],
         model: this.config.embeddingModel,
         input_type: this.config.queryInputType,
-      });
-
-      return response.data.embeddings[0];
-    });
+      },
+      options.signal,
+    );
+    return data.embeddings[0];
   }
 
   /**
@@ -104,26 +177,27 @@ class CohereConnector extends BaseConnector {
    * @param {Array<Object|string>} documents - Documents to rerank
    * @param {Object} [options={}]
    * @param {number} [options.topK] - Number of top results to return (defaults to all)
+   * @param {AbortSignal} [options.signal] - Cancels the request
    * @returns {Promise<Array<Object>>} Reranked documents with relevance scores
    */
   async rerank(query, documents, options = {}) {
-    return this.withRetry(async () => {
-      const docs = documents.map(
-        (doc) => doc.content || doc.text || String(doc),
-      );
+    const docs = documents.map((doc) => doc.content || doc.text || String(doc));
 
-      const response = await this.client.post('/rerank', {
+    const data = await this._post(
+      '/rerank',
+      {
         model: this.config.rerankModel,
         query,
         documents: docs,
         top_n: options.topK || documents.length,
-      });
+      },
+      options.signal,
+    );
 
-      return response.data.results.map((result) => ({
-        ...documents[result.index],
-        score: result.relevance_score,
-      }));
-    });
+    return data.results.map((result) => ({
+      ...documents[result.index],
+      score: result.relevance_score,
+    }));
   }
 
   /**

package/src/connectors/connector-utils.js

@@ -0,0 +1,302 @@
+'use strict';
+
+/**
+ * Shared helpers for the built-in connectors: AbortSignal handling, retry
+ * classification and backoff, and API key resolution.
+ *
+ * Everything here must run on Node 18 (no AbortSignal.any, no crypto.hash).
+ */
+
+/** HTTP status codes that are worth retrying. */
+const RETRYABLE_STATUS = new Set([408, 409, 425, 429]);
+
+/** Error codes produced by Node, axios and undici for transient network failures. */
+const RETRYABLE_CODES = new Set([
+  'ECONNRESET',
+  'ECONNREFUSED',
+  'ETIMEDOUT',
+  'ECONNABORTED',
+  'EPIPE',
+  'EAI_AGAIN',
+  'ENETUNREACH',
+  'ENETDOWN',
+  'EHOSTUNREACH',
+  'ERR_NETWORK',
+]);
+
+/** Error names (openai SDK) that signal a connection-level failure. */
+const RETRYABLE_NAMES = new Set([
+  'APIConnectionError',
+  'APIConnectionTimeoutError',
+]);
+
+/**
+ * Create an AbortError for the given signal.
+ * Reuses signal.reason when it already is an AbortError.
+ *
+ * @param {AbortSignal} [signal]
+ * @returns {Error} Error with name 'AbortError' and code 'ABORT_ERR'
+ */
+function createAbortError(signal) {
+  const reason = signal ? signal.reason : undefined;
+  if (reason instanceof Error && reason.name === 'AbortError') {
+    return reason;
+  }
+  const err = new Error('The operation was aborted');
+  err.name = 'AbortError';
+  err.code = 'ABORT_ERR';
+  if (reason !== undefined) err.cause = reason;
+  return err;
+}
+
+/**
+ * @param {*} err
+ * @returns {boolean} True when err represents a cancellation.
+ */
+function isAbortError(err) {
+  if (!err) return false;
+  return (
+    err.name === 'AbortError' ||
+    err.code === 'ABORT_ERR' ||
+    err.name === 'CanceledError' ||
+    err.code === 'ERR_CANCELED' ||
+    err.name === 'APIUserAbortError'
+  );
+}
+
+/**
+ * Throw an AbortError when the signal is already aborted.
+ * @param {AbortSignal} [signal]
+ */
+function throwIfAborted(signal) {
+  if (signal && signal.aborted) {
+    throw createAbortError(signal);
+  }
+}
+
+/**
+ * Race a promise against an AbortSignal so callers reject promptly even when
+ * the underlying SDK or HTTP client ignores the signal. The abort listener is
+ * removed once the promise settles.
+ *
+ * @template T
+ * @param {Promise<T>|(() => Promise<T>)} work - Promise or function returning one
+ * @param {AbortSignal} [signal]
+ * @returns {Promise<T>}
+ */
+function raceAbort(work, signal) {
+  if (!signal) {
+    return typeof work === 'function' ? Promise.resolve().then(work) : work;
+  }
+  if (signal.aborted) {
+    return Promise.reject(createAbortError(signal));
+  }
+  const promise =
+    typeof work === 'function' ? Promise.resolve().then(work) : work;
+  return new Promise((resolve, reject) => {
+    const onAbort = () => reject(createAbortError(signal));
+    signal.addEventListener('abort', onAbort, { once: true });
+    promise.then(
+      (value) => {
+        signal.removeEventListener('abort', onAbort);
+        resolve(value);
+      },
+      (err) => {
+        signal.removeEventListener('abort', onAbort);
+        // Normalise SDK/axios cancellation errors to AbortError.
+        reject(
+          signal.aborted || isAbortError(err) ? createAbortError(signal) : err,
+        );
+      },
+    );
+  });
+}
+
+/**
+ * Sleep that rejects with AbortError as soon as the signal aborts.
+ * @param {number} ms
+ * @param {AbortSignal} [signal]
+ * @returns {Promise<void>}
+ */
+function abortableSleep(ms, signal) {
+  if (signal && signal.aborted) {
+    return Promise.reject(createAbortError(signal));
+  }
+  return new Promise((resolve, reject) => {
+    let onAbort = null;
+    const timer = setTimeout(() => {
+      if (signal && onAbort) signal.removeEventListener('abort', onAbort);
+      resolve();
+    }, ms);
+    if (signal) {
+      onAbort = () => {
+        clearTimeout(timer);
+        reject(createAbortError(signal));
+      };
+      signal.addEventListener('abort', onAbort, { once: true });
+    }
+  });
+}
+
+/**
+ * Extract the HTTP status from an openai SDK error or an axios error.
+ * @param {*} err
+ * @returns {number|undefined}
+ */
+function getErrorStatus(err) {
+  if (!err) return undefined;
+  const status =
+    (err.response && err.response.status) ?? err.status ?? err.statusCode;
+  return typeof status === 'number' ? status : undefined;
+}
+
+/**
+ * Read a header from a plain object, an AxiosHeaders instance or a fetch Headers.
+ * @param {*} headers
+ * @param {string} name - Lower-case header name
+ * @returns {string|undefined}
+ */
+function getHeader(headers, name) {
+  if (!headers) return undefined;
+  if (typeof headers.get === 'function') {
+    const value = headers.get(name);
+    if (value != null) return String(value);
+  }
+  for (const key of Object.keys(headers)) {
+    if (key.toLowerCase() === name && headers[key] != null) {
+      return String(headers[key]);
+    }
+  }
+  return undefined;
+}
+
+/**
+ * Parse Retry-After / retry-after-ms from an error's response headers.
+ *
+ * @param {*} err
+ * @param {number} [now=Date.now()]
+ * @returns {number|undefined} Delay in milliseconds, or undefined when absent
+ */
+function getRetryAfterMs(err, now = Date.now()) {
+  if (!err) return undefined;
+  const headers = (err.response && err.response.headers) || err.headers;
+  const ms = getHeader(headers, 'retry-after-ms');
+  if (ms !== undefined && /^\d+(\.\d+)?$/.test(ms.trim())) {
+    return Math.max(0, Number(ms));
+  }
+  const value = getHeader(headers, 'retry-after');
+  if (value === undefined) return undefined;
+  const trimmed = value.trim();
+  if (/^\d+(\.\d+)?$/.test(trimmed)) {
+    return Math.max(0, Number(trimmed) * 1000);
+  }
+  const date = Date.parse(trimmed);
+  if (!Number.isNaN(date)) return Math.max(0, date - now);
+  return undefined;
+}
+
+/**
+ * Default retry classification: network failures and HTTP
+ * 408/409/425/429/5xx are retried; aborts, other 4xx (400, 401, 403, 404,
+ * 422, ...) and errors without a status or network code are not.
+ *
+ * @param {*} err
+ * @returns {boolean}
+ */
+function isRetryableError(err) {
+  if (!err || isAbortError(err)) return false;
+  if (err.code === 'MISSING_API_KEY') return false;
+  const status = getErrorStatus(err);
+  if (status !== undefined) {
+    return RETRYABLE_STATUS.has(status) || (status >= 500 && status <= 599);
+  }
+  const code = err.code || (err.cause && err.cause.code);
+  if (typeof code === 'string') {
+    if (RETRYABLE_CODES.has(code) || code.startsWith('UND_ERR_')) return true;
+  }
+  if (RETRYABLE_NAMES.has(err.name)) return true;
+  // fetch() network failures surface as TypeError('fetch failed').
+  if (err.name === 'TypeError' && /fetch failed/i.test(err.message || '')) {
+    return true;
+  }
+  return false;
+}
+
+/**
+ * Compute the delay before the next retry attempt.
+ * Honours Retry-After (capped at maxDelay); otherwise uses exponential backoff
+ * with full jitter: random() * min(maxDelay, baseDelay * 2^attempt).
+ *
+ * @param {*} err - The error that triggered the retry
+ * @param {number} attempt - Zero-based retry attempt
+ * @param {object} [opts]
+ * @param {number} [opts.baseDelay=1000]
+ * @param {number} [opts.maxDelay=30000]
+ * @param {() => number} [opts.random=Math.random]
+ * @returns {number} Delay in milliseconds
+ */
+function computeRetryDelay(err, attempt, opts = {}) {
+  const baseDelay = opts.baseDelay ?? 1000;
+  const maxDelay = opts.maxDelay ?? 30000;
+  const random = opts.random || Math.random;
+  const retryAfter = getRetryAfterMs(err);
+  if (retryAfter !== undefined) {
+    return Math.min(retryAfter, maxDelay);
+  }
+  const exp = Math.min(maxDelay, baseDelay * 2 ** attempt);
+  return Math.floor(random() * exp);
+}
+
+/**
+ * Resolve an API key at construction time. An explicit option (even '')
+ * wins over the environment variable.
+ *
+ * @param {object} options - Connector options
+ * @param {string} envVar - Environment variable name
+ * @returns {string}
+ */
+function resolveApiKey(options, envVar) {
+  if (options && options.apiKey !== undefined) return options.apiKey || '';
+  return process.env[envVar] || '';
+}
+
+/**
+ * Error thrown when a provider call is attempted without an API key.
+ * @param {string} provider - Display name, e.g. 'OpenAI'
+ * @param {string} envVar - Environment variable name
+ * @returns {Error}
+ */
+function missingApiKeyError(provider, envVar) {
+  const err = new Error(
+    `${provider} API key required. Set ${envVar} or pass the apiKey option.`,
+  );
+  err.code = 'MISSING_API_KEY';
+  return err;
+}
+
+/**
+ * Build request options carrying the signal only when one was supplied, so
+ * calls without a signal look exactly as before.
+ * @param {AbortSignal} [signal]
+ * @returns {object}
+ */
+function signalOptions(signal) {
+  return signal ? { signal } : {};
+}
+
+module.exports = {
+  createAbortError,
+  isAbortError,
+  throwIfAborted,
+  raceAbort,
+  abortableSleep,
+  getErrorStatus,
+  getRetryAfterMs,
+  isRetryableError,
+  computeRetryDelay,
+  resolveApiKey,
+  missingApiKeyError,
+  signalOptions,
+  RETRYABLE_STATUS,
+  RETRYABLE_CODES,
+};

package/src/connectors/local-embedder.js

@@ -1,7 +1,9 @@
 'use strict';
+const { digest } = require('../cache/component-keys');
 
 const { BaseConnector } = require('./base-connector');
 const { tokenize } = require('../evaluate/scoring');
+const { throwIfAborted } = require('./connector-utils');
 
 /**
  * Default configuration for the local TF-IDF embedder.
@@ -41,6 +43,7 @@ class LocalEmbedder extends BaseConnector {
    * @param {Array<string|{ content?: string, text?: string }>} documents - Training documents
    */
   train(documents) {
+    this._fingerprint = null; // the cache key changes with the training
     const docs = documents.map((doc) => {
       if (typeof doc === 'string') return doc;
       return doc.content || doc.text || String(doc);
@@ -61,7 +64,9 @@ class LocalEmbedder extends BaseConnector {
     const idfEntries = [];
     for (const [token, freq] of docFreq) {
       if (freq < this.config.minDocFrequency) continue;
-      const idfValue = Math.log(this.docCount / (freq + 1));
+      // Smoothed IDF stays positive; log(N / (df + 1)) is <= 0 for every
+      // token in small corpora (N <= 2), which produced all-zero vectors.
+      const idfValue = Math.log(1 + this.docCount / freq);
       idfEntries.push([token, idfValue]);
     }
 
@@ -86,11 +91,16 @@ class LocalEmbedder extends BaseConnector {
 
   /**
    * Compute a TF-IDF vector for the given text.
-   * The vector length equals the vocabulary size.
+   * The vector length equals the vocabulary size. Synchronous; when
+   * options.signal is already aborted it throws an AbortError (which an
+   * awaiting caller sees as a rejection).
    * @param {string} text - Text to embed
+   * @param {Object} [options={}]
+   * @param {AbortSignal} [options.signal] - Cancellation signal
    * @returns {number[]} L2-normalized TF-IDF vector
    */
-  embed(text) {
+  embed(text, options = {}) {
+    throwIfAborted(options && options.signal);
     const vocabSize = this.vocabulary.size;
 
     if (vocabSize === 0) {
@@ -137,10 +147,58 @@ class LocalEmbedder extends BaseConnector {
    * Compute a TF-IDF vector for a query string.
    * Alias for {@link embed}.
    * @param {string} query - Query text to embed
+   * @param {Object} [options={}] - Same as embed()
    * @returns {number[]} L2-normalized TF-IDF vector
    */
-  embedQuery(query) {
-    return this.embed(query);
+  embedQuery(query, options = {}) {
+    return this.embed(query, options);
+  }
+
+  /**
+   * Serializable vocabulary and IDF weights, so a trained embedder can be
+   * restored in another process (e.g. by FileVectorStore-backed CLI queries).
+   *
+   * @returns {{docCount: number, vocabulary: Array<[string, number]>, idf: Array<[string, number]>}}
+   */
+  getState() {
+    return {
+      docCount: this.docCount,
+      vocabulary: [...this.vocabulary],
+      idf: [...this.idf],
+    };
+  }
+
+  /**
+   * Restore state produced by getState().
+   * @param {{docCount: number, vocabulary: Array<[string, number]>, idf: Array<[string, number]>}} state
+   */
+  setState(state) {
+    this._fingerprint = null;
+    this.docCount = state.docCount || 0;
+    this.vocabulary = new Map(state.vocabulary || []);
+    this.idf = new Map(state.idf || []);
+  }
+
+  /**
+   * Cache identity for createRagPipeline({ cache }): a fingerprint of the
+   * trained vocabulary and IDF weights, so retraining invalidates cached
+   * query vectors (also across processes sharing a Redis cache).
+   * @param {'embedder'} role
+   * @returns {string|null}
+   */
+  getCacheKey(role) {
+    if (role !== 'embedder') return null;
+    if (!this._fingerprint) {
+      this._fingerprint = digest(
+        JSON.stringify([this.docCount, [...this.vocabulary], [...this.idf]]),
+      );
+    }
+    return `local-embedder|${this._fingerprint}`;
+  }
+
+  /** @returns {number} Vector length produced by embed() once trained */
+  get dimensions() {
+    return this.vocabulary.size;
   }
 
   /**