OpenHands/evaluation/biocoder/run_infer.py

import asyncio
import functools
import json
import os
import tempfile
from typing import Any

import pandas as pd
from datasets import load_dataset

from evaluation.biocoder.utils import BiocoderData
from evaluation.utils.shared import (
    EvalMetadata,
    EvalOutput,
    codeact_user_response,
    make_metadata,
    prepare_dataset,
    reset_logger_for_multiprocessing,
    run_evaluation,
)
from openhands.controller.state.state import State
from openhands.core.config import (
    AppConfig,
    SandboxConfig,
    get_llm_config_arg,
    parse_arguments,
)
from openhands.core.logger import openhands_logger as logger
from openhands.core.main import create_runtime, run_controller
from openhands.events.action import CmdRunAction
from openhands.events.observation import CmdOutputObservation
from openhands.runtime.runtime import Runtime

AGENT_CLS_TO_FAKE_USER_RESPONSE_FN = {
    'CodeActAgent': functools.partial(
        codeact_user_response, encapsulate_solution=True, try_parse=None
    ),
}

AGENT_CLS_TO_INST_SUFFIX = {
    'CodeActAgent': 'When you think you have fixed the issue through code changes, please run the following command: <execute_bash> exit </execute_bash>.\n'
}

FILE_EXT_MAP = {
    'python': 'py',
    'java': 'java',
    'c': 'c',
    'cpp': 'cpp',
    'javascript': 'js',
    'typescript': 'ts',
}


def get_config(
    metadata: EvalMetadata,
) -> AppConfig:
    BIOCODER_BENCH_CONTAINER_IMAGE = 'public.ecr.aws/i5g0m1f6/eval_biocoder:v1.0'

    config = AppConfig(
        default_agent=metadata.agent_class,
        run_as_openhands=False,
        runtime='eventstream',
        max_iterations=metadata.max_iterations,
        sandbox=SandboxConfig(
            base_container_image=BIOCODER_BENCH_CONTAINER_IMAGE,
            enable_auto_lint=True,
            use_host_network=False,
        ),
        # do not mount workspace
        workspace_base=None,
        workspace_mount_path=None,
    )
    config.set_llm_config(metadata.llm_config)
    return config


async def initialize_runtime(
    runtime: Runtime,
    instance: BiocoderData,  # this argument is not required
):
    """Initialize the runtime for the agent.

    This function is called before the runtime is used to run the agent.
    """
    logger.info(f"{'-' * 50} BEGIN Runtime Initialization Fn {'-' * 50}")
    obs: CmdOutputObservation

    file_ext = FILE_EXT_MAP[instance.language.lower()]

    action = CmdRunAction(command='mkdir -p /workspace && mkdir -p /testing_files')
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    assert obs.exit_code == 0

    with tempfile.TemporaryDirectory() as tmpdir:
        context_path = os.path.join(tmpdir, 'context.' + file_ext)
        with open(context_path, 'w') as f:
            f.write(instance.contextCode)
        await runtime.copy_to(context_path, '/testing_files')

        golden_path = os.path.join(tmpdir, 'golden.' + file_ext)
        with open(golden_path, 'w') as f:
            f.write(instance.goldenCode)
        await runtime.copy_to(golden_path, '/testing_files')

        testcase_json = {
            'test_case_id': instance.test_case_id,
            'num_cases': 1000,
            'language': instance.language.lower(),
        }
        testcase_path = os.path.join(tmpdir, 'testcase_biocoder.json')
        with open(testcase_path, 'w') as f:
            f.write(json.dumps(testcase_json, indent=4))

        await runtime.copy_to(testcase_path, '/testing_files')

    # setup paths
    remove_code_script = os.path.join(
        os.path.dirname(__file__), 'scripts', 'setup', 'remove_code.py'
    )
    await runtime.copy_to(remove_code_script, '/testing_files')

    action = CmdRunAction(command='cd /workspace')
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    assert obs.exit_code == 0

    # download repository archive
    repository_url = f"https://biocoder.lilbillbiscuit.com/repos/{instance.repository.split('/')[1]}.zip"
    action = CmdRunAction(command='wget -O repo.zip ' + repository_url)
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    assert obs.exit_code == 0, f'Failed to download the repository: {obs.content}'

    # unzip the repository
    action = CmdRunAction(command='unzip -o -q repo.zip && rm repo.zip')
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    assert obs.exit_code == 0, f'Failed to unzip the repository: {obs.content}'

    # chmod 777
    action = CmdRunAction(command='chmod -R 777 /workspace')
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    assert obs.exit_code == 0, f'Failed to chmod the files: {obs.content}'

    # remove code for evaluation instance
    target_filepath = os.path.join(
        '/workspace', instance.repository.split('/')[1], instance.filePath
    )
    line_start = instance.lineStart
    line_end = instance.lineEnd
    language = instance.language.lower()
    action = CmdRunAction(
        command=f'python3 /testing_files/remove_code.py --target_filepath {target_filepath} --line_start {line_start} --line_end {line_end} --language {language}'
    )
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    assert obs.exit_code == 0, f'Failed to remove the code: {obs.content}'

    logger.info(f"{'-' * 50} END Runtime Initialization Fn {'-' * 50}")


