@devilsdev/rag-pipeline-utils 2.3.1 → 2.4.0

package/src/cost/token-budget.js

@@ -0,0 +1,245 @@
+'use strict';
+
+const { EventEmitter } = require('events');
+
+/**
+ * Default budget configuration.
+ * @type {{maxTokens: number, maxCost: number, softLimitPercent: number, trackingWindow: number}}
+ */
+const DEFAULT_CONFIG = {
+  maxTokens: Infinity,
+  maxCost: Infinity,
+  softLimitPercent: 0.8, // warn at 80% of limit
+  trackingWindow: 0, // 0 = no windowing, else ms for sliding window
+};
+
+/**
+ * Configurable token and cost budgets with hard and soft limits.
+ *
+ * Emits:
+ * - `usage`    — after every {@link record} call with the recorded entry.
+ * - `warning`  — when usage crosses the soft limit threshold.
+ * - `exceeded` — when usage crosses a hard limit.
+ *
+ * @example
+ * const budget = new TokenBudget({ maxTokens: 100000, maxCost: 5.0 });
+ * budget.on('warning', (msg) => console.warn(msg));
+ * const status = budget.check(2000, 0.10);
+ * if (status.allowed) budget.record(2000, 0.10);
+ */
+class TokenBudget extends EventEmitter {
+  /**
+   * @param {object} [options]
+   * @param {number} [options.maxTokens=Infinity] - Hard token limit.
+   * @param {number} [options.maxCost=Infinity] - Hard cost limit (USD).
+   * @param {number} [options.softLimitPercent=0.8] - Fraction (0-1) at which soft-limit warnings fire.
+   * @param {number} [options.trackingWindow=0] - Sliding window in ms (0 disables).
+   */
+  constructor(options = {}) {
+    super();
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    this.usage = { tokens: 0, cost: 0, operations: 0, history: [] };
+  }
+
+  /**
+   * Prune history entries outside the sliding window.
+   * @private
+   */
+  _pruneWindow() {
+    if (this.config.trackingWindow <= 0) return;
+
+    const cutoff = Date.now() - this.config.trackingWindow;
+    this.usage.history = this.usage.history.filter(
+      (h) => h.timestamp >= cutoff,
+    );
+
+    // Recompute totals from remaining history
+    this.usage.tokens = this.usage.history.reduce(
+      (sum, h) => sum + h.tokens,
+      0,
+    );
+    this.usage.cost = this.usage.history.reduce((sum, h) => sum + h.cost, 0);
+    this.usage.operations = this.usage.history.length;
+  }
+
+  /**
+   * Check whether a prospective usage fits within the budget.
+   *
+   * @param {number} [tokens=0] - Tokens about to be consumed.
+   * @param {number} [cost=0] - Cost about to be incurred.
+   * @returns {{allowed: boolean, reason?: string, remaining: {tokens: number, cost: number},
+   *            warnings: string[], usage: {tokens: number, cost: number, operations: number}}}
+   */
+  check(tokens = 0, cost = 0) {
+    this._pruneWindow();
+
+    const projectedTokens = this.usage.tokens + tokens;
+    const projectedCost = this.usage.cost + cost;
+    const warnings = [];
+
+    // Hard-limit checks
+    if (projectedTokens > this.config.maxTokens) {
+      return {
+        allowed: false,
+        reason: `Token limit exceeded: ${projectedTokens} > ${this.config.maxTokens}`,
+        remaining: {
+          tokens: Math.max(0, this.config.maxTokens - this.usage.tokens),
+          cost: Math.max(0, this.config.maxCost - this.usage.cost),
+        },
+        warnings,
+        usage: {
+          tokens: this.usage.tokens,
+          cost: this.usage.cost,
+          operations: this.usage.operations,
+        },
+      };
+    }
+
+    if (projectedCost > this.config.maxCost) {
+      return {
+        allowed: false,
+        reason: `Cost limit exceeded: ${projectedCost.toFixed(4)} > ${this.config.maxCost}`,
+        remaining: {
+          tokens: Math.max(0, this.config.maxTokens - this.usage.tokens),
+          cost: Math.max(0, this.config.maxCost - this.usage.cost),
+        },
+        warnings,
+        usage: {
+          tokens: this.usage.tokens,
+          cost: this.usage.cost,
+          operations: this.usage.operations,
+        },
+      };
+    }
+
+    // Soft-limit checks
+    const softTokens = this.config.maxTokens * this.config.softLimitPercent;
+    const softCost = this.config.maxCost * this.config.softLimitPercent;
+
+    if (isFinite(this.config.maxTokens) && projectedTokens > softTokens) {
+      warnings.push(
+        `Approaching token limit: ${projectedTokens} / ${this.config.maxTokens} (${((projectedTokens / this.config.maxTokens) * 100).toFixed(1)}%)`,
+      );
+    }
+
+    if (isFinite(this.config.maxCost) && projectedCost > softCost) {
+      warnings.push(
+        `Approaching cost limit: $${projectedCost.toFixed(4)} / $${this.config.maxCost} (${((projectedCost / this.config.maxCost) * 100).toFixed(1)}%)`,
+      );
+    }
+
+    return {
+      allowed: true,
+      remaining: {
+        tokens: isFinite(this.config.maxTokens)
+          ? this.config.maxTokens - projectedTokens
+          : Infinity,
+        cost: isFinite(this.config.maxCost)
+          ? this.config.maxCost - projectedCost
+          : Infinity,
+      },
+      warnings,
+      usage: {
+        tokens: this.usage.tokens,
+        cost: this.usage.cost,
+        operations: this.usage.operations,
+      },
+    };
+  }
+
+  /**
+   * Record actual token and cost usage.
+   *
+   * @param {number} tokens - Tokens consumed.
+   * @param {number} [cost=0] - Cost incurred (USD).
+   */
+  record(tokens, cost = 0) {
+    const entry = { tokens, cost, timestamp: Date.now() };
+    this.usage.history.push(entry);
+    this.usage.tokens += tokens;
+    this.usage.cost += cost;
+    this.usage.operations += 1;
+
+    this.emit('usage', entry);
+
+    // Soft-limit warnings
+    if (isFinite(this.config.maxTokens)) {
+      const softTokens = this.config.maxTokens * this.config.softLimitPercent;
+      if (
+        this.usage.tokens > softTokens &&
+        this.usage.tokens <= this.config.maxTokens
+      ) {
+        this.emit(
+          'warning',
+          `Token usage at ${((this.usage.tokens / this.config.maxTokens) * 100).toFixed(1)}% of limit`,
+        );
+      }
+    }
+
+    if (isFinite(this.config.maxCost)) {
+      const softCost = this.config.maxCost * this.config.softLimitPercent;
+      if (
+        this.usage.cost > softCost &&
+        this.usage.cost <= this.config.maxCost
+      ) {
+        this.emit(
+          'warning',
+          `Cost usage at ${((this.usage.cost / this.config.maxCost) * 100).toFixed(1)}% of limit`,
+        );
+      }
+    }
+
+    // Hard-limit exceeded
+    if (this.usage.tokens > this.config.maxTokens) {
+      this.emit(
+        'exceeded',
+        `Token limit exceeded: ${this.usage.tokens} > ${this.config.maxTokens}`,
+      );
+    }
+
+    if (this.usage.cost > this.config.maxCost) {
+      this.emit(
+        'exceeded',
+        `Cost limit exceeded: ${this.usage.cost.toFixed(4)} > ${this.config.maxCost}`,
+      );
+    }
+  }
+
+  /**
+   * Reset all usage counters to zero.
+   */
+  reset() {
+    this.usage = { tokens: 0, cost: 0, operations: 0, history: [] };
+  }
+
+  /**
+   * Return current usage summary.
+   *
+   * @returns {{tokens: number, cost: number, operations: number,
+   *            limits: {maxTokens: number, maxCost: number},
+   *            remaining: {tokens: number, cost: number}}}
+   */
+  getUsage() {
+    this._pruneWindow();
+
+    return {
+      tokens: this.usage.tokens,
+      cost: this.usage.cost,
+      operations: this.usage.operations,
+      limits: {
+        maxTokens: this.config.maxTokens,
+        maxCost: this.config.maxCost,
+      },
+      remaining: {
+        tokens: isFinite(this.config.maxTokens)
+          ? Math.max(0, this.config.maxTokens - this.usage.tokens)
+          : Infinity,
+        cost: isFinite(this.config.maxCost)
+          ? Math.max(0, this.config.maxCost - this.usage.cost)
+          : Infinity,
+      },
+    };
+  }
+}
+
+module.exports = { TokenBudget, DEFAULT_CONFIG };

