@devilsdev/rag-pipeline-utils 2.3.1 → 2.4.0

package/src/streaming/stream-router.js

@@ -0,0 +1,210 @@
+'use strict';
+
+/**
+ * Stream Router
+ *
+ * Wraps a RAG pipeline and routes streaming output to the appropriate
+ * transport adapter (SSE or WebSocket) based on the incoming request.
+ *
+ * @module streaming/stream-router
+ * @since 2.4.0
+ */
+
+const { EventEmitter } = require('events');
+const { SSEAdapter } = require('./sse-adapter');
+const { WebSocketAdapter } = require('./websocket-adapter');
+
+/**
+ * Default configuration for the stream router
+ * @type {Object}
+ */
+const DEFAULT_CONFIG = {
+  /** @type {string} Transport to use when auto-detection is inconclusive */
+  defaultTransport: 'sse',
+  /** @type {boolean} Whether to set CORS headers on responses */
+  enableCORS: true,
+  /** @type {string} Value for the Access-Control-Allow-Origin header */
+  corsOrigin: '*',
+};
+
+/**
+ * Routes pipeline streaming output to SSE or WebSocket transports.
+ *
+ * Provides both an Express/Node-HTTP-compatible request handler via
+ * {@link StreamRouter#createHandler} and lower-level convenience methods
+ * ({@link StreamRouter#streamSSE}, {@link StreamRouter#streamWS}) for
+ * manual integration.
+ *
+ * @extends EventEmitter
+ *
+ * @example
+ *   const router = new StreamRouter({ enableCORS: true });
+ *   const handler = router.createHandler(pipeline);
+ *   app.post('/api/query', handler);
+ */
+class StreamRouter extends EventEmitter {
+  /**
+   * @param {Object} [options] - Configuration overrides
+   * @param {string} [options.defaultTransport='sse'] - Fallback transport type
+   * @param {boolean} [options.enableCORS=true] - Enable CORS headers
+   * @param {string} [options.corsOrigin='*'] - CORS origin value
+   * @param {Object} [options.sseOptions] - Options forwarded to SSEAdapter
+   * @param {Object} [options.wsOptions] - Options forwarded to WebSocketAdapter
+   */
+  constructor(options = {}) {
+    super();
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    this.sseAdapter = new SSEAdapter(options.sseOptions);
+    this.wsAdapter = new WebSocketAdapter(options.wsOptions);
+  }
+
+  /**
+   * Create an HTTP request handler that runs the pipeline in streaming mode
+   * and routes output to the appropriate transport.
+   *
+   * The returned function is compatible with Node's `http.createServer`,
+   * Express, and similar frameworks.
+   *
+   * Transport detection:
+   * - `Accept: text/event-stream` header selects SSE
+   * - `Upgrade: websocket` header selects WebSocket (note: actual upgrade
+   *   negotiation is typically handled at a higher level)
+   * - Otherwise falls back to {@link DEFAULT_CONFIG.defaultTransport}
+   *
+   * @param {Object} pipeline - A pipeline object with a `.run()` method
+   * @returns {Function} Async request handler `(req, res) => Promise<void>`
+   */
+  createHandler(pipeline) {
+    /**
+     * @param {import('http').IncomingMessage} req
+     * @param {import('http').ServerResponse} res
+     */
+    const handler = async (req, res) => {
+      try {
+        // --- 1. Parse query -----------------------------------------------
+        const query =
+          (req.body && req.body.query) || (req.query && req.query.q) || '';
+
+        if (!query) {
+          res.writeHead(400, { 'Content-Type': 'application/json' });
+          res.end(JSON.stringify({ error: 'Missing query parameter' }));
+          return;
+        }
+
+        // --- 2. CORS headers ----------------------------------------------
+        if (this.config.enableCORS) {
+          res.setHeader('Access-Control-Allow-Origin', this.config.corsOrigin);
+          res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
+          res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Accept');
+        }
+
+        // Handle CORS preflight
+        if (req.method === 'OPTIONS') {
+          res.writeHead(204);
+          res.end();
+          return;
+        }
+
+        // --- 3. Detect transport ------------------------------------------
+        const accept = (req.headers && req.headers.accept) || '';
+        const upgrade = (req.headers && req.headers.upgrade) || '';
+
+        let transport = this.config.defaultTransport;
+        if (accept.includes('text/event-stream')) {
+          transport = 'sse';
+        } else if (upgrade.toLowerCase() === 'websocket') {
+          transport = 'ws';
+        }
+
+        // --- 4. Route to adapter ------------------------------------------
+        if (transport === 'sse') {
+          await this.streamSSE(pipeline, query, res);
+        } else if (transport === 'ws') {
+          // WebSocket upgrade is normally handled at the server level;
+          // send a hint response when received through a regular handler.
+          res.writeHead(426, { 'Content-Type': 'application/json' });
+          res.end(
+            JSON.stringify({
+              error: 'WebSocket upgrade required',
+              hint: 'Use the streamWS() method with an upgraded WebSocket connection',
+            }),
+          );
+        } else {
+          await this.streamSSE(pipeline, query, res);
+        }
+      } catch (err) {
+        this.emit('error', err);
+
+        if (!res.headersSent) {
+          res.writeHead(500, { 'Content-Type': 'application/json' });
+        }
+        if (!res.writableEnded) {
+          res.end(
+            JSON.stringify({ error: err.message || 'Internal server error' }),
+          );
+        }
+      }
+    };
+
+    return handler;
+  }
+
+  /**
+   * Run the pipeline in streaming mode and pipe tokens to an HTTP response
+   * via Server-Sent Events.
+   *
+   * If the pipeline returns a non-iterable result, it is sent as a single
+   * SSE data frame followed by the done sentinel.
+   *
+   * @param {Object} pipeline - Pipeline with a `.run()` method
+   * @param {string} query - User query string
+   * @param {import('http').ServerResponse} httpResponse - Writable HTTP response
+   * @param {Object} [options={}] - Additional pipeline run options
+   * @returns {Promise<void>}
+   */
+  async streamSSE(pipeline, query, httpResponse, options = {}) {
+    const result = await pipeline.run({
+      query,
+      options: { ...options, stream: true },
+    });
+
+    if (result && Symbol.asyncIterator in Object(result)) {
+      await this.sseAdapter.stream(result, httpResponse);
+    } else {
+      // Non-streaming result: send as a single SSE frame then DONE
+      httpResponse.writeHead(200, { 'Content-Type': 'text/event-stream' });
+      httpResponse.write(`data: ${JSON.stringify(result)}\n\n`);
+      httpResponse.write(`data: [DONE]\n\n`);
+      httpResponse.end();
+    }
+  }
+
+  /**
+   * Run the pipeline in streaming mode and pipe tokens over a WebSocket
+   * connection.
+   *
+   * If the pipeline returns a non-iterable result, it is sent as a single
+   * `result` message followed by a `done` message.
+   *
+   * @param {Object} pipeline - Pipeline with a `.run()` method
+   * @param {string} query - User query string
+   * @param {Object} wsConnection - WebSocket-like connection with `.send()`
+   * @param {Object} [options={}] - Additional pipeline run options
+   * @returns {Promise<void>}
+   */
+  async streamWS(pipeline, query, wsConnection, options = {}) {
+    const result = await pipeline.run({
+      query,
+      options: { ...options, stream: true },
+    });
+
+    if (result && Symbol.asyncIterator in Object(result)) {
+      await this.wsAdapter.stream(result, wsConnection);
+    } else {
+      wsConnection.send(JSON.stringify({ type: 'result', data: result }));
+      wsConnection.send(JSON.stringify({ type: 'done' }));
+    }
+  }
+}
+
+module.exports = { StreamRouter, DEFAULT_CONFIG };

