Шаблон интерпретатора кода для анализа данных

Шаблон интерпретатора кода

Основной цикл: сгенерировать → выполнить → интерпретировать

def code_interpreter_agent(question, data_path):
    # Step 1: Generate code
    code = generate_analysis_code(question, data_path)
    print('Generated code:', code[:200])

    # Step 2: Execute in sandbox
    result = execute_in_sandbox(code)

    if result['error']:
        # Try to fix the error
        fixed_code = fix_code(code, result['error'])
        result = execute_in_sandbox(fixed_code)

    # Step 3: Interpret output
    return interpret_output(question, result)

print(code_interpreter_agent(
    'What is the average order value by customer segment?',
    'data/orders.csv'
))

Запрос для генерации кода

CODE_GEN_PROMPT = """You are a Python data analyst. Write Python code to answer the question.
Data available at: {data_path}
Question: {question}
Write ONLY Python code (no markdown, no explanation):"""

def llm_call(prompt):
    return "```python\nprint('df.describe() results')\n```"

def generate_analysis_code(question, data_path):
    response = llm_call(CODE_GEN_PROMPT.format(question=question, data_path=data_path))
    code = response.strip()
    if code.startswith('```'):
        code = code.split('```')[1]
        if code.startswith('python'):
            code = code[6:]
    return code.strip()

print(generate_analysis_code('What is the average price?', 'data.csv'))

Выполнение в изолированной среде через подпроцесс

import subprocess
import tempfile
import os

def execute_in_sandbox(code, timeout=30):
    # Write code to temp file
    with tempfile.NamedTemporaryFile(suffix='.py', mode='w', delete=False) as f:
        f.write(code)
        script_path = f.name

    try:
        result = subprocess.run(
            ['python3', script_path],
            capture_output=True,
            text=True,
            timeout=timeout,
            env={**os.environ, 'MPLBACKEND': 'Agg'}  # non-interactive matplotlib
        )
        return {
            'stdout': result.stdout,
            'stderr': result.stderr,
            'returncode': result.returncode,
            'error': result.stderr if result.returncode != 0 else None
        }
    except subprocess.TimeoutExpired:
        return {'stdout': '', 'stderr': 'Timeout', 'returncode': -1, 'error': 'Code timed out'}
    finally:
        os.unlink(script_path)

if __name__ == '__main__':
    result = execute_in_sandbox('print(2 + 2)')
    print('Sandbox stdout:', result['stdout'].strip())
    print('Return code   :', result['returncode'])

Облачная изолированная среда E2B

from e2b_code_interpreter import Sandbox
import os

def execute_with_e2b(code, data_bytes=None):
    with Sandbox(api_key=os.getenv('E2B_API_KEY')) as sandbox:
        # Upload data file if provided
        if data_bytes:
            sandbox.files.write('/home/user/data.csv', data_bytes)

        # Execute code
        execution = sandbox.run_code(code)

        result = {
            'stdout': '\n'.join(execution.logs.stdout),
            'stderr': '\n'.join(execution.logs.stderr),
            'error': None
        }

        # Check for errors
        if execution.error:
            result['error'] = str(execution.error)

        # Download any generated files
        result['files'] = []
        for output in execution.results:
            if hasattr(output, 'png'):
                result['files'].append({
                    'type': 'image/png',
                    'data': output.png  # base64 encoded
                })

    return result

Получение сгенерированных файлов

import os
import glob
import base64

OUTPUT_DIR = '/tmp/chart_output'

def execute_and_capture(code, timeout=30):
    # Create output dir
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    result = execute_in_sandbox(code, timeout=timeout)

    # Capture any generated image files
    generated_files = []
    for filepath in glob.glob(os.path.join(OUTPUT_DIR, '*.png')):
        with open(filepath, 'rb') as f:
            encoded = base64.b64encode(f.read()).decode('utf-8')
        generated_files.append({
            'filename': os.path.basename(filepath),
            'type': 'image/png',
            'base64': encoded
        })
        os.unlink(filepath)  # clean up

    result['generated_files'] = generated_files
    print(f'Captured {len(generated_files)} file(s) from sandbox')
    return result

Цикл восстановления после ошибок

FIX_PROMPT = '''The following Python code raised an error. Fix it.

Original code:
{code}

Error:
{error}

Return ONLY the fixed Python code (no explanation, no markdown):'''

