Add sanitization to strip LLM control tokens

acusti · May 10, 2024 · 25173c4 · 25173c4
1 parent 20a0e98
commit 25173c4
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/packages/parsing/src/parse-as-json.ts b/packages/parsing/src/parse-as-json.ts
@@ -183,13 +183,16 @@ const getObjectKeyFromIndex = (index: number) =>
 
 const OBJECT_KEY_REGEXP = /^\s*"[^"]+":/;
 
+const CONTROL_TOKENS_REGEXP = /(^<\|im_start\|>|<\|im_end\|>$)/;
+
 type ParsedValue = string | boolean | number | GenericObject | Array<unknown>;
 
 export function parseAsJSON(text: string): ParsedValue | null {
     // if the input is undefined/null, return null to indicate failure
     if (text == null) return null;
 
-    // attempt to parse the string as-is
+    text = text.replace(CONTROL_TOKENS_REGEXP, '');
+    // attempt to parse the string as-is (minus control tokens)
     try {
         return JSON.parse(text) as ParsedValue;
     } catch (error) {