Function execute_code

optillm/plugins/coc_plugin.py:123–231

Attempt to execute the code using Jupyter notebook kernel and return result or error.

(code: str)

    return '\n'.join(safe_lines)

def execute_code(code: str) -> Tuple[Any, str]:
    """Attempt to execute the code using Jupyter notebook kernel and return result or error."""
    logger.info("Attempting to execute code in notebook kernel")
    logger.info(f"Code:\n{code}")
    
    try:
        # Sanitize the code first
        sanitized_code = sanitize_code(code)
        
        # Create a notebook with the code
        notebook = nbformat.v4.new_notebook()
        
        # Add code that captures the answer variable
        enhanced_code = f"""
{sanitized_code}

# Capture the answer variable for output
if 'answer' in locals():
    print(f"ANSWER_RESULT: {{answer}}")
else:
    print("ANSWER_RESULT: No answer variable found")
"""
        
        notebook['cells'] = [nbformat.v4.new_code_cell(enhanced_code)]
        
        # Convert notebook to JSON string and then to bytes
        notebook_json = nbformat.writes(notebook)
        notebook_bytes = notebook_json.encode('utf-8')

        # Create temporary notebook file
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.ipynb', delete=False) as tmp:
            tmp.write(notebook_bytes)
            tmp.flush()
            tmp_name = tmp.name

        try:
            # Read and execute the notebook
            with open(tmp_name, 'r', encoding='utf-8') as f:
                nb = nbformat.read(f, as_version=4)
            
            # Execute with timeout and isolation
            ep = ExecutePreprocessor(timeout=30, kernel_name='python3')
            ep.preprocess(nb, {'metadata': {'path': './'}})

            # Extract the output
            output = ""
            error_output = ""
            
            for cell in nb.cells:
                if cell.cell_type == 'code' and cell.outputs:
                    for output_item in cell.outputs:
                        if output_item.output_type == 'stream':
                            if output_item.name == 'stdout':
                                output += output_item.text
                            elif output_item.name == 'stderr':
                                error_output += output_item.text
                        elif output_item.output_type == 'execute_result':
                            output += str(output_item.data.get('text/plain', ''))