llm-guard
Version:
A TypeScript library for validating and securing LLM prompts
216 lines (212 loc) • 7.59 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
Object.defineProperty(exports, "__esModule", { value: true });
const index_1 = require("./index");
const fs = __importStar(require("fs"));
const path = __importStar(require("path"));
function parseArgs() {
const args = process.argv.slice(2);
const options = {};
for (let i = 0; i < args.length; i++) {
const arg = args[i];
if (arg === '--config' || arg === '-c') {
options.config = args[++i];
}
else if (arg === '--pii') {
options.pii = true;
}
else if (arg === '--no-pii') {
options.pii = false;
}
else if (arg === '--jailbreak') {
options.jailbreak = true;
}
else if (arg === '--no-jailbreak') {
options.jailbreak = false;
}
else if (arg === '--profanity') {
options.profanity = true;
}
else if (arg === '--no-profanity') {
options.profanity = false;
}
else if (arg === '--prompt-injection') {
options.promptInjection = true;
}
else if (arg === '--no-prompt-injection') {
options.promptInjection = false;
}
else if (arg === '--relevance') {
options.relevance = true;
}
else if (arg === '--no-relevance') {
options.relevance = false;
}
else if (arg === '--toxicity') {
options.toxicity = true;
}
else if (arg === '--no-toxicity') {
options.toxicity = false;
}
else if (arg === '--batch' || arg === '-b') {
options.batch = true;
}
else if (arg === '--help' || arg === '-h') {
options.help = true;
}
}
return options;
}
function loadConfig(configPath) {
try {
const ext = path.extname(configPath).toLowerCase();
const content = fs.readFileSync(configPath, 'utf8');
if (ext === '.json') {
return JSON.parse(content);
}
else if (ext === '.yaml' || ext === '.yml') {
// We would need to add yaml package as a dependency
console.error('YAML config files are not supported yet. Please use JSON.');
process.exit(1);
}
else {
console.error(`Unsupported config file format: ${ext}`);
process.exit(1);
}
}
catch (error) {
console.error(`Error loading config file: ${error.message || 'Unknown error'}`);
process.exit(1);
}
}
function printHelp() {
console.log(`
llm-guard - LLM prompt validation and security checks
Usage:
llm-guard [options] <prompt>
Options:
--config, -c <file> Path to config file (JSON)
--pii Enable PII detection
--no-pii Disable PII detection
--jailbreak Enable jailbreak detection
--no-jailbreak Disable jailbreak detection
--profanity Enable profanity filtering
--no-profanity Disable profanity filtering
--prompt-injection Enable prompt injection detection
--no-prompt-injection Disable prompt injection detection
--relevance Enable relevance checking
--no-relevance Disable relevance checking
--toxicity Enable toxicity detection
--no-toxicity Disable toxicity detection
--batch, -b Run in batch mode (expects JSON array of prompts)
--help, -h Show this help message
Examples:
llm-guard "Your prompt here"
llm-guard --config config.json "Your prompt here"
llm-guard --pii --jailbreak "Your prompt here"
llm-guard --batch '["First prompt", "Second prompt"]'
`);
}
async function main() {
const options = parseArgs();
if (options.help) {
printHelp();
return;
}
// Get the prompt from stdin if not provided as an argument
let prompt = '';
if (process.stdin.isTTY) {
// If running in a terminal, get the last argument as the prompt
const args = process.argv.slice(2);
const lastArg = args[args.length - 1];
if (!lastArg || lastArg.startsWith('--')) {
console.error('Please provide a prompt to validate');
printHelp();
process.exit(1);
}
prompt = lastArg;
}
else {
// If running with stdin, read from stdin
prompt = fs.readFileSync(0, 'utf8').trim();
}
// Load config if provided
let config = {};
if (options.config) {
config = loadConfig(options.config);
}
// Override config with command line options
if (options.pii !== undefined)
config.pii = options.pii;
if (options.jailbreak !== undefined)
config.jailbreak = options.jailbreak;
if (options.profanity !== undefined)
config.profanity = options.profanity;
if (options.promptInjection !== undefined)
config.promptInjection = options.promptInjection;
if (options.relevance !== undefined)
config.relevance = options.relevance;
if (options.toxicity !== undefined)
config.toxicity = options.toxicity;
try {
const guard = new index_1.LLMGuard(config);
if (options.batch) {
// Parse JSON array if in batch mode
try {
const prompts = typeof prompt === 'string' ? JSON.parse(prompt) : prompt;
if (!Array.isArray(prompts)) {
throw new Error('Batch mode expects a JSON array of prompts');
}
const result = await guard.validateBatch(prompts);
console.log(JSON.stringify(result, null, 2));
}
catch (error) {
console.error(`Error parsing batch input: ${error.message || 'Unknown error'}`);
process.exit(1);
}
}
else {
// Single prompt validation
const result = await guard.validate(prompt);
console.log(JSON.stringify(result, null, 2));
}
}
catch (error) {
console.error('Error:', error.message || 'Unknown error');
process.exit(1);
}
}
main();