跨执行步骤管理状态

# Iteration 1 - works fine
df = pd.read_csv('data.csv')  # df is in memory
df_cleaned = df.dropna()
print('Rows after cleaning:', len(df_cleaned))

# Iteration 2 - NEW container, df_cleaned is GONE
result = df_cleaned.groupby('category').sum()  # NameError: df_cleaned is not defined
print(result)  # This will fail!

import pandas as pd
import json
from pathlib import Path

WORKSPACE = Path('/workspace')

# Iteration 1: process and SAVE
df = pd.read_csv(WORKSPACE / 'raw_data.csv')
df_cleaned = df.dropna().reset_index(drop=True)
df_cleaned.to_parquet(WORKSPACE / 'cleaned.parquet')  # save for next iteration

stats = {'rows': len(df_cleaned), 'columns': list(df_cleaned.columns)}
with open(WORKSPACE / 'stats.json', 'w') as f:
    json.dump(stats, f)

print('Saved cleaned data:', len(df_cleaned), 'rows')

# Iteration 2: LOAD and continue
df_cleaned = pd.read_parquet(WORKSPACE / 'cleaned.parquet')  # restore state
with open(WORKSPACE / 'stats.json') as f:
    stats = json.load(f)
result = df_cleaned.groupby('category')['value'].sum()
print(result)

STATE_MANAGEMENT_INSTRUCTIONS = '''
State persistence rules:
- You have a persistent workspace at /workspace/
- After computing any result you will need later, SAVE it to /workspace/
  - DataFrames: use .to_parquet('/workspace/name.parquet')
  - Dicts/lists: use json.dump to /workspace/name.json
  - Text: write to /workspace/name.txt
- At the start of each code block, LOAD the files you need from previous steps
- Use clear, descriptive filenames like 'cleaned_data.parquet', not 'tmp1.parquet'
- When listing files, use: import os; print(os.listdir('/workspace/'))
'''

import jupyter_client

class PersistentKernel:
    def __init__(self):
        km, self.kc = jupyter_client.manager.start_new_kernel(kernel_name='python3')
        self.kc.wait_for_ready(timeout=30)
        print('Kernel started')

    def execute(self, code: str, timeout=60) -> tuple[str, str]:
        msg_id = self.kc.execute(code)
        outputs, errors = [], []
        
        while True:
            msg = self.kc.get_iopub_msg(timeout=timeout)
            if msg['msg_type'] == 'stream':
                if msg['content']['name'] == 'stdout':
                    outputs.append(msg['content']['text'])
                else:
                    errors.append(msg['content']['text'])
            if msg['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
                break
        
        return ''.join(outputs), ''.join(errors)

    def shutdown(self):
        self.kc.shutdown()

# Variables persist across execute() calls!
kernel = PersistentKernel()
kernel.execute('x = 42')  # x is now defined in kernel
output, _ = kernel.execute('print(x)')  # prints 42 - state persists!
kernel.shutdown()

from dataclasses import dataclass, field
from typing import Any

@dataclass
class ExecutionState:
    task: str
    iteration: int = 0
    workspace_files: list[str] = field(default_factory=list)
    computed_values: dict[str, Any] = field(default_factory=dict)
    completed_steps: list[str] = field(default_factory=list)
    last_output: str = ''

    def to_context_summary(self) -> str:
        return f'''Current task: {self.task}
Iteration: {self.iteration}
Completed steps: {', '.join(self.completed_steps) or 'None yet'}
Workspace files: {', '.join(self.workspace_files) or 'None yet'}
Key values: {self.computed_values}
Last output: {self.last_output[:500]}'''

    def after_execution(self, output: str, step_name: str):
        import os
        self.workspace_files = os.listdir('/workspace')
        self.completed_steps.append(step_name)
        self.last_output = output
        self.iteration += 1

def build_iteration_prompt(task: str, state: ExecutionState, last_observation: str) -> str:
    return f'''ORIGINAL TASK: {task}

CURRENT STATE:
{state.to_context_summary()}

LAST EXECUTION OUTPUT:
{last_observation}

What is the next step? Write Python code to continue. 
Remember:
- Load data from workspace files if needed
- Save any results you will need in future steps
- If the task is complete, say TASK COMPLETE and summarize the result'''

# Use in the main loop
for step_num in range(max_iterations):
    context = build_iteration_prompt(task, state, last_observation)
    response = llm.complete(messages + [{'role': 'user', 'content': context}])
    code = extract_code_block(response)
    if not code:
        break  # done
    output, err = execute_in_sandbox(code)
    state.after_execution(output, step_name=f'step_{step_num}')
    last_observation = format_observation(output, err)

# Good: load only needed columns
df = pd.read_parquet('/workspace/full_data.parquet', columns=['date', 'revenue', 'region'])

# Bad: load everything even if you only need 2 columns
# df = pd.read_parquet('/workspace/full_data.parquet')  # loads 50 columns you don't need

# Good: filter early before loading full dataset
df = pd.read_parquet('/workspace/full_data.parquet', filters=[('region', '=', 'EMEA')])

# For very large files, tell the LLM about data shape upfront
df_info = {'shape': (1_000_000, 50), 'size_mb': 850, 'columns': [...]}
# Include df_info in state so LLM plans accordingly

import json
import time

def checkpoint_state(state: ExecutionState, checkpoint_id: str, db):
    data = {
        'task': state.task,
        'iteration': state.iteration,
        'completed_steps': state.completed_steps,
        'workspace_files': state.workspace_files,
        'timestamp': time.time()
    }
    db.execute(
        'INSERT INTO agent_checkpoints (id, state) VALUES (?, ?) ON CONFLICT(id) DO UPDATE SET state=excluded.state',
        (checkpoint_id, json.dumps(data))
    )
    print(f'Checkpoint saved at iteration {state.iteration}')

def resume_from_checkpoint(checkpoint_id: str, db) -> ExecutionState | None:
    row = db.execute('SELECT state FROM agent_checkpoints WHERE id = ?', (checkpoint_id,)).fetchone()
    if not row:
        return None
    data = json.loads(row[0])
    state = ExecutionState(task=data['task'])
    state.iteration = data['iteration']
    state.completed_steps = data['completed_steps']
    print(f'Resumed from iteration {state.iteration}')
    return state

import shutil
from pathlib import Path

INTERMEDIATE_PATTERNS = ['*.parquet', 'tmp_*.csv', 'step_*.json', 'debug_*.txt']
FINAL_OUTPUT_PATTERNS = ['report.pdf', 'final_*.csv', 'summary.json']

def cleanup_workspace(workspace_dir: str, keep_final=True):
    workspace = Path(workspace_dir)
    final_outputs = []
    
    if keep_final:
        for pattern in FINAL_OUTPUT_PATTERNS:
            final_outputs.extend(workspace.glob(pattern))
        # Move final outputs to output directory
        output_dir = workspace.parent / 'outputs'
        output_dir.mkdir(exist_ok=True)
        for f in final_outputs:
            shutil.move(str(f), str(output_dir / f.name))
    
    # Delete the workspace
    shutil.rmtree(workspace_dir)
    print(f'Workspace cleaned up. Kept {len(final_outputs)} output files.')
    return [str(f) for f in final_outputs]