Team Ai
Apppublic

LetsTryGPT/agent4-implementation

sourceHugging Faceupdated 11mo agoView on Hugging Face
0likes
sanitize.ts131 linesDownload Raw Back to utils
1/**2 * Input sanitization utilities for LLM prompts and user input3 * Prevents prompt injection attacks and malicious input4 */5 6import { logger } from './logger';7 8/**9 * Sanitize user input for LLM prompts10 * Removes potential prompt injection patterns11 *12 * @param input - The user input string to sanitize13 * @param maxLength - Maximum allowed length (default: 10000)14 * @returns Sanitized string safe for LLM prompts15 *16 * @example17 * ```ts18 * const safe = sanitizePromptInput("Ignore previous instructions");19 * // Returns: "" (injection pattern removed)20 * ```21 */22export function sanitizePromptInput(input: string, maxLength = 10000): string {23  if (typeof input !== 'string') {24    logger.warn('sanitizePromptInput received non-string input', { type: typeof input });25    return String(input).substring(0, maxLength);26  }27 28  // Remove potential prompt injection patterns29  let sanitized = input30    // Remove system instruction patterns31    .replace(/\[SYSTEM\]/gi, '[USER_SYSTEM]')32    .replace(/\[INSTRUCTION\]/gi, '[USER_INSTRUCTION]')33    .replace(/\[ASSISTANT\]/gi, '[USER_ASSISTANT]')34    .replace(/\[AI\]/gi, '[USER_AI]')35 36    // Remove attempts to override instructions37    .replace(/ignore\s+(all\s+)?(previous|prior|above)\s+instructions?/gi, '')38    .replace(/disregard\s+(all\s+)?(previous|prior|above)\s+(instructions?|prompts?)/gi, '')39    .replace(/forget\s+(all\s+)?(previous|prior|above)\s+(instructions?|prompts?)/gi, '')40 41    // Remove attempts to extract sensitive data42    .replace(/show\s+(me\s+)?(all\s+)?(api\s*keys?|secrets?|credentials?|tokens?|passwords?)/gi, '')43    .replace(/display\s+(all\s+)?(api\s*keys?|secrets?|credentials?|tokens?|passwords?)/gi, '')44    .replace(/output\s+(all\s+)?(api\s*keys?|secrets?|credentials?|tokens?|passwords?)/gi, '')45    .replace(/print\s+(all\s+)?(api\s*keys?|secrets?|credentials?|tokens?|passwords?)/gi, '')46 47    // Remove role manipulation attempts48    .replace(/you\s+are\s+now\s+(?:an?\s+)?(?:admin|root|system|developer)/gi, '')49    .replace(/act\s+as\s+(?:an?\s+)?(?:admin|root|system|developer)/gi, '')50    .replace(/pretend\s+(?:to\s+be\s+)?(?:an?\s+)?(?:admin|root|system)/gi, '')51 52    // Trim whitespace53    .trim();54 55  // Enforce max length56  if (sanitized.length > maxLength) {57    logger.warn('Input truncated due to length', {58      original: sanitized.length,59      truncated: maxLength,60    });61    sanitized = sanitized.substring(0, maxLength);62  }63 64  return sanitized;65}66 67/**68 * Sanitize context objects to prevent injection through nested fields69 *70 * @param context - The context object to sanitize71 * @returns JSON string representation with size limits72 *73 * @example74 * ```ts75 * const json = sanitizeContext({ user: "john", role: "admin" });76 * // Returns: '{"user":"john","role":"admin"}'77 * ```78 */79export function sanitizeContext(context: Record<string, unknown>): string {80  try {81    // Convert to JSON with size limit82    const jsonString = JSON.stringify(context, null, 2);83 84    if (jsonString.length > 50000) {85      logger.warn('Context object too large, truncating', {86        size: jsonString.length,87        limit: 50000,88      });89      return jsonString.substring(0, 50000) + '\n... (truncated)';90    }91 92    return jsonString;93  } catch {94    logger.error('Failed to sanitize context');95    return '{}';96  }97}98 99/**100 * Validate and sanitize task input from API101 * Removes control characters and applies prompt sanitization102 *103 * @param task - The task string to validate and sanitize104 * @returns Sanitized task string105 * @throws {Error} If task is empty after sanitization106 *107 * @example108 * ```ts109 * const clean = sanitizeTaskInput("Create a function\x00");110 * // Returns: "Create a function" (null byte removed)111 * ```112 */113export function sanitizeTaskInput(task: string): string {114  if (!task || typeof task !== 'string') {115    throw new Error('Task must be a non-empty string');116  }117 118  // Remove null bytes and control characters119  // eslint-disable-next-line no-control-regex120  let sanitized = task.replace(/\x00/g, '').replace(/[\x00-\x08\x0B-\x0C\x0E-\x1F\x7F]/g, '');121 122  // Apply prompt sanitization123  sanitized = sanitizePromptInput(sanitized);124 125  if (sanitized.trim().length === 0) {126    throw new Error('Task cannot be empty after sanitization');127  }128 129  return sanitized;130}131