jester1177/cloud-native-debug-env

1"""Heuristic baseline runner for the /baseline endpoint.2 3Applies expected_fixes directly to verify the environment + grader work e2e.4"""5 6 7from typing import List, Optional8 9from server.environment import CICDDebugEnvironment10from server.graders import run_grader11from server.models import Action, ActionType, FileEdit, GraderResult12from server.tasks.task_registry import TASK_REGISTRY13 14 15def _heuristic_episode(env: CICDDebugEnvironment, task_id: str, scenario_id: Optional[str] = None) -> GraderResult:16    """Run one episode using a heuristic that applies expected fixes."""17    obs = env.reset(task_id=task_id, scenario_id=scenario_id)18 19    # Apply each expected fix as an edit_file action20    for fix in env.expected_fixes:21        if env.done:22            break23        file_path = fix["file"]24        if file_path not in env.current_files:25            continue26 27        current_content = env.current_files[file_path].content28 29        if fix["type"] == "contains":30            # Need to ensure expected string is present31            if fix["expected"] not in current_content:32                # Try to find the broken line using hint33                hint_text = fix.get("hint", "")34                # Use edit_file with old/new content based on the fix35                # We look at original files to find what changed36                original_content = env.original_files.get(file_path)37                if original_content:38                    lines = current_content.split("\n")39                    expected = fix["expected"]40                    line_num = fix.get("line")41 42                    if line_num and 1 <= line_num <= len(lines):43                        old_line = lines[line_num - 1]44                        action = Action(45                            action_type=ActionType.REPLACE_LINE,46                            edits=[FileEdit(47                                file_path=file_path,48                                line_number=line_num,49                                new_content=expected,50                            )],51                        )52                    else:53                        # Find the line that's closest to expected but wrong54                        best_line = None55                        best_idx = None56                        for i, line in enumerate(lines):57                            stripped = line.strip()58                            exp_stripped = expected.strip()59                            # Check if this line is a broken version of expected60                            if (stripped and exp_stripped and61                                    len(set(stripped) & set(exp_stripped)) > len(exp_stripped) * 0.3):62                                if best_line is None:63                                    best_line = line64                                    best_idx = i65 66                        if best_line is not None:67                            action = Action(68                                action_type=ActionType.EDIT_FILE,69                                edits=[FileEdit(70                                    file_path=file_path,71                                    old_content=best_line,72                                    new_content=expected,73                                )],74                            )75                        else:76                            # Append the expected content77                            action = Action(78                                action_type=ActionType.ADD_LINE,79                                edits=[FileEdit(80                                    file_path=file_path,81                                    new_content=expected,82                                )],83                            )84                    env.step(action)85 86        elif fix["type"] == "not_contains":87            # Need to ensure expected string is NOT present88            if fix["expected"] in current_content:89                action = Action(90                    action_type=ActionType.DELETE_BLOCK,91                    edits=[FileEdit(92                        file_path=file_path,93                        old_content=fix["expected"],94                    )],95                )96                env.step(action)97 98        elif fix["type"] == "line_equals":99            line_num = int(fix.get("line", 0))100            if line_num >= 1:101                action = Action(102                    action_type=ActionType.REPLACE_LINE,103                    edits=[FileEdit(104                        file_path=file_path,105                        line_number=line_num,106                        new_content=str(fix["expected"]),107                    )],108                )109                env.step(action)110 111    # Submit if not already done112    if not env.done:113        env.step(Action(action_type=ActionType.SUBMIT))114 115    return run_grader(task_id, env.trajectory)116 117 118def run_baseline_episodes(task_id: Optional[str] = None, num_episodes: int = 1) -> List[GraderResult]:119    """Run baseline episodes across tasks.120 121    Args:122        task_id: Specific task to run, or None for all tasks.123        num_episodes: Number of episodes per task.124 125    Returns:126        List of GraderResult for each episode.127    """128    task_ids: List[str]129    if task_id:130        if task_id not in TASK_REGISTRY:131            raise ValueError(f"Unknown task: {task_id}")132        task_ids = [task_id]133    else:134        task_ids = list(TASK_REGISTRY.keys())135 136    results: List[GraderResult] = []137    for tid in task_ids:138        task_cls = TASK_REGISTRY[tid]139        scenarios = task_cls.SCENARIOS140        episodes_run = 0141        for scenario in scenarios:142            if episodes_run >= num_episodes:143                break144            env = CICDDebugEnvironment()145            result = _heuristic_episode(env, tid, scenario["id"])146            results.append(result)147            episodes_run += 1148 149    return results150