From 9f41663092869fce1dd7f622795db2e46797e0a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Dvo=C5=99=C3=A1k?= Date: Fri, 27 Sep 2024 14:20:33 +0200 Subject: [PATCH] feat(agent): switch to a new system prompt (#38) --- src/agents/bee/agent.ts | 6 +- src/agents/bee/parser.test.ts | 182 ------------------ src/agents/bee/parser.ts | 337 +--------------------------------- src/agents/bee/prompts.ts | 110 +++++------ src/agents/bee/runner.ts | 92 ++++++++-- src/agents/bee/types.ts | 12 +- 6 files changed, 147 insertions(+), 592 deletions(-) delete mode 100644 src/agents/bee/parser.test.ts diff --git a/src/agents/bee/agent.ts b/src/agents/bee/agent.ts index 41d06c8..e635d66 100644 --- a/src/agents/bee/agent.ts +++ b/src/agents/bee/agent.ts @@ -36,6 +36,7 @@ import { GetRunContext } from "@/context.js"; import { BeeAgentRunner } from "@/agents/bee/runner.js"; import { BeeAgentError } from "@/agents/bee/errors.js"; import { BeeIterationToolResult } from "@/agents/bee/parser.js"; +import { assign } from "@/internals/helpers/object.js"; export interface BeeInput { llm: ChatLLM; @@ -128,9 +129,10 @@ export class BeeAgent extends BaseAgent { - describe("Parsing", () => { - it("Basics", async () => { - const parser = new BeeOutputParser({}); - await parser.add("Final Answer: I need to find the current president of the Czech Republic."); - await parser.finalize(); - parser.validate(); - - const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy)); - expect(result).toMatchInlineSnapshot(` - { - "final_answer": "I need to find the current president of the Czech Republic.", - } - `); - }); - - it("Ends up with the same result", async () => { - const partial = new Map(); - const final = new Map(); - - const parser = new BeeOutputParser(); - parser.emitter.on("update", async ({ update, type }) => { - if (type === "full") { - final.set(update.key, update.value); - } else { - partial.set(update.key, (partial.get(update.key) ?? "").concat(update.value)); - } - }); - await parser.add("Thought: I "); - await parser.add("will do it."); - await parser.finalize(); - parser.validate(); - expect(partial).toStrictEqual(final); - }); - - it("Parses chunked JSON", async () => { - const parser = new BeeOutputParser({}); - await parser.add("Tool Name:\n"); - await parser.add("Goo"); - await parser.add("gle"); - await parser.add("\n"); - await parser.add("Tool "); - await parser.add("Input:\n"); - await parser.add('{"query'); - await parser.add('": "Czech President"'); - await parser.add("}"); - await parser.finalize(); - parser.validate(); - - const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy)); - expect(result).toMatchInlineSnapshot(` - { - "tool_input": { - "query": "Czech President", - }, - "tool_name": "Google", - } - `); - }); - - it("Handles newlines", async () => { - const parser = new BeeOutputParser({ - allowMultiLines: true, - preserveNewLines: false, - trimContent: true, - }); - await parser.add("Final Answer:\n\nI need to find\n\n the fastest car. \n "); - await parser.finalize(); - parser.validate(); - - const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy)); - expect(result).toMatchInlineSnapshot(` - { - "final_answer": "I need to find the fastest car.", - } - `); - }); - - it("Ignores newlines before first keyword occurrence", async () => { - const parser = new BeeOutputParser(); - await parser.add(""); - await parser.add(" \n"); - await parser.add(" \n "); - await parser.add(""); - await parser.add("\n Final Answer: Hello"); - await parser.finalize(); - parser.validate(); - - const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy)); - expect(result).toMatchInlineSnapshot(` - { - "final_answer": "Hello", - } - `); - }); - - it("Handles newlines with custom settings", async () => { - const parser = new BeeOutputParser({ - allowMultiLines: true, - preserveNewLines: true, - trimContent: true, - }); - await parser.add("Final Answer:\n\nI need to find\n\n the fastest car. \n "); - await parser.finalize(); - parser.validate(); - - const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy)); - expect(result).toMatchInlineSnapshot(` - { - "final_answer": "I need to find - the fastest car.", - } - `); - }); - }); - - describe("Chunking", () => { - it.each([ - " F#inal #answer : #I need to# search #Colorado, find#\n#\n the\n #area th#at th#e easter#n secto#r of# the Colora#do ex#tends i#nto, then find th#e elev#ation# #range #of the area.\n\n\n\nExtra Content.", - "\nfinal answer:A#B#C###", - ])("Text", async (text) => { - const parser = new BeeOutputParser({ - allowMultiLines: true, - preserveNewLines: true, - trimContent: false, - }); - - const chunks = text.split("#"); - for (const chunk of chunks) { - await parser.add(chunk); - } - - await parser.finalize(); - parser.validate(); - - const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy)); - expect(result).toMatchSnapshot(); - }); - }); - - describe("Validation", () => { - it("Throws when no data passed", () => { - const parser = new BeeOutputParser(); - expect(omitEmptyValues(parser.parse())).toStrictEqual({}); - expect(() => parser.validate()).toThrowError(BeeOutputParserError); - }); - - it.each(["Hello\nWorld", "Tool{\nxxx", "\n\n\nT"])( - "Throws on invalid data (%s)", - async (chunk) => { - const parser = new BeeOutputParser({ - allowMultiLines: false, - }); - await expect(parser.add(chunk)).rejects.toThrowError(BeeOutputParserError); - await expect(parser.finalize()).rejects.toThrowError(BeeOutputParserError); - expect(() => parser.validate()).toThrowError(BeeOutputParserError); - }, - ); - }); -}); diff --git a/src/agents/bee/parser.ts b/src/agents/bee/parser.ts index 059eed6..6f32297 100644 --- a/src/agents/bee/parser.ts +++ b/src/agents/bee/parser.ts @@ -14,338 +14,11 @@ * limitations under the License. */ -import * as R from "remeda"; -import { FrameworkError } from "@/errors.js"; -import { Cache } from "@/cache/decoratorCache.js"; -import { halveString } from "@/internals/helpers/string.js"; -import { parseBrokenJson } from "@/internals/helpers/schema.js"; -import { Emitter } from "@/emitter/emitter.js"; import { NonUndefined } from "@/internals/types.js"; -import { sumBy } from "remeda"; - -export interface BeeOutputParserOptions { - allowMultiLines: boolean; - preserveNewLines: boolean; - trimContent: boolean; -} - -interface BeeIterationSerializedResult { - thought?: string; - tool_name?: string; - tool_caption?: string; - tool_input?: string; - tool_output?: string; - final_answer?: string; -} - -export interface BeeIterationResult extends Omit { - tool_input?: unknown; -} +import type { LinePrefixParser } from "@/agents/parsers/linePrefix.js"; +import type { BeeAgentRunner } from "@/agents/bee/runner.js"; +type Parser = ReturnType["parser"]; +export type BeeIterationResult = LinePrefixParser.inferOutput; +export type BeeIterationResultPartial = LinePrefixParser.inferPartialOutput; export type BeeIterationToolResult = NonUndefined; - -export class BeeOutputParserError extends FrameworkError {} - -const NEW_LINE_CHARACTER = "\n"; - -interface Callbacks { - update: (data: { - type: "full" | "partial"; - state: BeeIterationResult; - update: { - key: keyof BeeIterationResult; - value: string; - }; - }) => Promise; -} - -class RepetitionChecker { - protected stash: string[] = []; - public enabled = true; - - constructor( - protected readonly size: number, - protected threshold: number, - ) {} - - add(chunk: string) { - if (!this.enabled) { - return false; - } - - if (this.stash.length > this.size) { - this.stash.shift(); - } - this.stash.push(chunk); - - const occurrences = sumBy(this.stash, (token) => Number(token === chunk)); - return occurrences >= this.threshold; - } - - reset() { - this.stash.length = 0; - } -} - -export class BeeOutputParser { - public isDone = false; - protected readonly lines: string[]; - protected lastKeyModified: keyof BeeIterationSerializedResult | null = null; - public stash: string; - public readonly emitter = new Emitter({ - creator: this, - namespace: ["agent", "bee", "parser"], - }); - - static Config = { - repetition: { - size: 20, - threshold: 10, - enabled: false, - }, - }; - - public repetitionChecker: RepetitionChecker; - protected readonly options: BeeOutputParserOptions; - - protected readonly result: BeeIterationSerializedResult = { - thought: undefined, - tool_name: undefined, - tool_caption: undefined, - tool_input: undefined, - tool_output: undefined, - final_answer: undefined, - }; - - constructor(options?: Partial) { - this.options = { - ...this._defaultOptions, - ...options, - }; - this.lines = []; - this.stash = ""; - this.repetitionChecker = new RepetitionChecker( - BeeOutputParser.Config.repetition.size, - BeeOutputParser.Config.repetition.threshold, - ); - this.repetitionChecker.enabled = BeeOutputParser.Config.repetition.enabled; - } - - async add(chunk: string) { - if (this.isDone || !chunk) { - return; - } - - chunk = chunk ?? ""; - - const isRepeating = this.repetitionChecker.add(chunk); - if (isRepeating) { - chunk = NEW_LINE_CHARACTER; - } - - this.stash += chunk; - if (!chunk.includes(NEW_LINE_CHARACTER)) { - return; - } - - while (this.stash.includes(NEW_LINE_CHARACTER)) { - this.repetitionChecker.reset(); - this.filterStash(); - const [line, stash = ""] = halveString(this.stash, NEW_LINE_CHARACTER); - this.stash = stash; - - await this._processChunk(line); - } - - if (isRepeating) { - this.isDone = true; - } - } - - protected filterStash() { - this.stash = this.stash.replaceAll("<|eom_id|>", ""); - this.stash = this.stash.replaceAll("<|eot_id|>", ""); - this.stash = this.stash.replaceAll("<|start_header_id|>assistant<|end_header_id|>", ""); - this.stash = this.stash.replaceAll("<|start_header_id|>", ""); - this.stash = this.stash.replaceAll("<|end_header_id|>", ""); - this.stash = this.stash.replaceAll("<|im_start|>", ""); - this.stash = this.stash.replaceAll("<|im_end|>", ""); - } - - async finalize() { - if (this.stash) { - await this._processChunk(this.stash); - this.stash = ""; - } - - if (this.isEmpty()) { - const response = this.lines.join(NEW_LINE_CHARACTER).concat(this.stash); - this.lines.length = 0; - this.stash = ""; - - await this.add(`Thought: ${response}${NEW_LINE_CHARACTER}`); - await this.add(`Final Answer: ${response}${NEW_LINE_CHARACTER}`); - } - if (this.result.thought && !this.result.final_answer && !this.result.tool_input) { - this.stash = ""; - await this.add(`Final Answer: ${this.result.thought}${NEW_LINE_CHARACTER}`); - } - - if (this.lastKeyModified) { - const parsed = this.parse(); - await this.emitter.emit("update", { - type: "full", - state: parsed, - update: { - key: this.lastKeyModified, - value: this.result[this.lastKeyModified]!, - }, - }); - } - this.lastKeyModified = null; - } - - isEmpty() { - return R.isEmpty(R.pickBy(this.result, R.isTruthy)); - } - - validate() { - if (this.isEmpty()) { - throw new BeeOutputParserError("Nothing valid has been parsed yet!", [], { - context: { - raw: this.lines.join(NEW_LINE_CHARACTER), - stash: this.stash, - }, - }); - } - - const { final_answer, tool_name, tool_input } = this.parse(); - const context = { - result: this.parse(), - stash: this.stash, - }; - if (!final_answer && !tool_input) { - if (this.result.tool_input) { - throw new BeeOutputParserError('Invalid "Tool Input" has been generated.', [], { - context: { - toolName: tool_name, - toolCaption: this.result.tool_caption, - toolInput: this.result.tool_input, - ...context, - }, - }); - } - - throw new BeeOutputParserError('Neither "Final Answer" nor "Tool Call" are present.', [], { - context, - }); - } - if (tool_input && final_answer) { - throw new BeeOutputParserError('Both "Final Answer" and "Tool Call" are present.', [], { - context, - }); - } - } - - @Cache() - protected get _defaultOptions(): BeeOutputParserOptions { - return { - allowMultiLines: true, - preserveNewLines: true, - trimContent: false, - }; - } - - protected async _processChunk(chunk: string) { - if (this.isDone) { - return; - } - - this.lines.push(chunk); - - let oldChunk = this.lastKeyModified ? this.result[this.lastKeyModified] : ""; - if (!this._extractStepPair(chunk) && this.options.allowMultiLines && this.lastKeyModified) { - const prev = this.result[this.lastKeyModified] || ""; - const newLine = this.options.preserveNewLines ? NEW_LINE_CHARACTER : ""; - chunk = `${this.lastKeyModified}:${prev}${newLine}${chunk}`; - } - - const step = this._extractStepPair(chunk); - if (!step && this.lastKeyModified === null && this.options.allowMultiLines) { - return; - } - - if (!step) { - throw new BeeOutputParserError(`No valid type has been detected in the chunk. (${chunk})}`); - } - - if (this.lastKeyModified && this.lastKeyModified !== step.type) { - this.isDone = Boolean(this.result[step.type]); - if (this.isDone) { - return; - } - - const state = this.parse(); - await this.emitter.emit("update", { - type: "full", - state, - update: { - key: this.lastKeyModified, - value: this.result[this.lastKeyModified]!, - }, - }); - oldChunk = this.result[step.type] ?? ""; - } - this.lastKeyModified = step.type; - - if (step.content) { - this.result[step.type] = step.content; - const state = this.parse(); - await this.emitter.emit("update", { - type: "partial", - state, - update: { - key: step.type, - value: step.content.replace(oldChunk ?? "", ""), - }, - }); - } - } - - parse(): BeeIterationResult { - const toolInput = parseBrokenJson(this?.result.tool_input?.trim?.(), { pair: ["{", "}"] }); - return R.pickBy( - Object.assign( - { ...this.result }, - { - tool_name: this.result.tool_name, - tool_input: toolInput ?? undefined, - }, - ), - R.isDefined, - ); - } - - protected _isValidStepType(type?: string | null): type is keyof BeeIterationSerializedResult { - return Boolean(type && type in this.result); - } - - protected _extractStepPair(line: string) { - let [, type, content] = line.match(/\s*([\w|\s]+?)\s*:\s*(.*)/ms) ?? [line, null, null]; - type = type ? type.trim().toLowerCase().replace(" ", "_") : null; - - if (!this._isValidStepType(type)) { - return null; - } - - content = content ?? ""; - if (this.options.trimContent) { - content = content.trim(); - } - - return { - type, - content, - }; - } -} diff --git a/src/agents/bee/prompts.ts b/src/agents/bee/prompts.ts index cd0d1ca..3e6af9d 100644 --- a/src/agents/bee/prompts.ts +++ b/src/agents/bee/prompts.ts @@ -17,7 +17,6 @@ import { PromptTemplate } from "@/template.js"; import { BaseMessageMeta } from "@/llms/primitives/message.js"; import { z } from "zod"; -import { PythonToolOutput } from "@/tools/python/output.js"; export const BeeSystemPrompt = new PromptTemplate({ schema: z.object({ @@ -33,72 +32,73 @@ export const BeeSystemPrompt = new PromptTemplate({ }) .passthrough(), ), - tool_names: z.string(), }), - template: `{{instructions}} - -# Tools - -Tools must be used to retrieve factual or historical information to answer the question. + template: `# Available functions {{#tools.length}} -A tool can be used by generating the following three lines: - -Tool Name: ZblorgColorLookup -Tool Caption: Searching Zblorg #178 -Tool Input: {"id":178} - -## Available tools {{#tools}} +Function Name: {{name}} +Description: {{description}} +Parameters: {{schema}} -Tool Name: {{name}} -Tool Description: {{description}} -Tool Input Schema: {{schema}} {{/tools}} {{/tools.length}} {{^tools.length}} +No functions are available. -## Available tools +{{/tools.length}} +# Communication structure +You communicate in instruction lines. +The format is: "Instruction: expected output". +You must not enter empty lines or anything else between instruction lines. +{{#tools.length}} +You must skip the instruction lines Function Name, Function Input, Function Caption and Function Output if no function use is required. +{{/tools.length}} -No tools are available at the moment therefore you mustn't provide any factual or historical information. -If you need to, you must respond that you cannot. +Question: User's question and other relevant input. You never use this instruction line. +{{^tools.length}} +Thought: A single-line explanation of what needs to happen next to be able to answer the user's question. It must be immediately followed by Final Answer. +{{/tools.length}} +{{#tools.length}} +Thought: A short plan of how to answer the user's question. It must be immediately followed by Function Name when one of available functions can be used to obtain more information, or by Final Answer when available information and capabilities are sufficient to provide the answer. +Function Name: Name of the function that can best answer the preceding Thought. It must be one of the available functions defined above. +Function Input: Parameters for the function to best answer the preceding Thought. You must follow the Parameters schema. +Function Caption: A short description of the function calling for the user. +Function Output: Output of the function in JSON format. +Thought: Repeat your thinking process. {{/tools.length}} +Final Answer: Either response to the original question and context once enough information is available or ask user for more information or clarification. + +## Examples +Question: What's your name? +Thought: The user wants to know my name. I have enough information to answer that. +Final Answer: My name is Bee. + +Question: Can you translate "How are you" into French? +Thought: The user wants to translate a text into French. I can do that. +Final Answer: Comment vas-tu? # Instructions +If you don't know the answer, say that you don't know. +{{^tools.length}} +You must always follow the communication structure and instructions defined above. Do not forget that Thought must be immediately followed by Final Answer. +{{/tools.length}} +{{#tools.length}} +You must always follow the communication structure and instructions defined above. Do not forget that Thought must be immediately followed by either Function Name or Final Answer. +Prefer to use your capabilities over functions. +Functions must be used to retrieve factual or historical information to answer the question. +{{/tools.length}} +If the user suggests using a function that is not available, answer politely that the function is not available. You can suggest alternatives if appropriate. +When the question is unclear or you need more information from the user, ask in Final Answer. + +# Your other capabilities +- You understand these languages: English, Spanish, French. +- You can translate and summarize, even long documents. + +# Notes +- Last message's time is the current date and time. -Responses must always have the following structure: -- The user's input starts with 'Question: ' followed by the question the user asked, for example, 'Question: What is the color of Zblorg #178?' - - The question may contain square brackets with a nested sentence, like 'What is the color of [The Zblorg with the highest score of the 2023 season is Zblorg #178.]?'. Just assume that the question regards the entity described in the bracketed sentence, in this case 'Zblorg #178'. -- Line starting 'Thought: ', explaining the thought, for example 'Thought: I don't know what Zblorg is, but given that I have a ZblorgColorLookup tool, I can assume that it is something that can have a color and I should use the ZblorgColorLookup tool to find out the color of Zblorg number 178.' - - In a 'Thought', it is either determined that a Tool Call will be performed to obtain more information, or that the available information is sufficient to provide the Final Answer. - - If a tool needs to be called and is available, the following lines will be: - - Line starting 'Tool Name: ' name of the tool that you want to use. - - Line starting 'Tool Caption: ' short description of the calling action. - - Line starting 'Tool Input: ' JSON formatted input adhering to the selected tool JSON Schema. - - Line starting 'Tool Output: ', containing the tool output, for example 'Tool Output: {"success": true, "color": "green"}' - - The 'Tool Output' may or may not bring useful information. The following 'Thought' must determine whether the information is relevant and how to proceed further. - - If enough information is available to provide the Final Answer, the following line will be: - - Line starting 'Final Answer: ' followed by a response to the original question and context, for example: 'Final Answer: Zblorg #178 is green.' - - Use markdown syntax for formatting code snippets, links, JSON, tables, images, files. - - To reference an internal file, use the markdown syntax [file_name.ext](${PythonToolOutput.FILE_PREFIX}:file_identifier). - - The bracketed part must contain the file name, verbatim. - - The parenthesis part must contain the file URN, which can be obtained from the user or from tools. - - The agent does not, under any circumstances, reference a URN that was not provided by the user or a tool in the current conversation. - - To show an image, prepend an exclamation mark, as usual in markdown: ![file_name.ext](urn:file_identifier). - - This only applies to internal files. HTTP(S) links must be provided as is, without any modifications. -- The sequence of lines will be 'Thought' - ['Tool Name' - 'Tool Caption' - 'Tool Input' - 'Tool Output' - 'Thought'] - 'Final Answer', with the bracketed part repeating one or more times (but never repeating them in a row). Do not use empty lines between instructions. -- Sometimes, things don't go as planned. Tools may not provide useful information on the first few tries. The agent always tries a few different approaches before declaring the problem unsolvable: -- When the tool doesn't give you what you were asking for, you MUST either use another tool or a different tool input. - - When using search engines, the assistant tries different formulations of the query, possibly even in a different language. -- When executing code, the assistant fixes and retries when the execution errors out and tries a completely different approach if the code does not seem to be working. - - When the problem seems too hard for the tool, the assistant tries to split the problem into a few smaller ones. - -## Notes - -- Any comparison table (including its content), file, image, link, or other asset must only be in the Final Answer. -- When the question is unclear, respond with a line starting with 'Final Answer:' followed by the information needed to solve the problem. -- When the user wants to chitchat instead, always respond politely. -- IMPORTANT: Lines 'Thought', 'Tool Name', 'Tool Caption', 'Tool Input', 'Tool Output' and 'Final Answer' must be sent within a single message. -`, +# Role +{{instructions}}`, }); export const BeeAssistantPrompt = new PromptTemplate({ @@ -112,7 +112,7 @@ export const BeeAssistantPrompt = new PromptTemplate({ finalAnswer: z.array(z.string()), }) .partial(), - template: `{{#thought}}Thought: {{.}}\n{{/thought}}{{#toolName}}Tool Name: {{.}}\n{{/toolName}}{{#toolCaption}}Tool Caption: {{.}}\n{{/toolCaption}}{{#toolInput}}Tool Input: {{.}}\n{{/toolInput}}{{#toolOutput}}Tool Output: {{.}}\n{{/toolOutput}}{{#finalAnswer}}Final Answer: {{.}}{{/finalAnswer}}`, + template: `{{#thought}}Thought: {{.}}\n{{/thought}}{{#toolName}}Function Name: {{.}}\n{{/toolName}}{{#toolInput}}Function Input: {{.}}\n{{/toolInput}}{{#toolCaption}}Function Caption: {{.}}\n{{/toolCaption}}{{#toolOutput}}Function Output: {{.}}\n{{/toolOutput}}{{#finalAnswer}}Final Answer: {{.}}{{/finalAnswer}}`, }); export const BeeUserPrompt = new PromptTemplate({ diff --git a/src/agents/bee/runner.ts b/src/agents/bee/runner.ts index 732e001..dc85aa0 100644 --- a/src/agents/bee/runner.ts +++ b/src/agents/bee/runner.ts @@ -20,6 +20,7 @@ import { BaseMessage, Role } from "@/llms/primitives/message.js"; import { TokenMemory } from "@/memory/tokenMemory.js"; import { BeeAgentRunIteration, BeeCallbacks, BeeMeta, BeeRunOptions } from "@/agents/bee/types.js"; import { + AnyTool, BaseToolRunOptions, ToolError, ToolInputValidationError, @@ -39,9 +40,12 @@ import { BeeUserEmptyPrompt, BeeUserPrompt, } from "@/agents/bee/prompts.js"; -import { BeeIterationToolResult, BeeOutputParser } from "@/agents/bee/parser.js"; +import { BeeIterationToolResult } from "@/agents/bee/parser.js"; import { AgentError } from "@/agents/base.js"; import { Emitter } from "@/emitter/emitter.js"; +import { LinePrefixParser } from "@/agents/parsers/linePrefix.js"; +import { JSONParserField, ZodParserField } from "@/agents/parsers/field.js"; +import { z } from "zod"; export class BeeAgentRunnerFatalError extends BeeAgentError { isFatal = true; @@ -126,7 +130,6 @@ export class BeeAgentRunner { schema: JSON.stringify(await tool.getInputJsonSchema()), })), ), - tool_names: input.tools.map((tool) => tool.name).join(","), instructions: undefined, }), meta: { @@ -154,6 +157,58 @@ export class BeeAgentRunner { return new BeeAgentRunner(input, options, memory); } + static createParser(tools: AnyTool[]) { + 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)), + }, + } as const); + + return { + parser, + parserRegex, + } as const; + } + async llm(input: { emitter: Emitter; signal: AbortSignal; @@ -170,45 +225,48 @@ export class BeeAgentRunner { executor: async () => { await emitter.emit("start", { meta }); - const outputParser = new BeeOutputParser(); + const { parser, parserRegex } = BeeAgentRunner.createParser(this.input.tools); const llmOutput = await this.input.llm .generate(this.memory.messages.slice(), { signal, stream: true, guided: { - regex: - /Thought:.{1,512}\n(?:Final Answer:[\S\s]+|Tool Name:.+\nTool Caption:.+\nTool Input:\{.+\}\nTool Output:)?/ - .source, + regex: parserRegex.source, }, }) .observe((llmEmitter) => { - outputParser.emitter.on("update", async ({ type, update, state }) => { - await emitter.emit(type === "full" ? "update" : "partialUpdate", { - data: state, - update, + parser.emitter.on("update", async ({ value, key, field }) => { + await emitter.emit("update", { + data: parser.finalState, + update: { key, value: field.raw, parsedValue: value }, + meta: { success: true, ...meta }, + }); + }); + parser.emitter.on("partialUpdate", async ({ key, delta, value }) => { + await emitter.emit("partialUpdate", { + data: parser.finalState, + update: { key, value: delta, parsedValue: value }, meta: { success: true, ...meta }, }); }); llmEmitter.on("newToken", async ({ value, callbacks }) => { - if (outputParser.isDone) { + if (parser.isDone) { callbacks.abort(); return; } - await outputParser.add(value.getTextContent()); - if (outputParser.stash.match(/^\s*Tool Output:/i)) { - outputParser.stash = ""; + await parser.add(value.getTextContent()); + if (parser.partialState.tool_output !== undefined) { callbacks.abort(); } }); }); - await outputParser.finalize(); - outputParser.validate(); + await parser.end(); return { - state: outputParser.parse(), + state: parser.finalState, raw: llmOutput, }; }, diff --git a/src/agents/bee/types.ts b/src/agents/bee/types.ts index 604652e..42fba5c 100644 --- a/src/agents/bee/types.ts +++ b/src/agents/bee/types.ts @@ -15,7 +15,11 @@ */ import { ChatLLMOutput } from "@/llms/chat.js"; -import { BeeIterationResult, BeeIterationToolResult } from "@/agents/bee/parser.js"; +import { + BeeIterationResult, + BeeIterationResultPartial, + BeeIterationToolResult, +} from "@/agents/bee/parser.js"; import { BaseMemory } from "@/memory/base.js"; import { BaseMessage } from "@/llms/primitives/message.js"; import { Callback } from "@/emitter/types.js"; @@ -83,12 +87,12 @@ export interface BeeCallbacks { }>; update?: Callback<{ data: BeeIterationResult; - update: { key: keyof BeeIterationResult; value: string }; + update: { key: keyof BeeIterationResult; value: string; parsedValue: unknown }; meta: BeeUpdateMeta; }>; partialUpdate?: Callback<{ - data: BeeIterationResult; - update: { key: keyof BeeIterationResult; value: string }; + data: BeeIterationResultPartial; + update: { key: keyof BeeIterationResult; value: string; parsedValue: unknown }; meta: BeeUpdateMeta; }>; toolStart?: Callback<{