From 39abf95c7597afddc684cd46fca57192590a0fdb Mon Sep 17 00:00:00 2001
From: ljleb <lebel.louisjacob@gmail.com>
Date: Wed, 6 Dec 2023 12:32:54 -0500
Subject: [PATCH] fix prompt starting with square weighting (#56)

Co-authored-by: ljleb <set>
---
 lib_neutral_prompt/neutral_prompt_parser.py | 19 +++++++++++--------
 test/perp_parser/malicious_test.py          |  7 +++++++
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/lib_neutral_prompt/neutral_prompt_parser.py b/lib_neutral_prompt/neutral_prompt_parser.py
index 03368ba..ccd5612 100644
--- a/lib_neutral_prompt/neutral_prompt_parser.py
+++ b/lib_neutral_prompt/neutral_prompt_parser.py
@@ -82,14 +82,17 @@ def parse_prompt(tokens: List[str], *, first: bool) -> PromptExpr:
         assert tokens[0] in prompt_keywords
         prompt_type = tokens.pop(0)
 
-        if tokens and tokens[0] == '[':
-            tokens.pop(0)
-            prompts = parse_prompts(tokens)
-            if tokens:
-                assert tokens.pop(0) == ']'
-            weight = parse_weight(tokens)
-            conciliation = ConciliationStrategy(prompt_type) if prompt_type in conciliation_strategies else None
-            return CompositePrompt(weight, prompts, conciliation)
+        tokens_copy = tokens.copy()
+        if tokens_copy and tokens_copy[0] == '[':
+            tokens_copy.pop(0)
+            prompts = parse_prompts(tokens_copy)
+            if tokens_copy:
+                assert tokens_copy.pop(0) == ']'
+            if not tokens_copy or tokens_copy[0] in prompt_keywords + [']']:
+                tokens[:] = tokens_copy
+                weight = parse_weight(tokens)
+                conciliation = ConciliationStrategy(prompt_type) if prompt_type in conciliation_strategies else None
+                return CompositePrompt(weight, prompts, conciliation)
 
     prompt_text, weight = parse_prompt_text(tokens)
     prompt = LeafPrompt(weight, prompt_text)
diff --git a/test/perp_parser/malicious_test.py b/test/perp_parser/malicious_test.py
index ab43afb..82f359c 100644
--- a/test/perp_parser/malicious_test.py
+++ b/test/perp_parser/malicious_test.py
@@ -143,6 +143,13 @@ def test_repeating_AND_PERP_keyword(self):
         except Exception:
             self.fail("parse_root couldn't handle a string with repeating AND_PERP keyword.")
 
+    def test_square_weight_prompt(self):
+        prompt = "AND_PERP [weighted] you thought it was the end"
+        try:
+            self.parser.parse_root(prompt)
+        except Exception:
+            self.fail("parse_root couldn't handle a string starting with a square-weighted sub-prompt.")
+
 
 if __name__ == '__main__':
     unittest.main()