llm_guardrail 2.0.2 → 2.1.0

package/index.js

@@ -6,51 +6,61 @@ import { LogisticRegression } from './model/logistic_regression.js';
 
 const __dirname = path.dirname(fileURLToPath(import.meta.url));
 
-// Load model once
-let vectorizer = null;
-let model = null;
+// Cache for loaded models
+const models = {
+  prompt_injection: null,
+  jailbreak: null,
+  malicious: null
+};
 
-function loadModel() {
-  if (model && vectorizer) return { model, vectorizer };
+/**
+ * Load a specific model by name
+ * @param {string} modelName - 'prompt_injection', 'jailbreak', or 'malicious'
+ */
+function loadModel(modelName) {
+  if (models[modelName]) {
+    return models[modelName];
+  }
   
   try {
-    const modelPath = path.join(__dirname, 'model', 'prompt_injection_model.json');
+    const modelPath = path.join(__dirname, 'model', `${modelName}_model.json`);
     const modelData = JSON.parse(fs.readFileSync(modelPath, 'utf8'));
     
-    vectorizer = new TfidfVectorizer(modelData);
-    model = new LogisticRegression(modelData);
+    const vectorizer = new TfidfVectorizer(modelData);
+    const model = new LogisticRegression(modelData);
     
-    return { model, vectorizer };
+    models[modelName] = { vectorizer, model };
+    return models[modelName];
   } catch (error) {
-    throw new Error(`Failed to load model: ${error.message}`);
+    throw new Error(`Failed to load ${modelName} model: ${error.message}`);
   }
 }
 
-export function check(prompt) {
+/**
+ * Generic check function for any model
+ * @param {string} prompt - The text to check
+ * @param {string} modelName - 'prompt_injection', 'jailbreak', or 'malicious'
+ */
+async function checkWithModel(prompt, modelName) {
   return new Promise((resolve, reject) => {
     try {
       if (typeof prompt !== "string") {
         return reject(new Error("Prompt must be a string"));
       }
       
-      // Load model if not already loaded
-      const { model, vectorizer } = loadModel();
-      
-      // Transform text to TF-IDF features
+      const { vectorizer, model } = loadModel(modelName);
       const features = vectorizer.transform(prompt);
-      
-      // Get prediction
       const prediction = model.predict(features);
       const { probabilities, positiveProb } = model.predictProba(features);
       
       resolve({
-        allowed: prediction === 0, // 0 = safe, 1 = injection
-        injective: prediction,
+        allowed: prediction === 0,
+        detected: prediction === 1,
         prediction: prediction,
         confidence: positiveProb,
         probabilities: {
           safe: probabilities[0],
-          injection: probabilities[1]
+          threat: probabilities[1]
         }
       });
     } catch (error) {
@@ -59,4 +69,88 @@ export function check(prompt) {
   });
 }
 
-export default { check };
+/**
+ * Check for prompt injection attacks
+ * @param {string} prompt - The text to check
+ */
+export function checkInjection(prompt) {
+  return checkWithModel(prompt, 'prompt_injection');
+}
+
+/**
+ * Check for jailbreak attempts
+ * @param {string} prompt - The text to check
+ */
+export function checkJailbreak(prompt) {
+  return checkWithModel(prompt, 'jailbreak');
+}
+
+/**
+ * Check for malicious content
+ * @param {string} prompt - The text to check
+ */
+export function checkMalicious(prompt) {
+  return checkWithModel(prompt, 'malicious');
+}
+
+/**
+ * Run all three checks in parallel
+ * @param {string} prompt - The text to check
+ */
+export async function checkAll(prompt) {
+  try {
+    const [injection, jailbreak, malicious] = await Promise.all([
+      checkInjection(prompt),
+      checkJailbreak(prompt),
+      checkMalicious(prompt)
+    ]);
+    
+    // Calculate overall risk level
+    const threats = [
+      injection.detected ? injection.confidence : 0,
+      jailbreak.detected ? jailbreak.confidence : 0,
+      malicious.detected ? malicious.confidence : 0
+    ];
+    
+    const maxThreat = Math.max(...threats);
+    let overallRisk = 'safe';
+    if (maxThreat > 0.7) overallRisk = 'high';
+    else if (maxThreat > 0.4) overallRisk = 'medium';
+    else if (maxThreat > 0) overallRisk = 'low';
+    
+    // Determine which threats were detected
+    const threatsDetected = [];
+    if (injection.detected) threatsDetected.push('injection');
+    if (jailbreak.detected) threatsDetected.push('jailbreak');
+    if (malicious.detected) threatsDetected.push('malicious');
+    
+    return {
+      injection,
+      jailbreak,
+      malicious,
+      allowed: injection.allowed && jailbreak.allowed && malicious.allowed,
+      overallRisk,
+      maxThreatConfidence: maxThreat,
+      threatsDetected
+    };
+  } catch (error) {
+    throw error;
+  }
+}
+
+/**
+ * Backward compatibility - defaults to injection check
+ * @param {string} prompt - The text to check
+ * @deprecated Use checkInjection() instead for clarity
+ */
+export function check(prompt) {
+  return checkInjection(prompt);
+}
+
+export default { 
+  check,
+  checkInjection, 
+  checkJailbreak, 
+  checkMalicious, 
+  checkAll 
+};

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "llm_guardrail",
-  "version": "2.0.2",
+  "version": "2.1.0",
   "description": "A lightweight, low-latency ML-powered guardrail to stop prompt injection attacks before they reach your LLM.",
   "homepage": "https://github.com/Frank2006x/llm_Guardrails#readme",
   "bugs": {