def fix_code(code, error):
    return llm_call(FIX_PROMPT.format(code=code, error=error)).strip()

def execute_with_retry(code, max_retries=2):
    for attempt in range(max_retries + 1):
        result = execute_and_capture(code)
        if not result['error']:
            return result
        print(f'Attempt {attempt + 1} failed: {result["error"][:100]}')
        if attempt < max_retries:
            code = fix_code(code, result['error'])
    return result  # return last result even if errored

Интерпретация вывода кода

INTERPRET_PROMPT = '''A Python script was executed to answer a data analysis question.
Explain the results in clear, non-technical language.

Original question: {question}

Code output (stdout):
{output}

Provide a clear, concise answer that directly addresses the question.
Highlight the most important numbers or findings.
Answer:'''

def interpret_output(question, execution_result):
    stdout = execution_result.get('stdout', '').strip()
    error = execution_result.get('error')

    if error and not stdout:
        return f'The analysis failed with error: {error}'

    if not stdout:
        return 'The code ran successfully but produced no output.'

    return llm_call(INTERPRET_PROMPT.format(
        question=question,
        output=stdout[:3000]  # truncate very long outputs
    ))

Ограничения безопасности при генерации кода

BLOCKED_IMPORTS = ['requests', 'httpx', 'urllib', 'socket', 'subprocess', 'os.system']

def pre_validate_code(code):
    errors = []
    for blocked in BLOCKED_IMPORTS:
        if f'import {blocked}' in code or f'from {blocked}' in code:
            errors.append(f'Blocked import: {blocked}')

    # Block shell execution
    import re
    if re.search(r'os\.system|subprocess\.run|subprocess\.call|eval\(|exec\(', code):
        errors.append('Blocked: shell execution or eval/exec')

    # Block reading outside allowed paths
    if re.search(r'open\([^)]*\.\./|open\([^)]*\/etc\/', code):
        errors.append('Blocked: unauthorized file access')

    if errors:
        raise ValueError('Security check failed:\n' + '\n'.join(errors))

    return True

if __name__ == '__main__':
    try:
        pre_validate_code('import requests\nrequests.get("http://x")')
    except ValueError as e:
        print('Rejected:', e)
    print('Safe code passed:', pre_validate_code('print(1 + 1)'))

Передача схемы данных

import pandas as pd

def get_data_schema(data_path):
    df = pd.read_csv(data_path, nrows=5)
    schema_lines = []
    for col in df.columns:
        dtype = str(df[col].dtype)
        sample = df[col].dropna().iloc[0] if len(df[col].dropna()) > 0 else 'N/A'
        schema_lines.append(f'  - {col} ({dtype}): sample={sample!r}')
    schema_text = '\n'.join(schema_lines)
    return f'CSV columns:\n{schema_text}\nTotal rows: {len(pd.read_csv(data_path))}'

ENHANCED_PROMPT = CODE_GEN_PROMPT + '\n\nData schema:\n{schema}'

def generate_analysis_code_with_schema(question, data_path):
    schema = get_data_schema(data_path)
    response = llm_call(ENHANCED_PROMPT.format(
        question=question, data_path=data_path, schema=schema
    ))
    return response.strip()

Отслеживание истории выполнения

from datetime import datetime

execution_history = []

def record_execution(question, code, result):
    execution_history.append({
        'timestamp': datetime.now().isoformat(),
        'question': question,
        'code_lines': len(code.splitlines()),
        'stdout_preview': result.get('stdout', '')[:200],
        'success': result.get('error') is None,
        'generated_files': len(result.get('generated_files', []))
    })

def get_session_summary():
    total = len(execution_history)
    successful = sum(1 for e in execution_history if e['success'])
    return {
        'total_executions': total,
        'successful': successful,
        'failed': total - successful,
        'success_rate': f'{successful/max(total,1)*100:.0f}%',
        'questions_answered': [e['question'][:60] for e in execution_history]
    }

# Usage in agent loop
def code_interpreter_agent_tracked(question, data_path):
    code = generate_analysis_code_with_schema(question, data_path)
    result = execute_with_retry(code)
    record_execution(question, code, result)
    return interpret_output(question, result)

Проверка знаний

Повторение: шаблон интерпретатора кода для анализа данных