FIX: support multiple tool calls (#502)

* FIX: support multiple tool calls

Prior to this change we had a hard limit of 1 tool call per llm
round trip. This meant you could not google multiple things at
once or perform searches across two tools.

Also:

- Hint when Google stops working
- Log topic_id / post_id when performing completions

* Also track id for title

lib/ai_bot/bot.rb

@@ -7,6 +7,7 @@ class Bot
 
       BOT_NOT_FOUND = Class.new(StandardError)
       MAX_COMPLETIONS = 5
+      MAX_TOOLS = 5
 
       def self.as(bot_user, persona: DiscourseAi::AiBot::Personas::General.new, model: nil)
         new(bot_user, persona, model)
@@ -21,14 +22,19 @@ def initialize(bot_user, persona, model = nil)
       attr_reader :bot_user
       attr_accessor :persona
 
-      def get_updated_title(conversation_context, post_user)
+      def get_updated_title(conversation_context, post)
         system_insts = <<~TEXT.strip
         You are titlebot. Given a topic, you will figure out a title.
         You will never respond with anything but 7 word topic title.
         TEXT
 
         title_prompt =
-          DiscourseAi::Completions::Prompt.new(system_insts, messages: conversation_context)
+          DiscourseAi::Completions::Prompt.new(
+            system_insts,
+            messages: conversation_context,
+            topic_id: post.topic_id,
+            post_id: post.id,
+          )
 
         title_prompt.push(
           type: :user,
@@ -38,7 +44,7 @@ def get_updated_title(conversation_context, post_user)
 
         DiscourseAi::Completions::Llm
           .proxy(model)
-          .generate(title_prompt, user: post_user)
+          .generate(title_prompt, user: post.user)
           .strip
           .split("\n")
           .last
@@ -64,37 +70,14 @@ def reply(context, &update_blk)
 
           result =
             llm.generate(prompt, **llm_kwargs) do |partial, cancel|
-              if (tool = persona.find_tool(partial))
+              tools = persona.find_tools(partial)
+
+              if (tools.present?)
                 tool_found = true
-                ongoing_chain = tool.chain_next_response?
-                tool_call_id = tool.tool_call_id
-                invocation_result_json = invoke_tool(tool, llm, cancel, &update_blk).to_json
-
-                tool_call_message = {
-                  type: :tool_call,
-                  id: tool_call_id,
-                  content: { name: tool.name, arguments: tool.parameters }.to_json,
-                }
-
-                tool_message = { type: :tool, id: tool_call_id, content: invocation_result_json }
-
-                if tool.standalone?
-                  standalone_context =
-                    context.dup.merge(
-                      conversation_context: [
-                        context[:conversation_context].last,
-                        tool_call_message,
-                        tool_message,
-                      ],
-                    )
-                  prompt = persona.craft_prompt(standalone_context)
-                else
-                  prompt.push(**tool_call_message)
-                  prompt.push(**tool_message)
+                tools[0..MAX_TOOLS].each do |tool|
+                  ongoing_chain &&= tool.chain_next_response?
+                  process_tool(tool, raw_context, llm, cancel, update_blk, prompt)
                 end
-
-                raw_context << [tool_call_message[:content], tool_call_id, "tool_call"]
-                raw_context << [invocation_result_json, tool_call_id, "tool"]
               else
                 update_blk.call(partial, cancel, nil)
               end
@@ -115,6 +98,37 @@ def reply(context, &update_blk)
 
       private
 
+      def process_tool(tool, raw_context, llm, cancel, update_blk, prompt)
+        tool_call_id = tool.tool_call_id
+        invocation_result_json = invoke_tool(tool, llm, cancel, &update_blk).to_json
+
+        tool_call_message = {
+          type: :tool_call,
+          id: tool_call_id,
+          content: { name: tool.name, arguments: tool.parameters }.to_json,
+        }
+
+        tool_message = { type: :tool, id: tool_call_id, content: invocation_result_json }
+
+        if tool.standalone?
+          standalone_context =
+            context.dup.merge(
+              conversation_context: [
+                context[:conversation_context].last,
+                tool_call_message,
+                tool_message,
+              ],
+            )
+          prompt = persona.craft_prompt(standalone_context)
+        else
+          prompt.push(**tool_call_message)
+          prompt.push(**tool_message)
+        end
+
+        raw_context << [tool_call_message[:content], tool_call_id, "tool_call"]
+        raw_context << [invocation_result_json, tool_call_id, "tool"]
+      end
+
       def invoke_tool(tool, llm, cancel, &update_blk)
         update_blk.call("", cancel, build_placeholder(tool.summary, ""))
 

lib/ai_bot/personas/persona.rb

@@ -117,15 +117,25 @@ def craft_prompt(context)
             #{available_tools.map(&:custom_system_message).compact_blank.join("\n")}
           TEXT
               messages: context[:conversation_context].to_a,
+              topic_id: context[:topic_id],
+              post_id: context[:post_id],
             )
 
           prompt.tools = available_tools.map(&:signature) if available_tools
 
           prompt
         end
 
-        def find_tool(partial)
+        def find_tools(partial)
+          return [] if !partial.include?("</invoke>")
+
           parsed_function = Nokogiri::HTML5.fragment(partial)
+          parsed_function.css("invoke").map { |fragment| find_tool(fragment) }.compact
+        end
+
+        protected
+
+        def find_tool(parsed_function)
           function_id = parsed_function.at("tool_id")&.text
           function_name = parsed_function.at("tool_name")&.text
           return false if function_name.nil?

lib/ai_bot/playground.rb

@@ -156,7 +156,7 @@ def title_playground(post)
         context = conversation_context(post)
 
         bot
-          .get_updated_title(context, post.user)
+          .get_updated_title(context, post)
           .tap do |new_title|
             PostRevisor.new(post.topic.first_post, post.topic).revise!(
               bot.bot_user,
@@ -182,6 +182,8 @@ def reply_to(post)
           participants: post.topic.allowed_users.map(&:username).join(", "),
           conversation_context: conversation_context(post),
           user: post.user,
+          post_id: post.id,
+          topic_id: post.topic_id,
         }
 
         reply_user = bot.bot_user

lib/ai_bot/tools/google.rb

@@ -37,6 +37,7 @@ def invoke(bot_user, llm)
             URI(
               "https://www.googleapis.com/customsearch/v1?key=#{api_key}&cx=#{cx}&q=#{escaped_query}&num=10",
             )
+
           body = Net::HTTP.get(uri)
 
           parse_search_json(body, escaped_query, llm)
@@ -65,6 +66,19 @@ def minimize_field(result, field, llm, max_tokens: 100)
 
         def parse_search_json(json_data, escaped_query, llm)
           parsed = JSON.parse(json_data)
+          error_code = parsed.dig("error", "code")
+          if error_code == 429
+            Rails.logger.warn(
+              "Google Custom Search is Rate Limited, no search can be performed at the moment. #{json_data[0..1000]}",
+            )
+            return(
+              "Google Custom Search is Rate Limited, no search can be performed at the moment. Let the user know there is a problem."
+            )
+          elsif error_code
+            Rails.logger.warn("Google Custom Search returned an error. #{json_data[0..1000]}")
+            return "Google Custom Search returned an error. Let the user know there is a problem."
+          end
+
           results = parsed["items"]
 
           @results_count = parsed.dig("searchInformation", "totalResults").to_i

lib/completions/dialects/dialect.rb

@@ -106,9 +106,11 @@ def max_prompt_tokens
           raise NotImplemented
         end
 
+        attr_reader :prompt
+
         private
 
-        attr_reader :prompt, :model_name, :opts
+        attr_reader :model_name, :opts
 
         def trim_messages(messages)
           prompt_limit = max_prompt_tokens

lib/completions/endpoints/base.rb

@@ -100,6 +100,8 @@ def perform_completion!(dialect, user, model_params = {})
                   user_id: user&.id,
                   raw_request_payload: request_body,
                   request_tokens: prompt_size(prompt),
+                  topic_id: dialect.prompt.topic_id,
+                  post_id: dialect.prompt.post_id,
                 )
 
               if !@streaming_mode
