Allow access to code from interpreter results (#191)

improve code interpreters
InternLM · Jun 13, 2024 · 264ae5c · 264ae5c
1 parent e1e25c3
commit 264ae5c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git 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
@@ -46,27 +46,30 @@ 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)
             if 'Error' in output or 'Traceback' in output:
                 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):