fix backtracking into prompt

controllers/pyctrl/samples/test.py

@@ -3,6 +3,7 @@
 
 # asserts for microsoft/Orca-2-13b
 
+
 async def test_backtrack_one():
     await aici.FixedTokens("3+")
     l = aici.Label()
@@ -124,6 +125,15 @@ def inst(s: str) -> str:
     )
 
 
+async def test_prompt_backtrack():
+    await aici.FixedTokens("Some test prompt for the model to generate more text.")
+    l = aici.Label()
+    await aici.FixedTokens("And then some more text.")
+    await aici.gen_tokens(max_tokens=2)
+    await aici.FixedTokens("Now different text.", following=l)
+    await aici.gen_tokens(max_tokens=2)
+
+
 async def test_sample():
     # initialization code
     print("I'm going in the logs!")
@@ -161,6 +171,7 @@ async def test_eos():
     await aici.gen_tokens(regex=r' "[^"]+"', max_tokens=6, store_var="french")
     aici.check_vars({"french": ' "bonjour"'})
 
+
 async def test_joke():
     await aici.FixedTokens("Do you want a joke or a poem? A")
     answer = await aici.gen_text(options=[" joke", " poem"])

rllm/rllm-base/src/seq.rs

@@ -153,6 +153,8 @@ impl Sequence {
     ) {
         self.tokens.truncate(self.get_len() - backtrack);
         self.output_ptr = std::cmp::min(self.output_ptr, self.get_len());
+        // backtracking can remove some tokens from the initial prompt
+        self.prompt_len = std::cmp::min(self.prompt_len, self.get_len());
         if backtrack > 0 {
             self.output_pending.clear();
             self.output_pending.extend_from_slice(" ↩ ".as_bytes());