@@ -273,16 +275,22 @@ def extract_prompt_for_tokenizer(prompt)
         def build_buffer
           Nokogiri::HTML5.fragment(<<~TEXT)
           <function_calls>
-          <invoke>
-          <tool_name></tool_name>
-          <tool_id></tool_id>
-          <parameters>
-          </parameters>
-          </invoke>
+          #{noop_function_call_text}
           </function_calls>
           TEXT
         end
 
+        def noop_function_call_text
+          (<<~TEXT).strip
+            <invoke>
+            <tool_name></tool_name>
+            <tool_id></tool_id>
+            <parameters>
+            </parameters>
+            </invoke>
+          TEXT
+        end
+
         def has_tool?(response)
           response.include?("<function")
         end

lib/completions/endpoints/open_ai.rb

@@ -167,8 +167,26 @@ def add_to_buffer(function_buffer, _response_data, partial)
           @args_buffer ||= +""
 
           f_name = partial.dig(:function, :name)
-          function_buffer.at("tool_name").content = f_name if f_name
-          function_buffer.at("tool_id").content = partial[:id] if partial[:id]
+
+          @current_function ||= function_buffer.at("invoke")
+
+          if f_name
+            current_name = function_buffer.at("tool_name").content
+
+            if current_name.blank?
+              # first call
+            else
+              # we have a previous function, so we need to add a noop
+              @args_buffer = +""
+              @current_function =
+                function_buffer.at("function_calls").add_child(
+                  Nokogiri::HTML5::DocumentFragment.parse(noop_function_call_text + "\n"),
+                )
+            end
+          end
+
+          @current_function.at("tool_name").content = f_name if f_name
+          @current_function.at("tool_id").content = partial[:id] if partial[:id]
 
           args = partial.dig(:function, :arguments)
 
