import { OutputParserArgs } from "../agent.js";
import { AgentActionOutputParser } from "../types.js";

import { FORMAT_INSTRUCTIONS } from "./prompt.js";

export const FINAL_ANSWER_ACTION = "Final Answer:";
export class ZeroShotAgentOutputParser extends AgentActionOutputParser {
  lc_namespace = ["langchain", "agents", "mrkl"];

  finishToolName: string;

  constructor(fields?: OutputParserArgs) {
    super(fields);
    this.finishToolName = fields?.finishToolName || FINAL_ANSWER_ACTION;
  }

  async parse(text: string) {
    if (text.includes(this.finishToolName)) {
      const parts = text.split(this.finishToolName);
      const output = parts[parts.length - 1].trim();
      return {
        returnValues: { output },
        log: text,
      };
    }

    const match = /Action: ([\s\S]*?)(?:\nAction Input: ([\s\S]*?))?$/.exec(
      text
    );
    if (!match) {
      throw new Error(`Could not parse LLM output: ${text}`);
    }

    return {
      tool: match[1].trim(),
      toolInput: match[2]
        ? match[2].trim().replace(/^("+)(.*?)(\1)$/, "$2")
        : "",
      log: text,
    };
  }

  getFormatInstructions(): string {
    return FORMAT_INSTRUCTIONS;
  }
}