package/src/dag/core/execution-scheduler.js

@@ -14,6 +14,7 @@
 'use strict';
 
 const ErrorContext = require('./error-context.js');
+const { logger } = require('../../utils/logger');
 
 /**
  * ExecutionScheduler class for managing DAG node execution
@@ -251,7 +252,7 @@ class ExecutionScheduler {
 
         // Handle non-critical failures with warnings and no result nullification
         if (isNonCritical || (continueOnError && !requiredIds.has(node.id))) {
-          console.warn(`Non-critical node failure: ${nodeErr.message}`);
+          logger.warn(`Non-critical node failure: ${nodeErr.message}`);
           errors.set(node.id, nodeErr);
           // Don't set result for failed optional nodes - return special symbol so they don't appear in results
           return Symbol('failed-optional-node');

package/src/dag/core/topology-validation.js

@@ -410,6 +410,3 @@ module.exports = {
   validateDAG,
   validateTopology,
 };
-
-// Default export for convenience
-module.exports.default = module.exports;

package/src/dag/dag-engine.js

@@ -15,6 +15,7 @@ const {
 } = require('./core/topology-validation.js');
 const ExecutionScheduler = require('./core/execution-scheduler.js');
 const ErrorContext = require('./core/error-context.js');
+const { logger } = require('../utils/logger');
 
 class DAG {
   constructor() {
@@ -458,14 +459,12 @@ class DAG {
               continue; // no error recorded yet
             }
             // FINAL failure path - warn EXACT string and keep going
-            // eslint-disable-next-line no-console
-            console.warn(`Node ${node.id} failed during resume: ${e.message}`);
+            logger.warn(`Node ${node.id} failed during resume: ${e.message}`);
             keepRetrying = false; // Don't throw, just continue
           }
         }
       } catch (error) {
-        // eslint-disable-next-line no-console
-        console.warn(`Node ${node.id} failed during resume: ${error.message}`);
+        logger.warn(`Node ${node.id} failed during resume: ${error.message}`);
         continue;
       }
     }
@@ -553,4 +552,3 @@ class DAG {
 
 // Export classes
 module.exports = { DAG, DAGNode };
-module.exports.default = module.exports;

package/src/dag/dag-node.js

@@ -73,4 +73,3 @@ class DAGNode {
 }
 
 module.exports = DAGNode;
-module.exports.default = module.exports;

package/src/dag/enhanced-dag-engine.js

@@ -589,108 +589,36 @@ class EnhancedDAGEngine {
         // Validate DAG before execution
         this.validate();
 
-        // Initialize execution state
+        // Initialize execution state and context
         const results = new Map();
         const errors = new Map();
         const order = this.getTopologicalOrder();
-
-        // Create execution context
-        const context = {
+        const context = this._buildExecutionContext(
           results,
           errors,
           seed,
-          continueOnError: config.continueOnError,
-          gracefulDegradation: config.gracefulDegradation,
-          retryFailedNodes: config.retryFailedNodes,
-          maxRetries: config.maxRetries,
-        };
-
-        // Execute nodes with concurrency control
-        const semaphore = new Semaphore(config.concurrency);
-        const nodePromises = new Map();
-
-        for (const node of order) {
-          const promise = semaphore.acquire().then(async (release) => {
-            try {
-              // Wait for dependencies to complete
-              const depPromises = node.inputs
-                .map((dep) => nodePromises.get(dep.id))
-                .filter(Boolean);
-
-              await Promise.all(depPromises);
-
-              // Execute the node
-              const result = await this.executeNode(node, context);
-
-              if (
-                typeof result !== 'symbol' ||
-                result.description !== 'failed-optional-node'
-              ) {
-                results.set(node.id, result);
-              }
-
-              return result;
-            } finally {
-              release();
-            }
-          });
-
-          nodePromises.set(node.id, promise);
-        }
-
-        // Wait for all nodes to complete
-        const allResults = await Promise.allSettled(
-          Array.from(nodePromises.values()),
+          config,
         );
 
+        // Schedule and run all nodes with concurrency control
+        await this._scheduleNodes(order, context, config);
+
         const executionTime = Date.now() - startTime;
 
         // Update engine metrics
-        if (this.enableMetrics) {
-          this.engineMetrics.totalExecutions++;
-          this.engineMetrics.totalExecutionTime += executionTime;
-          this.engineMetrics.avgExecutionTime =
-            this.engineMetrics.totalExecutionTime /
-            this.engineMetrics.totalExecutions;
-
-          if (errors.size === 0) {
-            this.engineMetrics.successfulExecutions++;
-          } else {
-            this.engineMetrics.failedExecutions++;
-          }
-        }
+        this._recordExecutionMetrics(executionTime, errors.size === 0);
 
-        // Create execution summary
-        const summary = {
+        // Build and store execution summary
+        const summary = this._buildExecutionSummary(
           executionId,
-          success: errors.size === 0,
           executionTime,
-          nodesExecuted: results.size,
-          totalNodes: this.nodes.size,
-          errors: Array.from(errors.entries()).map(([id, error]) => ({
-            nodeId: id,
-            message: error.message,
-            code: error.code,
-          })),
-        };
-
-        // Store execution history
-        if (this.enableTracing) {
-          this.executionHistory.push({
-            ...summary,
-            timestamp: new Date().toISOString(),
-            results: Object.fromEntries(results),
-          });
-
-          // Keep only last 100 executions
-          if (this.executionHistory.length > 100) {
-            this.executionHistory.shift();
-          }
-        }
+          results,
+          errors,
+        );
+        this._storeExecutionHistory(summary, results);
 
         this.logger.info('DAG execution completed', summary);
 
-        // Return results
         return {
           results: Object.fromEntries(results),
           errors: Object.fromEntries(errors),
@@ -699,21 +627,168 @@ class EnhancedDAGEngine {
         };
       } catch (error) {
         const executionTime = Date.now() - startTime;
+        this._recordExecutionFailure(executionId, error, executionTime);
+        throw error;
+      }
+    });
+  }
 
-        if (this.enableMetrics) {
-          this.engineMetrics.totalExecutions++;
-          this.engineMetrics.failedExecutions++;
+  /**
+   * Build the execution context passed to each node during DAG execution.
+   * @param {Map} results - Map to store completed node results
+   * @param {Map} errors - Map to store node execution errors
+   * @param {any} seed - Initial seed value for source nodes
+   * @param {Object} config - Merged execution configuration
+   * @returns {ExecutionContext} Execution context object
+   * @private
+   */
+  _buildExecutionContext(results, errors, seed, config) {
+    return {
+      results,
+      errors,
+      seed,
+      continueOnError: config.continueOnError,
+      gracefulDegradation: config.gracefulDegradation,
+      retryFailedNodes: config.retryFailedNodes,
+      maxRetries: config.maxRetries,
+    };
+  }
+
+  /**
+   * Schedule and execute all nodes in topological order with concurrency control.
+   * Each node waits for its dependencies before executing.
+   * @param {EnhancedDAGNode[]} order - Nodes in topological execution order
+   * @param {ExecutionContext} context - Execution context with results and errors maps
+   * @param {Object} config - Execution configuration including concurrency limit
+   * @returns {Promise<void>} Resolves when all nodes have settled
+   * @private
+   */
+  async _scheduleNodes(order, context, config) {
+    const { results } = context;
+    const semaphore = new Semaphore(config.concurrency);
+    const nodePromises = new Map();
+
+    for (const node of order) {
+      const promise = semaphore.acquire().then(async (release) => {
+        try {
+          // Wait for dependencies to complete
+          const depPromises = node.inputs
+            .map((dep) => nodePromises.get(dep.id))
+            .filter(Boolean);
+
+          await Promise.all(depPromises);
+
+          // Execute the node
+          const result = await this.executeNode(node, context);
+
+          if (
+            typeof result !== 'symbol' ||
+            result.description !== 'failed-optional-node'
+          ) {
+            results.set(node.id, result);
+          }
+
+          return result;
+        } finally {
+          release();
         }
+      });
 
-        this.logger.error('DAG execution failed', {
-          executionId,
-          error: error.message,
-          executionTime,
-          stack: error.stack,
-        });
+      nodePromises.set(node.id, promise);
+    }
 
-        throw error;
-      }
+    await Promise.allSettled(Array.from(nodePromises.values()));
+  }
+
+  /**
+   * Record engine-level metrics after a completed execution.
+   * @param {number} executionTime - Total execution time in milliseconds
+   * @param {boolean} success - Whether the execution completed without errors
+   * @private
+   */
+  _recordExecutionMetrics(executionTime, success) {
+    if (!this.enableMetrics) {
+      return;
+    }
+
+    this.engineMetrics.totalExecutions++;
+    this.engineMetrics.totalExecutionTime += executionTime;
+    this.engineMetrics.avgExecutionTime =
+      this.engineMetrics.totalExecutionTime /
+      this.engineMetrics.totalExecutions;
+
+    if (success) {
+      this.engineMetrics.successfulExecutions++;
+    } else {
+      this.engineMetrics.failedExecutions++;
+    }
+  }
+
+  /**
+   * Build a summary object describing the outcome of a DAG execution.
+   * @param {string} executionId - Unique execution identifier
+   * @param {number} executionTime - Total execution time in milliseconds
+   * @param {Map} results - Map of successful node results
+   * @param {Map} errors - Map of node execution errors
+   * @returns {Object} Execution summary
+   * @private
+   */
+  _buildExecutionSummary(executionId, executionTime, results, errors) {
+    return {
+      executionId,
+      success: errors.size === 0,
+      executionTime,
+      nodesExecuted: results.size,
+      totalNodes: this.nodes.size,
+      errors: Array.from(errors.entries()).map(([id, error]) => ({
+        nodeId: id,
+        message: error.message,
+        code: error.code,
+      })),
+    };
+  }
+
+  /**
+   * Store an execution summary in the tracing history buffer.
+   * Keeps only the most recent 100 entries.
+   * @param {Object} summary - Execution summary to store
+   * @param {Map} results - Map of node results to include in history
+   * @private
+   */
+  _storeExecutionHistory(summary, results) {
+    if (!this.enableTracing) {
+      return;
+    }
+
+    this.executionHistory.push({
+      ...summary,
+      timestamp: new Date().toISOString(),
+      results: Object.fromEntries(results),
+    });
+
+    if (this.executionHistory.length > 100) {
+      this.executionHistory.shift();
+    }
+  }
+
+  /**
+   * Record metrics and log details when a DAG execution fails with an exception.
+   * @param {string} executionId - Unique execution identifier
+   * @param {Error} error - The error that caused the failure
+   * @param {number} executionTime - Time elapsed before the failure
+   * @private
+   */
+  _recordExecutionFailure(executionId, error, executionTime) {
+    if (this.enableMetrics) {
+      this.engineMetrics.totalExecutions++;
+      this.engineMetrics.failedExecutions++;
+    }
+
+    this.logger.error('DAG execution failed', {
+      executionId,
+      error: error.message,
+      executionTime,
+      stack: error.stack,
     });
   }
 
@@ -817,6 +892,3 @@ module.exports = {
   Semaphore,
   createEnhancedDAG,
 };
-
-// CommonJS/ESM interop
-module.exports.default = module.exports;