ReWOO agents in Ruby

actor.rb

# https://arxiv.org/pdf/2305.18323.pdf
# ReWOO is a silly name for "do the reasoning before you make observations". The idea
# is to first generate a plan, using placeholder values for the results of actual action runs,
# then run the actions in one go, then make a second LLM call to come up with a conclusion.
#
# It's much quicker than a ReAct agent, since it only requires two LLM calls, but if it doesn't get
# the plan right the first time (e.g. it tries to grep for the wrong thing) it's game over
require_relative './tools'

module Rewoo
    class ActorError < StandardError
    end

    class Actor
        attr_reader :question, :tools, :client, :prompt

        def initialize(client, question)
            @question = question
            @client = client
        end

        def self.run_for_objective(client, question)
            actor = new(client, question)
            plan = actor.make_plan
            puts "Plan:\n#{plan}"
            variables = actor.run_tools(plan)
            answer = actor.solve(question, plan, variables)
            puts "Solved:\n#{answer}"
            [answer, plan]
        rescue ActorError => e
            ["Error: #{e.message}", plan]
        end

        def make_plan
            fewshot_examples = """
        Question: How is the Widget component implemented in this repo?
        Plan: Find which file defines the Widget component
        #E1 = ls-files[Widget]
        Plan: Find which files could reference the Widget component
        #E2 = grep[Widget]
        Plan: Identify the file most likely to define the Widget component.
        #E3 = LLM[Of these files, which file defines the Widget component? Given context: #E1]
        Plan: Read the source of the Widget component.
        #E4 = read-file[#E3]
        Plan: Ask the LLM to summarize the Widget component source, given the context of how the Widget component is referenced.
        #E4 = LLM[How is the Widget component implemented in #E4? Given how it's referenced: #E2]

        Question: Which files are relevant to marking a user as spammy in this repo?
        Plan: Search the codebase for the spammy keyword.
        #E1 = grep[spammy]
        Plan: Identify the files that seem to mark a user as spammy.
        #E2 = LLM[Which file is the most relevant to marking a user as spammy? Given context: #E1]
    """
            prompt = """
    You are a software engineering AI with access to a codebase. For the following task, make plans that can solve the problem step-by-step.
    For each plan, indicate which external tool should be used to retrieve evidence.
    You can store the evidence into a variable #E that can be called by later tools. (Plan, #E1, Plan, #E2, Plan, ...)

    Your available tools are:
    #{tools.tool_descriptions}

    Some examples of previous plans: #{fewshot_examples}

    Begin! Describe your plans with rich details. Each Plan should be followed by only one #E.

    Question: #{question}
            """

            client.chat(prompt)
        end

        def run_tools(plan)
            tasks = plan.lines.filter { |l| l.start_with?("#E") }
            variables = {}
            tasks.each do |t|
                puts "Processing task line: #{t}"
                variable_name, action = t.split("=", 2)
                variable_name = variable_name.strip
                action = action.strip

                tool_name, input = action.split("[", 2)
                input = input.delete_suffix("]")

                input = input.gsub(/#E\d+/) do |match|
                    variables[match]
                end


                tool = tools.tools.find { |t| t.name == tool_name }

                begin
                    variables[variable_name] = tool.method.call(input)
                rescue StandardError => e
                    raise ActorError, e.inspect
                end
            end

            variables
        end

        def solve(question, plan, tool_output_variables)

            evidence = plan.lines.map do |l|
                match = l.match(/#E\d+/)
                next l unless match
                val = tool_output_variables[match[0]] || "Variable missing!"
                "\n#{match}:\n#{val[0..1000]}"
            end.join

            prompt = """
Solve the following task or problem:

#{question}

To assist you, we provide some plans and corresponding evidences that might be helpful.
Notice that some of these information contain noise so you should trust them with caution.

#{evidence}

Now begin to solve the task or problem. Respond with the answer directly with no extra words.
            """

            puts "Passing prompt to solver:\n#{prompt}"

            client.chat(prompt)


        end

        private
        
        def tools
            @tools ||= Rewoo::Tools.new(client)
        end
    end
end

main.rb

require_relative './actor'
Rewoo::Actor.run_for_objective(proxy_api_client, "How does auth work in this codebase?")

tools.rb

# Note that the ReWOO pattern can't recover from errors, so we don't want to raise
# useful error messages for the model. Instead, we just explode.

module Rewoo
    class Toolkit
        attr_accessor :llm_client

        Tool = Struct.new(:name, :description, :method)

        def initialize(client)
           @llm_client = client
        end
    
        def tool_descriptions
            tools.map do |t|
                "#{t.name}: #{t.description}"
            end.join("\n")
        end
    
        def tool_names
            tools.map { |t| t.name }.join(", ")
        end

        def tools
            [
                Tool.new(
                    "LLM",
                    """Ask a LLM a question.
                    Input should be the question
                    Result will be the answer.""",
                    ->(input) {
                        llm_client.chat(input[0..15000] + "\nRespond with the answer directly with no extra words.")
                    }
                ),
                Tool.new(
                    "ls-files",
                    """Search a git repository with the git ls-files utility.
                    Input should be the arguments passed to git ls-files (e.g. *filename.ext)
                    Result will be a list of files, or an error message.""",
                    ->(input) {
                        input = input.strip.delete_prefix("git ls-files ")

                        puts "Running 'git --no-pager ls-files #{input}'..."

                        stdout, stderr, status = Open3.capture3("git --no-pager ls-files #{input}")
                        if !stderr.empty?
                            raise Rewoo::ActorError, "ls-files, #{stderr}"
                        end

                        if stdout.empty?
                            raise Rewoo::ActorError, "ls-files, no results found"
                        end

                        stdout
                    }
                ),
                Tool.new(
                    "grep",
                    """Search a git repository with the git grep utility.
                    Input should be the arguments passed to git grep
                    Result will be a data structure with this format: `filename: line\nfilename2: line2`, or an error message.""",
                    ->(input) {
                        input = input.strip.delete_prefix("git grep ")
                        puts "Running 'git --no-pager grep '#{input}''..."
                        stdout, stderr, status = Open3.capture3("git --no-pager grep '#{input}'")
                        if !stderr.empty?
                            raise Rewoo::ActorError, "grep, #{stderr}"
                        end
    
                        if stdout.empty?
                            raise Rewoo::ActorError, "grep, #{stderr}"
                        end
    
                        stdout = Utils.to_utf8(stdout)
    
                        stdout.split("\n").filter do |line|
                            line.length < 400
                        end.join("\n")
                    }
                ),
                Tool.new(
                    "read-file",
                    """Read the first 200 lines of a file on disk.
                    Input should be the path to the file
                    Result will be 200 lines from the file, or an exception if one was raised.""",
                    ->(input) {
                        # let's try and find the first file in the input, which might be noisy LLM output
                        line = input.lines.find { |l| l.include?("/") && l.include?(".") }
                        puts "Could not find file to read in :#{input}".red unless line
                        line = line.strip
                        path = line.split(" ").find{ |w| w.include?("/") && w.include?(".") }
                        path = path.delete_suffix(":")
                            .delete_suffix(",")
                            .delete_suffix(".")

                        puts "read-file: reading file at #{path}".green

                        begin
                            file = File.read(File.expand_path(path))
                            file.lines[0..200].join
                        rescue Exception => e
                            raise Rewoo::ActorError, "read-file, #{e.inspect}"
                        end
                    }
                ),
            ]
        end
    end
end