@devilsdev/rag-pipeline-utils 2.3.1 → 2.4.0

package/src/guardrails/guardrails-pipeline.js

@@ -0,0 +1,136 @@
+'use strict';
+
+const { EventEmitter } = require('events');
+const { PreRetrievalGuard } = require('./pre-retrieval-guard');
+const { RetrievalGuard } = require('./retrieval-guard');
+const { PostGenerationGuard } = require('./post-generation-guard');
+
+/**
+ * Default configuration for the GuardrailsPipeline.
+ * @type {object}
+ */
+const DEFAULT_CONFIG = {
+  preRetrieval: {}, // PreRetrievalGuard options, or false to disable
+  retrieval: {}, // RetrievalGuard options, or false to disable
+  postGeneration: {}, // PostGenerationGuard options, or false to disable
+  strict: false, // If true, block on any guard failure; if false, warn and continue
+};
+
+/**
+ * Pipeline wrapper that applies all three guardrail layers (pre-retrieval,
+ * retrieval, and post-generation) around an existing RAG pipeline. In strict
+ * mode, any guard failure aborts the pipeline; otherwise issues are collected
+ * as warnings and execution continues.
+ *
+ * @extends EventEmitter
+ */
+class GuardrailsPipeline extends EventEmitter {
+  /**
+   * @param {object} pipeline - The underlying RAG pipeline (must have a .run() method)
+   * @param {object} [options] - Override default configuration
+   * @param {object|false} [options.preRetrieval={}] - PreRetrievalGuard options, or false to disable
+   * @param {object|false} [options.retrieval={}] - RetrievalGuard options, or false to disable
+   * @param {object|false} [options.postGeneration={}] - PostGenerationGuard options, or false to disable
+   * @param {boolean} [options.strict=false] - Whether to block on any guard failure
+   */
+  constructor(pipeline, options = {}) {
+    super();
+
+    if (!pipeline || typeof pipeline.run !== 'function') {
+      throw new Error(
+        'GuardrailsPipeline requires a pipeline with a .run() method',
+      );
+    }
+
+    this.pipeline = pipeline;
+    this.config = { ...DEFAULT_CONFIG, ...options };
+
+    this.preGuard =
+      this.config.preRetrieval !== false
+        ? new PreRetrievalGuard(this.config.preRetrieval || {})
+        : null;
+    this.retrievalGuard =
+      this.config.retrieval !== false
+        ? new RetrievalGuard(this.config.retrieval || {})
+        : null;
+    this.postGuard =
+      this.config.postGeneration !== false
+        ? new PostGenerationGuard(this.config.postGeneration || {})
+        : null;
+  }
+
+  /**
+   * Run the full guarded pipeline.
+   *
+   * @param {object} input - Pipeline input
+   * @param {string} input.query - The user query
+   * @param {Float32Array} [input.queryVector] - Optional precomputed query vector
+   * @param {object} [input.options] - Additional pipeline options
+   * @param {object} [input.context] - Request context (e.g. userPermissions for ACL)
+   * @returns {Promise<object>} Pipeline result with guardrails metadata attached
+   */
+  async run(input) {
+    const guardrails = {
+      preRetrieval: null,
+      retrieval: null,
+      postGeneration: null,
+    };
+
+    // 1. Pre-retrieval guard
+    let query = input.query;
+    if (this.preGuard) {
+      guardrails.preRetrieval = this.preGuard.check(query);
+      this.emit('preRetrieval', guardrails.preRetrieval);
+
+      if (!guardrails.preRetrieval.safe && this.config.strict) {
+        return {
+          success: false,
+          error: 'Pre-retrieval guard blocked query',
+          guardrails,
+        };
+      }
+      query = guardrails.preRetrieval.sanitizedQuery;
+    }
+
+    // 2. Run the underlying pipeline with (possibly sanitized) query
+    const result = await this.pipeline.run({
+      ...input,
+      query,
+      options: { ...input.options, citations: true },
+    });
+
+    // 3. Retrieval guard on results
+    if (this.retrievalGuard && result.results) {
+      guardrails.retrieval = this.retrievalGuard.filter(
+        result.results,
+        input.context,
+      );
+      this.emit('retrieval', guardrails.retrieval);
+      result.results = guardrails.retrieval.results;
+    }
+
+    // 4. Post-generation guard on answer
+    if (this.postGuard && result.answer) {
+      guardrails.postGeneration = this.postGuard.check(
+        result.answer,
+        result.results,
+      );
+      this.emit('postGeneration', guardrails.postGeneration);
+
+      if (!guardrails.postGeneration.safe && this.config.strict) {
+        return {
+          ...result,
+          success: false,
+          error: 'Post-generation guard flagged output',
+          guardrails,
+        };
+      }
+      result.answer = guardrails.postGeneration.sanitizedOutput;
+    }
+
+    result.guardrails = guardrails;
+    return result;
+  }
+}
+
+module.exports = { GuardrailsPipeline, DEFAULT_CONFIG };