async def complete_runtime(
    runtime: Runtime,
    instance: pd.Series,  # this argument is not required, but it is used to get the workspace_dir_name
) -> dict[str, Any]:
    """Complete the runtime for the agent.

    This function is called before the runtime is used to run the agent.
    If you need to do something in the sandbox to get the correctness metric after
    the agent has run, modify this function.
    """
    logger.info(f"{'-' * 50} BEGIN Runtime Completion Fn {'-' * 50}")
    obs: CmdOutputObservation

    test_result = {'result': {}, 'metadata': {}}

    copy_changed_code_script = os.path.join(
        os.path.dirname(__file__), 'scripts', 'setup', 'copy_changed_code.py'
    )
    await runtime.copy_to(copy_changed_code_script, '/testing_files')

    file_ext = FILE_EXT_MAP[instance.language.lower()]
    target_filepath = os.path.join(
        '/workspace', instance.repository.split('/')[1], instance.filePath
    )
    generated_path = os.path.join('/testing_files', 'generated.' + file_ext)

    action = CmdRunAction(
        command=f'python3 /testing_files/copy_changed_code.py --target_filepath {target_filepath} --generated_code_filepath {generated_path} --line_start {instance.lineStart} --include_signature'
    )
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    if obs.exit_code == 0:
        test_result['metadata']['1_copy_change_success'] = True

        action = CmdRunAction(command=f'cat {generated_path}', keep_prompt=False)
        logger.info(action, extra={'msg_type': 'ACTION'})
        obs = await runtime.run_action(action)
        assert obs.exit_code == 0

        code = obs.content
        test_result['metadata']['1_copy_change_code'] = code
    else:
        test_result['metadata']['1_copy_change_success'] = False
        test_result['metadata']['1_copy_change_code'] = None

    action = CmdRunAction(command='cd /testing_files')
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    assert obs.exit_code == 0

    action = CmdRunAction(
        command='/home/openhands/mambaforge/bin/mamba run -n test python3 /testing/start_test_openhands.py'
    )
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    logger.info(obs, extra={'msg_type': 'OBSERVATION'})
    assert obs.exit_code == 0

    action = CmdRunAction(
        command='cat /testing_files/results_biocoder.json', keep_prompt=False
    )
    logger.info(action, extra={'msg_type': 'ACTION'})
    obs = await runtime.run_action(action)
    if obs.exit_code == 0:
        test_result['metadata']['2_run_test_success'] = True
        test_result['metadata']['2_run_test_result'] = str(obs.content)
        json_obj = json.loads(obs.content)
        test_result['result'] = json_obj['result']
    else:
        test_result['metadata']['2_run_test_success'] = False
        test_result['metadata']['2_run_test_result'] = str(obs.content)

    logger.info(f"{'-' * 50} END Runtime Completion Fn {'-' * 50}")
    return test_result


