llm_guardrail 1.0.0 → 2.0.0

package/index.js

@@ -1,42 +1,62 @@
-import { execFile } from "child_process";
-import path from "path";
-import { fileURLToPath } from "url";
-
-const __filename = fileURLToPath(import.meta.url);
-const __dirname = path.dirname(__filename);
-
-const PYTHON = "D:/npm/llm_Guardrails/venv/Scripts/python.exe";
-
-
-export function check(prompt) {
-  return new Promise((resolve, reject) => {
-    if (typeof prompt !== "string") {
-      return reject(new Error("Prompt must be a string"));
-    }
-
-    const scriptPath = path.join(__dirname, "model", "predict.py");
-    
-
-    execFile(
-      PYTHON,
-      [scriptPath, prompt],
-      {},
-      (err, stdout, stderr) => {
-        if (err) return reject(err);
-        if (stderr) return reject(new Error(stderr));
-
-        try {
-          const result = JSON.parse(stdout);
-          resolve({
-            allowed: result.injective === 0,
-            ...result
-          });
-        } catch (e) {
-          reject(new Error("Invalid JSON from Python"));
-        }
-      }
-    );
-  });
-}
-
-export default { check };
+import fs from 'fs';
+import path from 'path';
+import { fileURLToPath } from 'url';
+import { TfidfVectorizer } from './model/tfidf.js';
+import { LogisticRegression } from './model/logistic_regression.js';
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+
+// Load model once
+let vectorizer = null;
+let model = null;
+
+function loadModel() {
+  if (model && vectorizer) return { model, vectorizer };
+  
+  try {
+    const modelPath = path.join(__dirname, 'model', 'model_data.json');
+    const modelData = JSON.parse(fs.readFileSync(modelPath, 'utf8'));
+    
+    vectorizer = new TfidfVectorizer(modelData);
+    model = new LogisticRegression(modelData);
+    
+    return { model, vectorizer };
+  } catch (error) {
+    throw new Error(`Failed to load model: ${error.message}`);
+  }
+}
+
+export function check(prompt) {
+  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 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,
+        prediction: prediction,
+        confidence: positiveProb,
+        probabilities: {
+          safe: probabilities[0],
+          injection: probabilities[1]
+        }
+      });
+    } catch (error) {
+      reject(error);
+    }
+  });
+}
+
+export default { check };

package/model/logistic_regression.js

@@ -0,0 +1,42 @@
+export class LogisticRegression {
+  constructor(modelData) {
+    this.coefficients = modelData.coefficients[0]; // For binary classification
+    this.intercept = modelData.intercept[0];
+    this.classes = modelData.classes;
+    this.threshold = modelData.threshold || 0.5;
+  }
+
+  /**
+   * Sigmoid function
+   */
+  sigmoid(z) {
+    return 1 / (1 + Math.exp(-z));
+  }
+
+  /**
+   * Predict probability for a feature vector
+   */
+  predictProba(features) {
+    // Calculate z = w·x + b
+    let z = this.intercept;
+    for (let i = 0; i < features.length; i++) {
+      z += this.coefficients[i] * features[i];
+    }
+    
+    // Apply sigmoid
+    const probPositive = this.sigmoid(z);
+    
+    return {
+      probabilities: [1 - probPositive, probPositive],
+      positiveProb: probPositive
+    };
+  }
+
+  /**
+   * Predict class (0 or 1)
+   */
+  predict(features) {
+    const { positiveProb } = this.predictProba(features);
+    return positiveProb >= this.threshold ? 1 : 0;
+  }
+}