| 7 | from bmtools.agent.translator import Translator |
| 8 | |
| 9 | class AgentExecutorWithTranslation(AgentExecutor): |
| 10 | |
| 11 | translator: Translator = Translator() |
| 12 | |
| 13 | def prep_outputs( |
| 14 | self, |
| 15 | inputs: Dict[str, str], |
| 16 | outputs: Dict[str, str], |
| 17 | return_only_outputs: bool = False, |
| 18 | ) -> Dict[str, str]: |
| 19 | try: |
| 20 | outputs = super().prep_outputs(inputs, outputs, return_only_outputs) |
| 21 | except ValueError as e: |
| 22 | return outputs |
| 23 | else: |
| 24 | if "input" in outputs: |
| 25 | outputs = self.translator(outputs) |
| 26 | return outputs |
| 27 | |
| 28 | class Executor(AgentExecutorWithTranslation): |
| 29 | def _call(self, inputs: Dict[str, str]) -> Dict[str, Any]: |