diff --git a/lagent/actions/ipython_interactive.py b/lagent/actions/ipython_interactive.py index c39c8893..09dc8702 100644 --- a/lagent/actions/ipython_interactive.py +++ b/lagent/actions/ipython_interactive.py @@ -52,8 +52,13 @@ def __init__( ): super().__init__(description, parser, enable) from IPython import InteractiveShell + from traitlets.config import Config + self.timeout = timeout - self._executor = InteractiveShell() + c = Config() + c.HistoryManager.enabled = False + c.HistoryManager.hist_file = ':memory:' + self._executor = InteractiveShell(config=c) self._highlighting = re.compile(r'\x1b\[\d{,3}(;\d{,3}){,3}m') self._max_out_len = max_out_len if max_out_len >= 0 else None self._use_signals = use_signals diff --git a/lagent/actions/ipython_manager.py b/lagent/actions/ipython_manager.py index 4085f8cd..1b32e190 100644 --- a/lagent/actions/ipython_manager.py +++ b/lagent/actions/ipython_manager.py @@ -46,16 +46,19 @@ def run(self): self.out_q.put('ok') elif isinstance(msg, tuple) and len(msg) == 3: i, session_id, code = msg - res = tm(self.timeout)(self.exec)(session_id, code) + res = self.exec(session_id, code) self.out_q.put((i, session_id, res)) def exec(self, session_id, code): try: + shell = self.session_id2shell[session_id] with StringIO() as io: old_stdout = sys.stdout sys.stdout = io - self.session_id2shell[session_id].run_cell( - self.extract_code(code)) + if self.timeout is False or self.timeout < 0: + shell.run_cell(self.extract_code(code)) + else: + tm(self.timeout)(shell.run_cell)(self.extract_code(code)) sys.stdout = old_stdout output = self._highlighting.sub('', io.getvalue().strip()) output = re.sub(r'^Out\[\d+\]: ', '', output) @@ -63,10 +66,10 @@ def exec(self, session_id, code): output = output.lstrip('-').strip() if output.startswith('TimeoutError'): output = 'The code interpreter encountered a timeout error.' - return {'status': 'FAILURE', 'msg': output} - return {'status': 'SUCCESS', 'value': output} + return {'status': 'FAILURE', 'msg': output, 'code': code} + return {'status': 'SUCCESS', 'value': output, 'code': code} except Exception as e: - return {'status': 'FAILURE', 'msg': str(e)} + return {'status': 'FAILURE', 'msg': str(e), 'code': code} @staticmethod def create_shell(enable_history: bool = False, in_memory: bool = True):