package/src/streaming/streaming-embedder.js

@@ -0,0 +1,231 @@
+'use strict';
+
+/**
+ * Streaming Embedder
+ *
+ * Class wrapper that provides an EventEmitter interface around the
+ * {@link module:streaming/embedding-stream createEmbeddingStream} async
+ * generator.  Emits lifecycle events (`start`, `embedding`, `batch`,
+ * `backpressure`, `done`, `error`) so that callers can monitor progress
+ * without manually iterating the generator.
+ *
+ * Also exposes a convenience `embedAll()` method that collects every
+ * embedding into an array, and tracks cumulative statistics across calls.
+ *
+ * @module streaming/streaming-embedder
+ * @since 2.5.0
+ */
+
+const { EventEmitter } = require('events');
+const { createEmbeddingStream } = require('./embedding-stream');
+
+/**
+ * Default configuration for the StreamingEmbedder
+ * @type {Object}
+ */
+const DEFAULT_CONFIG = {
+  /** @type {number} Number of documents per embedding batch */
+  batchSize: 50,
+  /** @type {number} Maximum concurrent batch operations */
+  maxConcurrency: 3,
+  /** @type {number} Maximum heap memory budget in megabytes */
+  maxMemoryMB: 512,
+  /** @type {number} Fraction of maxMemoryMB that triggers backpressure (0-1) */
+  backpressureThreshold: 0.8,
+};
+
+/**
+ * EventEmitter-based wrapper around {@link createEmbeddingStream}.
+ *
+ * @extends EventEmitter
+ *
+ * @fires StreamingEmbedder#start
+ * @fires StreamingEmbedder#embedding
+ * @fires StreamingEmbedder#batch
+ * @fires StreamingEmbedder#backpressure
+ * @fires StreamingEmbedder#done
+ * @fires StreamingEmbedder#error
+ *
+ * @example
+ *   const se = new StreamingEmbedder(myEmbedder, { batchSize: 25 });
+ *   se.on('embedding', (e) => console.log(e.id, e.progress));
+ *   se.on('done', (stats) => console.log('finished', stats));
+ *   const all = await se.embedAll(documents);
+ */
+class StreamingEmbedder extends EventEmitter {
+  /**
+   * @param {Object} embedder - Embedder instance with `.embed()` or `.embedQuery()`
+   * @param {Object} [options] - Override any key from {@link DEFAULT_CONFIG}
+   */
+  constructor(embedder, options = {}) {
+    super();
+
+    if (
+      !embedder ||
+      (typeof embedder.embed !== 'function' &&
+        typeof embedder.embedQuery !== 'function')
+    ) {
+      throw new TypeError(
+        'Embedder must expose an .embed() or .embedQuery() method',
+      );
+    }
+
+    /** @type {Object} */
+    this.embedder = embedder;
+
+    /** @type {Object} */
+    this.config = { ...DEFAULT_CONFIG, ...options };
+
+    /** @type {{ processed: number, failed: number, totalTime: number }} */
+    this.stats = { processed: 0, failed: 0, totalTime: 0 };
+  }
+
+  /**
+   * Async generator that yields embeddings while emitting lifecycle events.
+   *
+   * @param {Array<{id: string, content: string}>} documents - Documents to embed
+   * @yields {{ id: string, vector: number[], index: number, progress: number }}
+   *
+   * @fires StreamingEmbedder#start
+   * @fires StreamingEmbedder#embedding
+   * @fires StreamingEmbedder#batch
+   * @fires StreamingEmbedder#backpressure
+   * @fires StreamingEmbedder#done
+   * @fires StreamingEmbedder#error
+   */
+  async *embedStream(documents) {
+    if (!Array.isArray(documents)) {
+      throw new TypeError('documents must be an array');
+    }
+
+    const startTime = Date.now();
+    const documentCount = documents.length;
+
+    /**
+     * @event StreamingEmbedder#start
+     * @type {{ documentCount: number }}
+     */
+    this.emit('start', { documentCount });
+
+    let batchCount = 0;
+    let itemsInCurrentBatch = 0;
+
+    try {
+      const stream = createEmbeddingStream(
+        documents,
+        this.embedder,
+        this.config,
+      );
+
+      for await (const embedding of stream) {
+        // Track backpressure events
+        if (embedding.backpressure) {
+          /**
+           * @event StreamingEmbedder#backpressure
+           * @type {{ index: number, progress: number }}
+           */
+          this.emit('backpressure', {
+            index: embedding.index,
+            progress: embedding.progress,
+          });
+        }
+
+        // Track errors
+        if (embedding.error) {
+          this.stats.failed += 1;
+          /**
+           * @event StreamingEmbedder#error
+           * @type {{ id: string, error: Error, index: number }}
+           */
+          this.emit('error', {
+            id: embedding.id,
+            error: embedding.error,
+            index: embedding.index,
+          });
+        } else {
+          this.stats.processed += 1;
+        }
+
+        /**
+         * @event StreamingEmbedder#embedding
+         * @type {{ id: string, vector?: number[], index: number, progress: number }}
+         */
+        this.emit('embedding', embedding);
+
+        // Emit batch event every batchSize items
+        itemsInCurrentBatch += 1;
+        if (itemsInCurrentBatch >= this.config.batchSize) {
+          batchCount += 1;
+          /**
+           * @event StreamingEmbedder#batch
+           * @type {{ batchNumber: number, processedSoFar: number, progress: number }}
+           */
+          this.emit('batch', {
+            batchNumber: batchCount,
+            processedSoFar: embedding.index + 1,
+            progress: embedding.progress,
+          });
+          itemsInCurrentBatch = 0;
+        }
+
+        yield embedding;
+      }
+
+      // Emit a final batch event for any remaining items
+      if (itemsInCurrentBatch > 0) {
+        batchCount += 1;
+        this.emit('batch', {
+          batchNumber: batchCount,
+          processedSoFar: documentCount,
+          progress: 1,
+        });
+      }
+    } finally {
+      const elapsed = Date.now() - startTime;
+      this.stats.totalTime += elapsed;
+
+      /**
+       * @event StreamingEmbedder#done
+       * @type {{ processed: number, failed: number, totalTime: number, documentCount: number }}
+       */
+      this.emit('done', {
+        processed: this.stats.processed,
+        failed: this.stats.failed,
+        totalTime: this.stats.totalTime,
+        documentCount,
+      });
+    }
+  }
+
+  /**
+   * Convenience method that collects all embeddings into an array.
+   *
+   * @param {Array<{id: string, content: string}>} documents - Documents to embed
+   * @returns {Promise<Array<{ id: string, vector: number[], index: number, progress: number }>>}
+   */
+  async embedAll(documents) {
+    const results = [];
+    for await (const embedding of this.embedStream(documents)) {
+      results.push(embedding);
+    }
+    return results;
+  }
+
+  /**
+   * Returns a snapshot of the cumulative statistics.
+   *
+   * @returns {{ processed: number, failed: number, totalTime: number }}
+   */
+  getStats() {
+    return { ...this.stats };
+  }
+
+  /**
+   * Resets the cumulative statistics to zero.
+   */
+  resetStats() {
+    this.stats = { processed: 0, failed: 0, totalTime: 0 };
+  }
+}
+
+module.exports = { StreamingEmbedder, DEFAULT_CONFIG };

