Spaces:
Running
Running
from typing import Union | |
import re | |
from langchain.agents import AgentOutputParser | |
from langchain.schema import AgentAction, AgentFinish | |
class CustomOutputParser(AgentOutputParser): | |
""" | |
This is the output parser for the multi-tool agent. It parses the output from the LLM model. | |
""" | |
def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]: | |
""" | |
This function is used to parse the output from the LLM model. It checks if the output is the final answer or an action. | |
""" | |
if "Final Answer:" in llm_output: | |
return AgentFinish( | |
return_values={"output": llm_output.split("Final Answer:")[-1].strip()}, | |
log=llm_output, | |
) | |
regex = r"Action: (.*?)[\n]*Action Input:[\s]*(.*)" | |
match = re.search(regex, llm_output, re.DOTALL) | |
if not match: | |
raise ValueError(f"Could not parse LLM output: `{llm_output}`") | |
action = match.group(1).strip() | |
action_input = match.group(2) | |
return AgentAction( | |
tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output | |
) |