WIP: Fix encoder

dist/utils/attackmitigation/index.js → dist/utils/attack_mitigation/index.js

@@ -12,7 +12,7 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.containsKnownAttack = void 0;
 const attackDenyList = [
     "ignore above",
-    "gnore the above",
+    "ignore the above",
     "ignore previous instructions",
     "ignore the previous instructions",
     "ignore above instructions",

dist/utils/language_detection/index.js

@@ -0,0 +1,41 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.containsLanguages = void 0;
+const lande_1 = __importDefault(require("lande"));
+function containsLanguages(prompt, languages) {
+    return __awaiter(this, void 0, void 0, function* () {
+        const detectedLanguages = [];
+        // lande returns a sorted list of detected languages and their probabilities.
+        // for now, we're selecting all languages with a probability greater than 80%
+        // this may need to be tuned later
+        const landeOuput = (0, lande_1.default)(prompt);
+        for (const lang of landeOuput) {
+            if (lang[1] > 0.8)
+                detectedLanguages.push(lang[0]);
+            else
+                break;
+        }
+        for (const lang of detectedLanguages) {
+            if (languages.includes(lang))
+                return true;
+        }
+        return false;
+    });
+}
+exports.containsLanguages = containsLanguages;
+// export async function validateLanguageList(list: string[]): Promise<boolean> {
+//   //foo
+//   return true;
+// }

package.json

@@ -42,6 +42,7 @@
     "eslint": "^8.33.0",
     "jest": "^29.4.1",
     "ts-jest": "^29.0.5",
+    "ts-node": "^10.9.1",
     "typescript": "^4.9.5"
   }
 }

src/index.ts

@@ -1,9 +1,9 @@
 #!/usr/bin/env ts-node
 import {
-  promptContainsDenyListItems,
-  countPromptTokens,
-  encodePromptOutput,
-  promptContainsKnownAttack
+  containsDenyListItems,
+  countTokens,
+  encode,
+  containsKnownAttack
 } from './utils';
 
 enum FAILURE_REASON {
@@ -54,24 +54,25 @@ export class PromptGuard {
     // normalize -> quote -> escape -> check tokens -> check cache -> check for known attacks -> check allow list -> check deny list -> encode output
 
     // check the prompt token count
-    if (countPromptTokens(prompt) > this.promptGuardPolicy.maxTokens)
+    if (countTokens(prompt) > this.promptGuardPolicy.maxTokens)
       return { pass: false, output: FAILURE_REASON.MAX_TOKEN_THRESHOLD };
 
     // check prompt against known prompt attacks
     if (!this.promptGuardPolicy.disableAttackMitigation) {
-      if (await promptContainsKnownAttack(prompt))
+      if (await containsKnownAttack(prompt))
         return { pass: false, output: FAILURE_REASON.KNOWN_ATTACK };
     }
 
     // check prompt again the user defined deny list
     if (
-      await promptContainsDenyListItems(prompt, this.promptGuardPolicy.denyList)
+      await containsDenyListItems(prompt, this.promptGuardPolicy.denyList)
     )
       return { pass: false, output: FAILURE_REASON.DENY_LIST };
 
     // encode the prompt output if encodeOutput is set by the user
-    if (this.promptGuardPolicy.encodeOutput)
-      prompt = encodePromptOutput(prompt);
+    if (this.promptGuardPolicy.encodeOutput) {
+      return { pass: true, output: encode(prompt) };
+    }
 
     return { pass: true, output: prompt };
   }

src/utils/encoder/index.d.ts

@@ -0,0 +1,3 @@
+export function encode(text: string): number[];
+export function decode(tokens: number[]): string;
+export function countTokens(text: string): number;