package/src/guardrails/post-generation-guard.js

@@ -0,0 +1,170 @@
+'use strict';
+
+const { EventEmitter } = require('events');
+
+/**
+ * Default configuration for the PostGenerationGuard.
+ * @type {object}
+ */
+const DEFAULT_CONFIG = {
+  enablePIIDetection: true,
+  enableGroundednessCheck: true,
+  groundednessThreshold: 0.3,
+  maxResponseLength: 50000,
+  minResponseLength: 0,
+  piiPatterns: {
+    email: /[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}/g,
+    phone: /\b\d{3}[-.]?\d{3}[-.]?\d{4}\b/g,
+    ssn: /\b\d{3}-\d{2}-\d{4}\b/g,
+    creditCard: /\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b/g,
+  },
+};
+
+/**
+ * Post-generation guard that checks LLM output for PII leakage,
+ * groundedness against source documents, and length constraints.
+ * Sanitizes output by redacting detected PII.
+ *
+ * @extends EventEmitter
+ */
+class PostGenerationGuard extends EventEmitter {
+  /**
+   * @param {object} [options] - Override default configuration
+   * @param {boolean} [options.enablePIIDetection=true] - Whether to detect PII in output
+   * @param {boolean} [options.enableGroundednessCheck=true] - Whether to check groundedness
+   * @param {number} [options.groundednessThreshold=0.3] - Minimum groundedness score
+   * @param {number} [options.maxResponseLength=50000] - Maximum allowed response length
+   * @param {number} [options.minResponseLength=0] - Minimum response length
+   * @param {object} [options.piiPatterns] - Custom PII regex patterns (merged with defaults)
+   */
+  constructor(options = {}) {
+    super();
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    if (options.piiPatterns) {
+      this.config.piiPatterns = {
+        ...DEFAULT_CONFIG.piiPatterns,
+        ...options.piiPatterns,
+      };
+    }
+  }
+
+  /**
+   * Check generated output for safety issues, PII, and groundedness.
+   *
+   * @param {string} output - The LLM-generated output text
+   * @param {Array<object>} [retrievedDocs=[]] - Source documents for groundedness checking
+   * @returns {{
+   *   safe: boolean,
+   *   issues: string[],
+   *   sanitizedOutput: string,
+   *   groundednessScore?: number
+   * }}
+   * @fires PostGenerationGuard#checked
+   */
+  check(output, retrievedDocs = []) {
+    const issues = [];
+
+    // 1. Validate output is a string
+    if (output == null || typeof output !== 'string') {
+      issues.push('Output must be a string');
+      const result = { safe: false, issues, sanitizedOutput: '' };
+      this.emit('checked', result);
+      return result;
+    }
+
+    // 2. Check length constraints
+    if (
+      this.config.maxResponseLength > 0 &&
+      output.length > this.config.maxResponseLength
+    ) {
+      issues.push(
+        `Output exceeds maximum length of ${this.config.maxResponseLength} characters`,
+      );
+    }
+    if (
+      this.config.minResponseLength > 0 &&
+      output.length < this.config.minResponseLength
+    ) {
+      issues.push(
+        `Output is shorter than minimum length of ${this.config.minResponseLength} characters`,
+      );
+    }
+
+    // 3. PII detection
+    const piiMatches = {};
+    if (this.config.enablePIIDetection) {
+      for (const [name, pattern] of Object.entries(this.config.piiPatterns)) {
+        // Reset lastIndex for global regex patterns
+        pattern.lastIndex = 0;
+        const matches = output.match(pattern);
+        if (matches && matches.length > 0) {
+          piiMatches[name] = matches;
+          issues.push(
+            `PII detected (${name}): ${matches.length} occurrence(s)`,
+          );
+        }
+      }
+    }
+
+    // 4. Groundedness check
+    let groundednessScore;
+    if (
+      this.config.enableGroundednessCheck &&
+      Array.isArray(retrievedDocs) &&
+      retrievedDocs.length > 0
+    ) {
+      try {
+        // Lazy-require CitationTracker to avoid circular dependencies
+        const { CitationTracker } = require('../citation/citation-tracker');
+        const tracker = new CitationTracker();
+        const tracking = tracker.track(output, retrievedDocs);
+        groundednessScore = tracking.groundednessScore;
+
+        if (groundednessScore < this.config.groundednessThreshold) {
+          issues.push(
+            `Groundedness score (${groundednessScore.toFixed(2)}) is below threshold (${this.config.groundednessThreshold})`,
+          );
+        }
+      } catch (err) {
+        // If citation tracking fails, log as issue but don't block
+        issues.push(`Groundedness check failed: ${err.message}`);
+      }
+    }
+
+    // 5. Sanitize: replace PII matches with [REDACTED]
+    let sanitizedOutput = output;
+    if (Object.keys(piiMatches).length > 0) {
+      for (const [, matches] of Object.entries(piiMatches)) {
+        for (const match of matches) {
+          sanitizedOutput = sanitizedOutput.split(match).join('[REDACTED]');
+        }
+      }
+    }
+
+    // Truncate if over max length
+    if (
+      this.config.maxResponseLength > 0 &&
+      sanitizedOutput.length > this.config.maxResponseLength
+    ) {
+      sanitizedOutput = sanitizedOutput.substring(
+        0,
+        this.config.maxResponseLength,
+      );
+    }
+
+    const result = {
+      safe: issues.length === 0,
+      issues,
+      sanitizedOutput,
+    };
+
+    if (groundednessScore !== undefined) {
+      result.groundednessScore = groundednessScore;
+    }
+
+    this.emit('checked', result);
+    return result;
+  }
+}
+
+module.exports = { PostGenerationGuard, DEFAULT_CONFIG };

