Checkpointing og genoptagelse af opgaver

from dataclasses import dataclass, field
from typing import List, Any, Optional

@dataclass
class AgentStep:
    step_index: int
    thought: str
    tool_name: str
    tool_args: dict
    observation: str
    tokens_used: int
    completed_at: str

@dataclass
class AgentCheckpoint:
    task_id: str
    task_description: str
    status: str  # 'running', 'completed', 'failed'
    current_step: int
    completed_steps: List[AgentStep] = field(default_factory=list)
    accumulated_results: dict = field(default_factory=dict)
    total_tokens: int = 0
    final_answer: Optional[str] = None

import json
import redis
from dataclasses import asdict

redis_client = redis.Redis()

def save_checkpoint(checkpoint: AgentCheckpoint):
    key = f'agent:checkpoint:{checkpoint.task_id}'
    data = json.dumps(asdict(checkpoint), default=str)
    redis_client.set(key, data, ex=86400)  # 24h TTL
    # Also append to step log
    log_key = f'agent:log:{checkpoint.task_id}'
    if checkpoint.completed_steps:
        last = checkpoint.completed_steps[-1]
        redis_client.rpush(log_key, json.dumps(asdict(last), default=str))

def load_checkpoint(task_id: str) -> AgentCheckpoint | None:
    key = f'agent:checkpoint:{task_id}'
    data = redis_client.get(key)
    if data:
        return AgentCheckpoint(**json.loads(data))
    return None

async def resume_or_start(task_id: str, task_description: str) -> str:
    checkpoint = load_checkpoint(task_id)
    if checkpoint and checkpoint.status == 'running':
        print(f'Resuming task {task_id} from step {checkpoint.current_step}')
        messages = rebuild_history(checkpoint)
        start_step = checkpoint.current_step
    else:
        print(f'Starting new task {task_id}')
        checkpoint = AgentCheckpoint(task_id=task_id, task_description=task_description, status='running', current_step=0)
        messages = [{'role': 'user', 'content': task_description}]
        start_step = 0
        save_checkpoint(checkpoint)
    return await run_agent_from(checkpoint, messages, start_step)

def rebuild_history(checkpoint: AgentCheckpoint) -> list:
    messages = [{'role': 'user', 'content': checkpoint.task_description}]
    for step in checkpoint.completed_steps:
        # Reconstruct the agent's reasoning message
        messages.append({
            'role': 'assistant',
            'content': f'Thought: {step.thought}\nAction: {step.tool_name}({step.tool_args})'
        })
        # Reconstruct the tool observation
        messages.append({
            'role': 'user',
            'content': f'Observation: {step.observation}'
        })
    return messages

async def idempotent_write_file(content: str, path: str, task_id: str, step: int) -> str:
    dedup_key = f'{task_id}:step_{step}:write:{path}'
    if redis_client.exists(dedup_key):
        return f'File {path} already written (dedup key present)'
    with open(path, 'w') as f:
        f.write(content)
    redis_client.set(dedup_key, '1', ex=3600)
    return f'Successfully wrote {len(content)} chars to {path}'

from datetime import datetime, timedelta

RETENTION = {
    'completed': timedelta(hours=24),
    'failed': timedelta(days=7),
    'running': None  # never auto-delete
}

def cleanup_expired_checkpoints():
    now = datetime.utcnow()
    for key in redis_client.scan_iter('agent:checkpoint:*'):
        data = json.loads(redis_client.get(key))
        status = data.get('status', 'running')
        retention = RETENTION.get(status)
        if retention is None:
            continue
        started = datetime.fromisoformat(data.get('started_at', str(now)))
        if now - started > retention:
            redis_client.delete(key)

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import StateGraph

# Create graph with persistent checkpointer
checkpointer = SqliteSaver.from_conn_string('/tmp/agent_state.db')
graph = StateGraph(AgentState)
graph.add_node('reason', reason_node)
graph.add_node('act', act_node)
# ... add edges ...
app = graph.compile(checkpointer=checkpointer)

# Run with thread_id - LangGraph auto-checkpoints
config = {'configurable': {'thread_id': 'task_abc123'}}
result = await app.ainvoke({'task': 'Research climate change'}, config)

# Resume same thread - LangGraph loads from checkpoint
result = await app.ainvoke({'task': 'Continue'}, config)

async def parallel_with_checkpoints(parent_id: str, subtasks: list) -> list:
    results = []
    for i, subtask in enumerate(subtasks):
        sub_id = f'{parent_id}:sub_{i}'
        # Check if subtask already completed
        existing = load_checkpoint(sub_id)
        if existing and existing.status == 'completed':
            print(f'Sub-task {i} already done, using cached result')
            results.append(existing.final_answer)
        else:
            result = await run_agent(sub_id, subtask)
            results.append(result)
    return results

import pytest

@pytest.mark.asyncio
async def test_resumption_from_step_3():
    task_id = 'test_resume_001'
    # Run until step 3, then crash
    with pytest.raises(SimulatedCrash):
        await run_agent_crashing_at(task_id, 'Research AI trends', crash_at_step=3)

    checkpoint = load_checkpoint(task_id)
    assert checkpoint.current_step == 3
    assert len(checkpoint.completed_steps) == 3

    # Resume and complete
    result = await resume_or_start(task_id, 'Research AI trends')
    assert result is not None
    # Verify only steps 4+ were re-executed
    assert checkpoint_step_was_not_replayed(task_id, step=0)

def load_and_migrate_checkpoint(task_id: str) -> AgentCheckpoint:
    raw = json.loads(redis_client.get(f'agent:checkpoint:{task_id}'))
    version = raw.get('checkpoint_version', '1.0')
    if version == '1.0':
        # Migrate: add new fields added in v2.0
        raw['checkpoint_version'] = '2.0'
        raw['accumulated_results'] = raw.get('accumulated_results', {})
        raw['total_tokens'] = raw.get('total_tokens', 0)
    return AgentCheckpoint(**raw)