feat(agent): add fallback for invalid llm output

Ref: #55
i-am-bee · Oct 3, 2024 · 5ac89c1 · 5ac89c1
1 parent f7d5128
commit 5ac89c1
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 42 deletions.
diff --git a/src/agents/bee/runner.ts b/src/agents/bee/runner.ts
@@ -161,47 +161,59 @@ export class BeeAgentRunner {
     const parserRegex =
       /Thought:.+\n(?:Final Answer:[\S\s]+|Function Name:.+\nFunction Input:\{.+\}\nFunction Caption:.+\nFunction Output:)?/;
 
-    const parser = new LinePrefixParser({
-      thought: {
-        prefix: "Thought:",
-        next: ["tool_name", "final_answer"],
-        isStart: true,
-        field: new ZodParserField(z.string().min(1)),
-      },
-      tool_name: {
-        prefix: "Function Name:",
-        next: ["tool_input"],
-        field: new ZodParserField(z.enum(tools.map((tool) => tool.name) as [string, ...string[]])),
-      },
-      tool_input: {
-        prefix: "Function Input:",
-        next: ["tool_caption", "tool_output"],
-        isEnd: true,
-        field: new JSONParserField({
-          schema: z.object({}).passthrough(),
-          base: {},
-        }),
-      },
-      tool_caption: {
-        prefix: "Function Caption:",
-        next: ["tool_output"],
-        isEnd: true,
-        field: new ZodParserField(z.string()),
-      },
-      tool_output: {
-        prefix: "Function Output:",
-        next: ["final_answer"],
-        isEnd: true,
-        field: new ZodParserField(z.string()),
-      },
-      final_answer: {
-        prefix: "Final Answer:",
-        next: [],
-        isStart: true,
-        isEnd: true,
-        field: new ZodParserField(z.string().min(1)),
+    const parser = new LinePrefixParser(
+      {
+        thought: {
+          prefix: "Thought:",
+          next: ["tool_name", "final_answer"],
+          isStart: true,
+          field: new ZodParserField(z.string().min(1)),
+        },
+        tool_name: {
+          prefix: "Function Name:",
+          next: ["tool_input"],
+          field: new ZodParserField(
+            z.enum(tools.map((tool) => tool.name) as [string, ...string[]]),
+          ),
+        },
+        tool_input: {
+          prefix: "Function Input:",
+          next: ["tool_caption", "tool_output"],
+          isEnd: true,
+          field: new JSONParserField({
+            schema: z.object({}).passthrough(),
+            base: {},
+            matchPair: ["{", "}"],
+          }),
+        },
+        tool_caption: {
+          prefix: "Function Caption:",
+          next: ["tool_output"],
+          isEnd: true,
+          field: new ZodParserField(z.string()),
+        },
+        tool_output: {
+          prefix: "Function Output:",
+          next: ["final_answer"],
+          isEnd: true,
+          field: new ZodParserField(z.string()),
+        },
+        final_answer: {
+          prefix: "Final Answer:",
+          next: [],
+          isStart: true,
+          isEnd: true,
+          field: new ZodParserField(z.string().min(1)),
+        },
+      } as const,
+      {
+        fallback: (stash) =>
+          [
+            { key: "thought", value: "I now know the final answer." },
+            { key: "final_answer", value: stash },
+          ] as const,
       },
-    } as const);
+    );
 
     return {
       parser,

diff --git a/src/agents/parsers/linePrefix.ts b/src/agents/parsers/linePrefix.ts
@@ -100,7 +100,12 @@ export class LinePrefixParser<
     return this.done;
   }
 
-  constructor(protected readonly nodes: T) {
+  constructor(
+    protected readonly nodes: T,
+    protected readonly options: {
+      fallback?: (value: string) => readonly { key: StringKey<T>; value: string }[];
+    } = {},
+  ) {
     super();
 
     let hasStartNode = false;
@@ -219,6 +224,8 @@ export class LinePrefixParser<
         context: {
           lines: linesToString(this.lines.concat(extra.line ? [extra.line] : [])),
           excludedLines: linesToString(this.excludedLines),
+          finalState: this.finalState,
+          partialState: this.partialState,
         },
       },
     );
@@ -228,6 +235,18 @@ export class LinePrefixParser<
     if (this.done) {
       return this.finalState;
     }
+
+    if (!this.lastNodeKey && this.options.fallback) {
+      const stash = linesToString(this.excludedLines);
+      this.lines.length = 0;
+      this.excludedLines.length = 0;
+
+      const nodes = this.options.fallback(stash);
+      await this.add(
+        nodes.map((node) => `${this.nodes[node.key].prefix}${node.value}`).join(NEW_LINE_CHARACTER),
+      );
+    }
+
     this.done = true;
 
     if (!this.lastNodeKey) {
@@ -285,7 +304,7 @@ export class LinePrefixParser<
     } catch (e) {
       if (e instanceof ZodError) {
         this.throwWithContext(
-          `Value for ${key} cannot be retrieved because it's value does not adhere to the appropriate schema.`,
+          `Value for '${key}' cannot be retrieved because it's value does not adhere to the appropriate schema.`,
           { errors: [e] },
         );
       }
@@ -332,6 +351,7 @@ export class LinePrefixParser<
       emitter: this.emitter,
       done: this.done,
       lastNodeKey: this.lastNodeKey,
+      options: this.options,
     };
   }