@@ -185,7 +203,7 @@ def add_to_buffer(function_buffer, _response_data, partial)
                 end
               argument_fragments << "\n"
 
-              function_buffer.at("parameters").children =
+              @current_function.at("parameters").children =
                 Nokogiri::HTML5::DocumentFragment.parse(argument_fragments)
             rescue JSON::ParserError
               return function_buffer

lib/completions/prompt.rb

@@ -6,12 +6,22 @@ class Prompt
       INVALID_TURN = Class.new(StandardError)
 
       attr_reader :messages
-      attr_accessor :tools
-
-      def initialize(system_message_text = nil, messages: [], tools: [], skip_validations: false)
+      attr_accessor :tools, :topic_id, :post_id
+
+      def initialize(
+        system_message_text = nil,
+        messages: [],
+        tools: [],
+        skip_validations: false,
+        topic_id: nil,
+        post_id: nil
+      )
         raise ArgumentError, "messages must be an array" if !messages.is_a?(Array)
         raise ArgumentError, "tools must be an array" if !tools.is_a?(Array)
 
+        @topic_id = topic_id
+        @post_id = post_id
+
         @messages = []
         @skip_validations = skip_validations
 

spec/lib/completions/endpoints/endpoint_compliance.rb

@@ -178,11 +178,9 @@ def regular_mode_simple_prompt(mock)
   def regular_mode_tools(mock)
     prompt = generic_prompt(tools: [mock.tool])
     a_dialect = dialect(prompt: prompt)
-
     mock.stub_tool_call(a_dialect.translate)
 
     completion_response = endpoint.perform_completion!(a_dialect, user)
-
     expect(completion_response).to eq(mock.invocation_response)
   end