Execute Python code in a Jupyter notebook environment.
(code: str)
| 23 | return re.findall(pattern, text, re.DOTALL) |
| 24 | |
| 25 | def execute_code(code: str) -> str: |
| 26 | """Execute Python code in a Jupyter notebook environment.""" |
| 27 | |
| 28 | notebook = nbformat.v4.new_notebook() |
| 29 | notebook['cells'] = [nbformat.v4.new_code_cell(code)] |
| 30 | |
| 31 | # Convert notebook to JSON string |
| 32 | notebook_json = nbformat.writes(notebook) |
| 33 | |
| 34 | # Convert JSON string to bytes |
| 35 | notebook_bytes = notebook_json.encode('utf-8') |
| 36 | |
| 37 | with tempfile.NamedTemporaryFile(mode='wb', suffix='.ipynb', delete=False) as tmp: |
| 38 | tmp.write(notebook_bytes) |
| 39 | tmp.flush() |
| 40 | tmp_name = tmp.name |
| 41 | |
| 42 | try: |
| 43 | with open(tmp_name, 'r', encoding='utf-8') as f: |
| 44 | nb = nbformat.read(f, as_version=4) |
| 45 | ep = ExecutePreprocessor(timeout=30, kernel_name='python3') |
| 46 | ep.preprocess(nb, {'metadata': {'path': './'}}) |
| 47 | |
| 48 | # Extract the output |
| 49 | output = "" |
| 50 | for cell in nb.cells: |
| 51 | if cell.cell_type == 'code' and cell.outputs: |
| 52 | for output_item in cell.outputs: |
| 53 | if output_item.output_type == 'stream': |
| 54 | output += output_item.text |
| 55 | elif output_item.output_type == 'execute_result': |
| 56 | output += str(output_item.data.get('text/plain', '')) |
| 57 | |
| 58 | return output.strip() |
| 59 | finally: |
| 60 | os.unlink(tmp_name) |
| 61 | |
| 62 | def should_execute_request_code(query: str) -> bool: |
| 63 | """Decide whether to execute code from the request based on the query.""" |