feat(agent): update line prefix parser fallback and add tests

Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>
i-am-bee · Oct 3, 2024 · 6932379 · 6932379
1 parent a6e582d
commit 6932379
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 3 deletions.
diff --git 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
@@ -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(