package/src/guardrails/pre-retrieval-guard.js

@@ -0,0 +1,161 @@
+'use strict';
+
+const { EventEmitter } = require('events');
+const { validateInput } = require('../security/validator');
+const { tokenize } = require('../evaluate/scoring');
+
+/**
+ * Default configuration for the PreRetrievalGuard.
+ * @type {object}
+ */
+const DEFAULT_CONFIG = {
+  maxQueryLength: 10000,
+  enableInjectionDetection: true,
+  enableTopicFiltering: false,
+  blockedTopics: [],
+  allowedTopics: [],
+  injectionPatterns: [
+    /ignore\s+(all\s+)?previous\s+instructions/i,
+    /you\s+are\s+now/i,
+    /forget\s+everything/i,
+    /act\s+as\s+(a|an)?\s*/i,
+    /pretend\s+(you\s+are|to\s+be)/i,
+    /system\s*:/i,
+    /\bDAN\b/,
+    /bypass\s+(all\s+)?restrictions/i,
+    /override\s+(safety|content|system)/i,
+    /jailbreak/i,
+  ],
+};
+
+/**
+ * Pre-retrieval guard that validates and sanitizes user queries before
+ * they reach the retrieval layer. Checks for injection attacks, topic
+ * violations, and input length constraints.
+ *
+ * @extends EventEmitter
+ */
+class PreRetrievalGuard extends EventEmitter {
+  /**
+   * @param {object} [options] - Override default configuration
+   * @param {number} [options.maxQueryLength=10000] - Maximum allowed query length
+   * @param {boolean} [options.enableInjectionDetection=true] - Whether to check for prompt injection
+   * @param {boolean} [options.enableTopicFiltering=false] - Whether to enforce topic constraints
+   * @param {string[]} [options.blockedTopics=[]] - Topics to reject
+   * @param {string[]} [options.allowedTopics=[]] - If non-empty, only these topics are permitted
+   * @param {RegExp[]} [options.injectionPatterns] - Additional injection patterns (merged with defaults)
+   */
+  constructor(options = {}) {
+    super();
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    if (options.injectionPatterns) {
+      this.config.injectionPatterns = [
+        ...DEFAULT_CONFIG.injectionPatterns,
+        ...options.injectionPatterns,
+      ];
+    }
+  }
+
+  /**
+   * Check a query for safety issues, injection attempts, and topic violations.
+   *
+   * @param {string} query - The user query to validate
+   * @returns {{
+   *   safe: boolean,
+   *   issues: string[],
+   *   sanitizedQuery: string
+   * }}
+   * @fires PreRetrievalGuard#checked
+   */
+  check(query) {
+    const issues = [];
+
+    // 1. Validate query is a non-empty string
+    if (
+      query == null ||
+      typeof query !== 'string' ||
+      query.trim().length === 0
+    ) {
+      issues.push('Query must be a non-empty string');
+      const result = { safe: false, issues, sanitizedQuery: '' };
+      this.emit('checked', result);
+      return result;
+    }
+
+    // 2. Run basic input validation (length + disallowed tokens)
+    const validation = validateInput(query, {
+      maxLength: this.config.maxQueryLength,
+    });
+    if (!validation.valid) {
+      for (const reason of validation.reasons) {
+        issues.push(reason);
+      }
+    }
+
+    let sanitizedQuery = validation.sanitized;
+
+    // 3. Injection detection
+    const matchedPatterns = [];
+    if (this.config.enableInjectionDetection) {
+      for (const pattern of this.config.injectionPatterns) {
+        if (pattern.test(sanitizedQuery)) {
+          matchedPatterns.push(pattern);
+          issues.push(`Potential prompt injection detected: ${pattern}`);
+        }
+      }
+    }
+
+    // 4. Topic filtering
+    if (this.config.enableTopicFiltering) {
+      const tokens = tokenize(sanitizedQuery);
+
+      // Check blocked topics
+      if (this.config.blockedTopics.length > 0) {
+        for (const topic of this.config.blockedTopics) {
+          const topicLower = topic.toLowerCase();
+          if (tokens.includes(topicLower)) {
+            issues.push(`Query contains blocked topic: ${topic}`);
+          }
+        }
+      }
+
+      // Check allowed topics (whitelist mode)
+      if (this.config.allowedTopics.length > 0) {
+        const allowedLower = this.config.allowedTopics.map((t) =>
+          t.toLowerCase(),
+        );
+        const hasAllowedTopic = tokens.some((token) =>
+          allowedLower.includes(token),
+        );
+        if (!hasAllowedTopic) {
+          issues.push('Query does not match any allowed topics');
+        }
+      }
+    }
+
+    // 5. Sanitize: strip matched injection patterns from query
+    if (matchedPatterns.length > 0) {
+      for (const pattern of matchedPatterns) {
+        // Create a global version of the pattern for replacement
+        const flags = pattern.flags.includes('g')
+          ? pattern.flags
+          : pattern.flags + 'g';
+        const globalPattern = new RegExp(pattern.source, flags);
+        sanitizedQuery = sanitizedQuery.replace(globalPattern, '').trim();
+      }
+      // Collapse multiple spaces left after stripping
+      sanitizedQuery = sanitizedQuery.replace(/\s{2,}/g, ' ').trim();
+    }
+
+    const result = {
+      safe: issues.length === 0,
+      issues,
+      sanitizedQuery,
+    };
+
+    this.emit('checked', result);
+    return result;
+  }
+}
+
+module.exports = { PreRetrievalGuard, DEFAULT_CONFIG };

