Coverage for mindsdb / interfaces / agents / safe_output_parser.py: 50%

1import re 

2from typing import Union 

3 

4from langchain_core.agents import AgentAction, AgentFinish 

5 

6from langchain.agents.agent import AgentOutputParser 

7from langchain.agents.conversational.prompt import FORMAT_INSTRUCTIONS 

8 

9from mindsdb.utilities import log 

10 

11logger = log.getLogger(__name__) 

12 

13 

14class SafeOutputParser(AgentOutputParser): 

15 '''Output parser for the conversational agent that does not throw OutputParserException.''' 

16 

17 ai_prefix: str = 'AI' 

18 '''Prefix to use before AI output.''' 

19 

20 format_instructions: str = FORMAT_INSTRUCTIONS 

21 '''Default formatting instructions''' 

22 

23 def get_format_instructions(self) -> str: 

24 '''Returns formatting instructions for the given output parser.''' 

25 return self.format_instructions 

26 

27 def parse(self, text: str) -> Union[AgentAction, AgentFinish]: 

28 '''Parses outputted text from an LLM. 

29 

30 Args: 

31 text (str): Outputted text to parse. 

32 

33 Returns: 

34 output (str): Parsed text to an Agent step. 

35 ''' 

36 regex = r'Action: (.*?)[\n]*Action Input:([\s\S]*)' 

37 match = re.search(regex, text, re.DOTALL) 

38 if match is not None: 

39 action = match.group(1) 

40 action_input = match.group(2) 

41 return AgentAction(action.strip(), action_input.strip(' ').strip('"'), text) 

42 output = text 

43 if f'{self.ai_prefix}:' in text: 

44 output = text.split(f'{self.ai_prefix}:')[-1].strip() 

45 return AgentFinish( 

46 {'output': output}, text 

47 ) 

48 

49 @property 

50 def _type(self) -> str: 

51 return 'conversational'