feat(agents): add retry for parser errors (i-am-bee#204)

Signed-off-by: Tomas Dvorak <[email protected]>
tonxxd · Dec 3, 2024 · f9de559 · f9de559
1 parent c176d05
commit f9de559
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 1 deletion.
diff --git a/src/agents/bee/prompts.ts b/src/agents/bee/prompts.ts
@@ -186,3 +186,9 @@ export const BeeToolNotFoundPrompt = new PromptTemplate({
 Use one of the following functions: {{#trim}}{{#tools}}{{name}},{{/tools}}{{/trim}}
 {{/tools.length}}`,
 });
+
+export const BeeSchemaErrorPrompt = new PromptTemplate({
+  schema: z.object({}).passthrough(),
+  template: `Error: The generated response does not adhere to the communication structure mentioned in the system prompt.
+You communicate only in instruction lines. Valid instruction lines are 'Thought' followed by either 'Function Name' + 'Function Input' + 'Function Output' or 'Final Answer'.`,
+});
diff --git a/src/agents/bee/runners/granite/prompts.ts b/src/agents/bee/runners/granite/prompts.ts
@@ -16,7 +16,7 @@
 
 import { PromptTemplate } from "@/template.js";
 import { z } from "zod";
-import { BeeAssistantPrompt } from "@/agents/bee/prompts.js";
+import { BeeAssistantPrompt, BeeSchemaErrorPrompt } from "@/agents/bee/prompts.js";
 
 export const GraniteBeeAssistantPrompt = BeeAssistantPrompt.fork((config) => ({
   ...config,
@@ -78,3 +78,9 @@ The current date and time is contained in the most recent user message.
 {{instructions}}
 `,
 });
+
+export const GraniteBeeSchemaErrorPrompt = BeeSchemaErrorPrompt.fork((config) => ({
+  ...config,
+  template: `Error: The generated response does not adhere to the communication structure mentioned in the system prompt.
+You communicate only in instruction lines. Valid instruction lines are 'Thought' followed by either 'Tool Name' + 'Tool Input' or 'Final Answer'.`,
+}));
diff --git a/src/agents/bee/runners/granite/runner.ts b/src/agents/bee/runners/granite/runner.ts
@@ -23,6 +23,7 @@ import { BeeAgent, BeeInput } from "@/agents/bee/agent.js";
 import type { GetRunContext } from "@/context.js";
 import {
   GraniteBeeAssistantPrompt,
+  GraniteBeeSchemaErrorPrompt,
   GraniteBeeSystemPrompt,
 } from "@/agents/bee/runners/granite/prompts.js";
 
@@ -39,6 +40,7 @@ export class GraniteRunner extends DefaultRunner {
           ...input.templates,
           system: input.templates?.system ?? GraniteBeeSystemPrompt,
           assistant: input.templates?.assistant ?? GraniteBeeAssistantPrompt,
+          schemaError: input.templates?.schemaError ?? GraniteBeeSchemaErrorPrompt,
         },
       },
       options,

diff --git a/src/agents/bee/types.ts b/src/agents/bee/types.ts
@@ -21,6 +21,7 @@ import { Callback } from "@/emitter/types.js";
 import { AnyTool, BaseToolRunOptions, ToolError, ToolOutput } from "@/tools/base.js";
 import {
   BeeAssistantPrompt,
+  BeeSchemaErrorPrompt,
   BeeSystemPrompt,
   BeeToolErrorPrompt,
   BeeToolInputErrorPrompt,
@@ -128,6 +129,7 @@ export interface BeeAgentTemplates {
   toolInputError: typeof BeeToolInputErrorPrompt;
   toolNoResultError: typeof BeeToolNoResultsPrompt;
   toolNotFoundError: typeof BeeToolNotFoundPrompt;
+  schemaError: typeof BeeSchemaErrorPrompt;
 }
 
 export type BeeParserInput = LinePrefixParser.define<{