build_agent.py

from hackbot.src.pocgen.runtime_environment import RuntimeEnvironment, ExecutionResults
from typing import Sequence
from hackbot.utils.dspy.tools import enrich_tools, dspy_tool, async_execute_command_tool, async_common_tools
from common.config import global_config

import dspy
from hackbot.utils.dspy.dspy_inference import DSPYInference


@dspy_tool(runtime_env=RuntimeEnvironment, approved_commands=Sequence[str])
async def final_check_before_completion(*, runtime_env: RuntimeEnvironment, approved_commands: Sequence[str]) -> str:
    """Always run this tool before finishing. This tool will tell you whether you are done or not."""
    print("Trying to finish up...")

    result1 = await async_execute_command_tool(
        "forge build", runtime_env=runtime_env, approved_commands=approved_commands
    )
    if result1.startswith("Error"):
        return (
            "Failure: Forge build failed! Please continue getting the repo to build. Here is your output: "
            + result1
        )

    print("Forge build passed!")

    return "Success: Forge build passed!"


class RunCommands(dspy.Signature):
    """
    Run commands in the given directory.

    IMPORTANT: Your response MUST strictly follow the format below. Each field, including the headers, must be on a new line.

    Example:
    [[ ## next_thought ## ]]
    I need to see what files are in the current directory.
    [[ ## next_tool_name ## ]]
    async_list_files_tool
    [[ ## next_tool_args ## ]]
    {"relative_path": "."}
    """

    high_level_instructions: str = dspy.InputField()  # type: ignore # noqa
    forge_build_runs: bool = dspy.OutputField()  # type: ignore # noqa
    forge_test_runs: bool = dspy.OutputField()  # type: ignore # noqa
    shell_commands_ran: list[str] = dspy.OutputField()  # type: ignore # noqa


async def compile_codebase(runtime_env: RuntimeEnvironment) -> ExecutionResults:
    """Compile the codebase using a dspy-based LLM agent."""

    approved_commands = global_config.build_agent.approved_commands

    await runtime_env.execute_command(["git", "reset", "--hard"])
    await runtime_env.execute_command(["git", "clean", "-ffdx"])

    # All dspy inference and tool execution must happen within this block
    with enrich_tools(runtime_env=runtime_env, approved_commands=approved_commands):
        inf_module = DSPYInference(
            pred_signature=RunCommands,
            tools=[
                *async_common_tools,
                final_check_before_completion,
            ],
            max_iters=50,
        )
        
        result = await inf_module.run(
            high_level_instructions=f"""
                You are a professional smart contract security engineer.
                Your current directory is inside of a git repository of a Foundry Forge Solidity project.
                Your job is to build the project and run the tests.
                You are allowed to read through files and folders in the project to understand it.
                You can additionally only run the following commands at the root of the project:
                {', '.join(approved_commands)}
                You will not combine commands with && or ||.
                You will always try to use the non-interactive or json versions of the commands. This is because the output is passed back to you in the same format.
                Sometimes repos fail to give instructions on how to set up dependencies. In this case you should figure out how to pull them from one of the supplied commands.
                For example, if you find out that forge-std doesn't exist, you can pull it in via `forge install foundry-rs/forge-std`.
                Please be mindful of required versions and stick to them.
                You should run a command with `--help` if you are unsure about its capabilities and don't have any ideas of how to proceed.
                When supplying any relative path to a tool you will never use `..`. This is forbidden and will cause the tool to fail.
                You are not allowed to move freely using `cd`.
                You can additionally modify any file in the project that you have already read.
                Your goal is to run `forge build` successfully and to run `forge test` (which must complete but may have failing tests so it may not necessarily "succeed").
                Keep in mind that the `forge soldeer` subcommand exists and is the dependency manager for some projects. You should expect to see a `soldeer.lock` in these cases.
                You will not give up until these commands run successfully.
                You will never run any commands that open ports or a shell.
                You only have one shot at this so make sure to do it right.
                You should start by finding and opening the readme. Then, you should see what other files and folders are in the project.
                Your response will always finish with a call to `final_check_before_completion` that returns a success.
                You will not execute any further commands after you have verified that `forge build` and `forge test` work.
            """
        )
        print(result.shell_commands_ran)
    
    if not result.forge_build_runs:
        return ExecutionResults(
            stdout=",".join(result.shell_commands_ran),
            stderr="Build agent failed to build the repository.",
            exit_code=1,
            cmd_full="build_agent"
        )
    
    return ExecutionResults(
        stdout=",".join(result.shell_commands_ran),
        stderr="",
        exit_code=0,
        cmd_full="build_agent"
    )