The Code Execution Loop

Code Agents: Writing and Running Code

The Write-Execute-Observe Loop

def code_execution_loop(task: str, max_iterations=10):
    messages = [
        {'role': 'system', 'content': 'You are a Python coding agent. Write code to solve tasks step by step.'},
        {'role': 'user', 'content': task}
    ]
    
    for i in range(max_iterations):
        # WRITE: LLM generates code
        response = llm.complete(messages)
        code = extract_code_block(response)
        
        if not code:
            return response  # LLM gave a final answer without code
        
        # EXECUTE: run the code
        output, error = execute_safely(code)
        
        # OBSERVE: feed output back
        observation = f'Output:\n{output}' if not error else f'Error:\n{error}'
        messages.append({'role': 'assistant', 'content': response})
        messages.append({'role': 'user', 'content': observation})
    
    return 'Max iterations reached'

Extracting Code Blocks from LLM Output

import re

def extract_code_block(response: str) -> str | None:
    # Match ```python ... ``` or ``` ... ```
    pattern = r'```(?:python)?\n(.*?)```'
    matches = re.findall(pattern, response, re.DOTALL)
    if not matches:
        return None
    # Return the LAST code block (often the most complete version)
    return matches[-1].strip()

# Test
sample_response = '''I will calculate the sum:\n\n```python\nnumbers = [1, 2, 3, 4, 5]\nprint(sum(numbers))\n```'''
code = extract_code_block(sample_response)
print(repr(code))  # 'numbers = [1, 2, 3, 4, 5]\nprint(sum(numbers))'

Capturing stdout and stderr

import subprocess
import tempfile
import os

def execute_code(code: str, timeout_seconds=30) -> tuple[str, str]:
    # Write code to a temp file
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
        f.write(code)
        tmp_path = f.name
    
    try:
        result = subprocess.run(
            ['python', tmp_path],
            capture_output=True,
            text=True,
            timeout=timeout_seconds
        )
        stdout = result.stdout[:5000]  # cap output length
        stderr = result.stderr[:2000]  # cap error length
        return stdout, stderr
    except subprocess.TimeoutExpired:
        return '', f'TimeoutError: code exceeded {timeout_seconds}s limit'
    finally:
        os.unlink(tmp_path)  # always clean up temp file

Handling Errors Intelligently

def format_error_observation(code: str, error: str) -> str:
    lines = code.split('\n')
    
    # Extract line number from traceback if available
    line_match = re.search(r'line (\d+)', error)
    error_context = ''
    if line_match:
        line_num = int(line_match.group(1))
        if 1 <= line_num <= len(lines):
            error_context = f'\nFailing line: {lines[line_num - 1].strip()}'
    
    return f'''Execution failed with error:{error_context}\n\n{error}\n\nPlease analyze the error and write corrected code.'''

# Usage in loop
output, err = execute_code(code)
if err:
    observation = format_error_observation(code, err)
else:
    observation = f'Output:\n{output}\n\nContinue with the next step or provide the final answer.'

Detecting Task Completion

COMPLETION_MARKERS = ['TASK COMPLETE', 'FINAL ANSWER:', 'DONE:']

def is_task_complete(response: str, code: str | None, output: str, success_fn=None) -> bool:
    # Check for explicit completion markers
    for marker in COMPLETION_MARKERS:
        if marker in response.upper():
            return True
    
    # No code generated - LLM is done
    if code is None:
        return True
    
    # Custom success function (e.g., check output file exists)
    if success_fn and success_fn(output):
        return True
    
    return False

# Example success function
def report_was_generated(output: str) -> bool:
    import os
    return os.path.exists('analysis_report.pdf')

Passing Data Between Iterations

# Strategy 1: Save to files between iterations
# Iteration 1: LLM writes
'''
import pandas as pd
df = pd.read_csv('data.csv')
df_cleaned = df.dropna()
df_cleaned.to_parquet('cleaned.parquet')  # save for next iteration
print('Cleaned rows:', len(df_cleaned))
'''

# Iteration 2: LLM reads previous output
'''
import pandas as pd
df = pd.read_parquet('cleaned.parquet')  # reload from file
result = df.groupby('category').sum()
result.to_csv('result.csv')
print(result.head())
'''

Structuring the System Prompt

CODE_AGENT_SYSTEM_PROMPT = '''
You are a Python code execution agent. Solve tasks by writing and running code.

Rules:
1. Write ALL code inside ```python ... ``` blocks.
2. Write small, testable code blocks - do not try to do everything in one block.
3. When you see execution output, analyze it and decide whether to continue or fix errors.
4. Available libraries: pandas, numpy, matplotlib, requests, json, os, pathlib.
5. You can read/write files only in /workspace/ directory.
6. When the task is fully complete, say TASK COMPLETE and summarize the result.
7. If you are stuck after 3 attempts at the same error, say ESCALATE and explain the problem.
'''

Output Truncation and Token Management

MAX_OUTPUT_CHARS = 3000

def format_observation(stdout: str, stderr: str) -> str:
    if stderr:
        # Errors take priority - show full error
        return f'Error:\n{stderr[:MAX_OUTPUT_CHARS]}'
    
    if len(stdout) > MAX_OUTPUT_CHARS:
        truncated = stdout[:MAX_OUTPUT_CHARS]
        remaining = len(stdout) - MAX_OUTPUT_CHARS
        return f'{truncated}\n... [output truncated, {remaining} more characters]'
    
    return f'Output:\n{stdout}' if stdout else 'Code executed successfully (no output)'

The Code Loop in Practice

Security Considerations in Code Loops

Quick Check

Lesson Recap