package/src/guardrails/retrieval-guard.js

@@ -0,0 +1,126 @@
+'use strict';
+
+const { EventEmitter } = require('events');
+
+/**
+ * Default configuration for the RetrievalGuard.
+ * @type {object}
+ */
+const DEFAULT_CONFIG = {
+  minRelevanceScore: 0,
+  maxAgeDays: 0, // 0 = disabled
+  enforceACL: false,
+};
+
+/**
+ * Retrieval guard that filters search results based on relevance score,
+ * document freshness, and access control lists. Applied between the
+ * retrieval and generation stages of the pipeline.
+ *
+ * @extends EventEmitter
+ */
+class RetrievalGuard extends EventEmitter {
+  /**
+   * @param {object} [options] - Override default configuration
+   * @param {number} [options.minRelevanceScore=0] - Minimum relevance score threshold
+   * @param {number} [options.maxAgeDays=0] - Maximum document age in days (0 = disabled)
+   * @param {boolean} [options.enforceACL=false] - Whether to enforce access control
+   */
+  constructor(options = {}) {
+    super();
+    this.config = { ...DEFAULT_CONFIG, ...options };
+  }
+
+  /**
+   * Filter retrieval results based on configured guards.
+   *
+   * @param {Array<object>} results - Array of retrieval results, each with .score and optional .metadata
+   * @param {object} [context={}] - Request context
+   * @param {string[]} [context.userPermissions] - User's permission roles for ACL filtering
+   * @returns {{
+   *   results: Array<object>,
+   *   removed: Array<object>,
+   *   metadata: { reasons: Object<string, number> }
+   * }}
+   * @fires RetrievalGuard#filtered
+   */
+  filter(results, context = {}) {
+    if (!Array.isArray(results)) {
+      const output = { results: [], removed: [], metadata: { reasons: {} } };
+      this.emit('filtered', output);
+      return output;
+    }
+
+    const kept = [];
+    const removed = [];
+    const reasons = {};
+
+    const now = Date.now();
+    const maxAgeMs =
+      this.config.maxAgeDays > 0
+        ? this.config.maxAgeDays * 24 * 60 * 60 * 1000
+        : 0;
+
+    for (const result of results) {
+      let dominated = false;
+      const doc = result.metadata || result.doc || result;
+
+      // 1. Relevance score filter
+      if (this.config.minRelevanceScore > 0) {
+        const score = result.score != null ? result.score : 0;
+        if (score < this.config.minRelevanceScore) {
+          dominated = true;
+          reasons.relevance = (reasons.relevance || 0) + 1;
+        }
+      }
+
+      // 2. Freshness filter
+      if (!dominated && maxAgeMs > 0) {
+        const timestamp = doc.updatedAt || doc.timestamp;
+        if (timestamp) {
+          const docTime = new Date(timestamp).getTime();
+          if (!isNaN(docTime) && now - docTime > maxAgeMs) {
+            dominated = true;
+            reasons.freshness = (reasons.freshness || 0) + 1;
+          }
+        }
+      }
+
+      // 3. ACL filter
+      if (!dominated && this.config.enforceACL && context.userPermissions) {
+        const requiredRoles = doc.requiredRole
+          ? [].concat(doc.requiredRole)
+          : doc.acl
+            ? [].concat(doc.acl)
+            : null;
+
+        if (requiredRoles && requiredRoles.length > 0) {
+          const hasAccess = requiredRoles.some((role) =>
+            context.userPermissions.includes(role),
+          );
+          if (!hasAccess) {
+            dominated = true;
+            reasons.acl = (reasons.acl || 0) + 1;
+          }
+        }
+      }
+
+      if (dominated) {
+        removed.push(result);
+      } else {
+        kept.push(result);
+      }
+    }
+
+    const output = {
+      results: kept,
+      removed,
+      metadata: { reasons },
+    };
+
+    this.emit('filtered', { original: results, filtered: kept, removed });
+    return output;
+  }
+}
+
+module.exports = { RetrievalGuard, DEFAULT_CONFIG };