forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
27 lines
978 B
Python
27 lines
978 B
Python
import re
|
|
from typing import Union
|
|
|
|
from langchain.agents.agent import AgentOutputParser
|
|
from langchain.agents.conversational.prompt import FORMAT_INSTRUCTIONS
|
|
from langchain.schema import AgentAction, AgentFinish, OutputParserException
|
|
|
|
|
|
class ConvoOutputParser(AgentOutputParser):
|
|
ai_prefix: str = "AI"
|
|
|
|
def get_format_instructions(self) -> str:
|
|
return FORMAT_INSTRUCTIONS
|
|
|
|
def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
|
|
if f"{self.ai_prefix}:" in text:
|
|
return AgentFinish(
|
|
{"output": text.split(f"{self.ai_prefix}:")[-1].strip()}, text
|
|
)
|
|
regex = r"Action: (.*?)[\n]*Action Input: (.*)"
|
|
match = re.search(regex, text)
|
|
if not match:
|
|
raise OutputParserException(f"Could not parse LLM output: `{text}`")
|
|
action = match.group(1)
|
|
action_input = match.group(2)
|
|
return AgentAction(action.strip(), action_input.strip(" ").strip('"'), text)
|