From 693237954d7a35b1e557b5f4f5e88168c7dbae9f Mon Sep 17 00:00:00 2001
From: Tomas Dvorak <toomas2d@gmail.com>
Date: Thu, 3 Oct 2024 19:28:38 +0200
Subject: [PATCH] feat(agent): update line prefix parser fallback and add tests

Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>
---
 src/agents/parsers/linePrefix.test.ts | 65 +++++++++++++++++++++++++++
 src/agents/parsers/linePrefix.ts      |  7 +--
 2 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/src/agents/parsers/linePrefix.test.ts b/src/agents/parsers/linePrefix.test.ts
index af591a67..0f22ed73 100644
--- a/src/agents/parsers/linePrefix.test.ts
+++ b/src/agents/parsers/linePrefix.test.ts
@@ -167,6 +167,71 @@ But I can use GoogleSearch to find out.`,
     });
   });
 
+  describe("Fallback", () => {
+    const getParser = () =>
+      new LinePrefixParser(
+        {
+          thought: {
+            prefix: "Thought:",
+            field: new ZodParserField(z.string()),
+            isStart: true,
+            next: ["final_answer"],
+          },
+          final_answer: {
+            prefix: "Final Answer:",
+            field: new ZodParserField(z.string()),
+            isEnd: true,
+            next: [],
+          },
+        },
+        {
+          fallback: (value) =>
+            [
+              { key: "thought", value: "I now know the final answer." },
+              {
+                key: "final_answer",
+                value,
+              },
+            ] as const,
+        },
+      );
+
+    it("Process", async () => {
+      const parser = getParser();
+      await parser.add("2+2=4");
+      await expect(parser.end()).resolves.toMatchInlineSnapshot(`
+        {
+          "final_answer": "2+2=4",
+          "thought": "I now know the final answer.",
+        }
+      `);
+    });
+
+    it("Process #2", async () => {
+      const parser = getParser();
+      await parser.add("A\nB\nC");
+      await expect(parser.end()).resolves.toMatchInlineSnapshot(`
+        {
+          "final_answer": "A
+        B
+        C",
+          "thought": "I now know the final answer.",
+        }
+      `);
+    });
+
+    it("Process #3", async () => {
+      const parser = getParser();
+      await parser.add("Thought");
+      await expect(parser.end()).resolves.toMatchInlineSnapshot(`
+        {
+          "final_answer": "Thought",
+          "thought": "I now know the final answer.",
+        }
+      `);
+    });
+  });
+
   describe("Edge cases", () => {
     it("Prevents processing when line starts with a potential prefix", async () => {
       const parser = new LinePrefixParser({
diff --git a/src/agents/parsers/linePrefix.ts b/src/agents/parsers/linePrefix.ts
index 2041e9ee..191bd3f9 100644
--- a/src/agents/parsers/linePrefix.ts
+++ b/src/agents/parsers/linePrefix.ts
@@ -237,9 +237,10 @@ export class LinePrefixParser<
     }
 
     if (!this.lastNodeKey && this.options.fallback) {
-      const stash = linesToString(this.excludedLines);
-      this.lines.length = 0;
-      this.excludedLines.length = 0;
+      const stash = linesToString([
+        ...this.excludedLines.splice(0, Infinity),
+        ...this.lines.splice(0, Infinity),
+      ]);
 
       const nodes = this.options.fallback(stash);
       await this.add(