package/src/streaming/websocket-adapter.js

@@ -0,0 +1,171 @@
+'use strict';
+
+/**
+ * WebSocket Adapter
+ *
+ * Sends AsyncGenerator tokens over a WebSocket connection
+ * for real-time bidirectional streaming of RAG pipeline responses.
+ *
+ * @module streaming/websocket-adapter
+ * @since 2.4.0
+ */
+
+const { EventEmitter } = require('events');
+
+/**
+ * Default configuration for WebSocket streaming
+ * @type {Object}
+ */
+const DEFAULT_CONFIG = {
+  /** @type {string} Message type identifier for token payloads */
+  messageType: 'token',
+  /** @type {string} Message type identifier for stream completion */
+  doneType: 'done',
+  /** @type {string} Message type identifier for error payloads */
+  errorType: 'error',
+};
+
+/**
+ * WebSocket ready-state constants (mirrors the W3C WebSocket API).
+ * @enum {number}
+ * @private
+ */
+const WS_READY_STATE = {
+  OPEN: 1,
+};
+
+/**
+ * Adapts an AsyncGenerator to a WebSocket connection.
+ *
+ * Sends JSON-framed messages for each token, a completion message when
+ * the generator is exhausted, and an error message on failure. Handles
+ * premature connection closure gracefully.
+ *
+ * @extends EventEmitter
+ *
+ * @example
+ *   const adapter = new WebSocketAdapter();
+ *   adapter.on('done', () => console.log('stream complete'));
+ *   await adapter.stream(asyncTokenGenerator, ws);
+ */
+class WebSocketAdapter extends EventEmitter {
+  /**
+   * @param {Object} [options] - Configuration overrides
+   * @param {string} [options.messageType='token'] - Type field for token messages
+   * @param {string} [options.doneType='done'] - Type field for completion message
+   * @param {string} [options.errorType='error'] - Type field for error messages
+   */
+  constructor(options = {}) {
+    super();
+    this.config = { ...DEFAULT_CONFIG, ...options };
+  }
+
+  /**
+   * Stream tokens from an AsyncGenerator over a WebSocket connection.
+   *
+   * Each yielded value should be an object with at least a `token` property
+   * and an optional `done` boolean. Messages are sent as JSON strings via
+   * `wsConnection.send()`.
+   *
+   * @param {AsyncGenerator<{token: string, done?: boolean}>} asyncGenerator - Token source
+   * @param {Object} wsConnection - WebSocket-like object with a `send(data)` method
+   * @param {Function} wsConnection.send - Sends a string message over the socket
+   * @param {number} [wsConnection.readyState] - Optional W3C readyState property
+   * @returns {Promise<void>} Resolves when the stream ends or the connection closes
+   * @throws {TypeError} If wsConnection lacks a send method
+   */
+  async stream(asyncGenerator, wsConnection) {
+    // --- 1. Validate connection -----------------------------------------
+    if (!wsConnection || typeof wsConnection.send !== 'function') {
+      throw new TypeError('wsConnection must have a send() method');
+    }
+
+    let connectionClosed = false;
+
+    // --- 2. Connection lifecycle handlers --------------------------------
+    const onClose = () => {
+      connectionClosed = true;
+    };
+    const onError = (err) => {
+      connectionClosed = true;
+      this.emit('error', err);
+    };
+
+    if (typeof wsConnection.on === 'function') {
+      wsConnection.on('close', onClose);
+      wsConnection.on('error', onError);
+    }
+
+    // --- 3. Iterate and send --------------------------------------------
+    try {
+      this.emit('started');
+
+      for await (const value of asyncGenerator) {
+        if (connectionClosed || !this._isOpen(wsConnection)) {
+          break;
+        }
+
+        if (value && value.done) {
+          this._send(wsConnection, { type: this.config.doneType });
+          this.emit('done');
+          break;
+        }
+
+        const token = value && value.token !== undefined ? value.token : value;
+        this._send(wsConnection, {
+          type: this.config.messageType,
+          data: { token },
+        });
+        this.emit('token', token);
+      }
+
+      // Generator exhausted without explicit done flag
+      if (!connectionClosed && this._isOpen(wsConnection)) {
+        this._send(wsConnection, { type: this.config.doneType });
+        this.emit('done');
+      }
+    } catch (err) {
+      this.emit('error', err);
+
+      if (!connectionClosed && this._isOpen(wsConnection)) {
+        this._send(wsConnection, {
+          type: this.config.errorType,
+          data: { message: err.message || 'Internal streaming error' },
+        });
+      }
+    } finally {
+      if (typeof wsConnection.removeListener === 'function') {
+        wsConnection.removeListener('close', onClose);
+        wsConnection.removeListener('error', onError);
+      }
+    }
+  }
+
+  /**
+   * Check whether the WebSocket connection is still open.
+   *
+   * @param {Object} wsConnection
+   * @returns {boolean}
+   * @private
+   */
+  _isOpen(wsConnection) {
+    if (typeof wsConnection.readyState === 'number') {
+      return wsConnection.readyState === WS_READY_STATE.OPEN;
+    }
+    // If readyState is not available, assume open
+    return true;
+  }
+
+  /**
+   * Send a JSON-serialised message over the WebSocket.
+   *
+   * @param {Object} wsConnection
+   * @param {Object} payload
+   * @private
+   */
+  _send(wsConnection, payload) {
+    wsConnection.send(JSON.stringify(payload));
+  }
+}
+
+module.exports = { WebSocketAdapter, DEFAULT_CONFIG };