async def process_instance(
    instance: pd.Series,
    metadata: EvalMetadata,
    reset_logger: bool = True,
) -> EvalOutput:
    config = get_config(metadata)
    instance = BiocoderData(**instance)
    print(instance)
    instance_id = f'{instance.repository}__{instance.instance_id[:10]}'

    # Setup the logger properly, so you can run multi-processing to parallelize the evaluation
    if reset_logger:
        log_dir = os.path.join(metadata.eval_output_dir, 'infer_logs')
        reset_logger_for_multiprocessing(logger, instance_id, log_dir)
    else:
        logger.info(f'Starting evaluation for instance {instance_id}.')

    # Prepare instruction
    instruction = (
        f'Please complete the function "{instance.signature}" in the file /workspace/{instance.repository.split("/")[1]}/{instance.filePath}.\n'
        f'The environment has been set up for you to start working. You may assume all necessary tools are installed.\n'
        f'To complete the task, you must directly modify the file and fill in the function, keeping in mind that the function signature is on line {instance.lineStart-1}\n\n'
        f'The function should do the following:\n'
        f'{instance.promptSummaryOnly}\n\n'
    )

    instruction += (
        'IMPORTANT: You should ONLY interact with the environment provided to you AND NEVER ASK FOR HUMAN HELP.\n'
        'You should NOT modify any other files other than the file intended. This means that you should NOT write any test cases.\n'
        'You may need context from other files in the repository to complete this task.'
        'Do NOT add any import statements or change anything else other than the writing the function body.\n'
        'You do not need to run the code to check if it works. \n'
        'Make sure to include proper formatting in Java and Python, including correct braces and/or indentation.\n'
    )
    # NOTE: You can actually set slightly different instruction for different agents
    instruction += AGENT_CLS_TO_INST_SUFFIX[metadata.agent_class]

    # use a session id for concurrent evaluation
    sid = instance.instance_id.replace('/', '__')

    runtime = await create_runtime(config, sid=sid)

    await initialize_runtime(runtime, instance)

    # Here's how you can run the agent (similar to the `main` function) and get the final task state
    state: State | None = await run_controller(
        config=config,
        task_str=instruction,
        runtime=runtime,
        fake_user_response_fn=AGENT_CLS_TO_FAKE_USER_RESPONSE_FN[metadata.agent_class],
    )

    if state is None:
        raise ValueError('State should not be None.')

    test_result = await complete_runtime(runtime, instance)
    metrics = state.metrics.get() if state.metrics else None
    # history is now available as a stream of events, rather than list of pairs of (Action, Observation)
    # for compatibility with the existing output format, we can remake the pairs here
    # remove when it becomes unnecessary
    histories = state.history.compatibility_for_eval_history_pairs()

    test_result['generated'] = test_result['metadata']['1_copy_change_code']

    # Save the output
    output = EvalOutput(
        instance_id=instance.instance_id,
        instance=instance.to_dict(),
        instruction=instruction,
        metadata=metadata,
        history=histories,
        metrics=metrics,
        error=state.last_error if state and state.last_error else None,
        test_result=test_result,
    )
    return output


if __name__ == '__main__':
    args = parse_arguments()

    dataset = load_dataset('lilbillbiscuit/biocoder_public')
    biocoder_tests = dataset['train'].to_pandas()
    biocoder_tests['instance_id'] = biocoder_tests['test_case_id']

    llm_config = None
    if args.llm_config:
        llm_config = get_llm_config_arg(args.llm_config)

    if llm_config is None:
        raise ValueError(f'Could not find LLM config: --llm_config {args.llm_config}')

    metadata = make_metadata(
        llm_config,
        'biocoder',
        args.agent_cls,
        args.max_iterations,
        args.eval_note,
        args.eval_output_dir,
    )
    output_file = os.path.join(metadata.eval_output_dir, 'output.jsonl')
    instances = prepare_dataset(biocoder_tests, output_file, args.eval_n_limit)

    asyncio.run(
        run_evaluation(
            instances, metadata, output_file, args.eval_num_workers, process_instance
        )
    )