FEATURE: support tool progress callbacks

This is anthropic only for now, but we can get a callback as tool
is completing, this gives us the ability to show progress to user
as the function is populating.

Author: SamSaffron

lib/ai_bot/bot.rb

@@ -105,8 +105,15 @@ def reply(context, &update_blk)
           tool_found = false
           force_tool_if_needed(prompt, context)
 
+          tool_progress = proc { |progress| p progress }
+
           result =
-            llm.generate(prompt, feature_name: "bot", **llm_kwargs) do |partial, cancel|
+            llm.generate(
+              prompt,
+              feature_name: "bot",
+              tool_progress: tool_progress,
+              **llm_kwargs,
+            ) do |partial, cancel|
               tools = persona.find_tools(partial, bot_user: user, llm: llm, context: context)
 
               if (tools.present?)

lib/completions/anthropic_message_processor.rb

@@ -1,25 +1,83 @@
 # frozen_string_literal: true
 
 class DiscourseAi::Completions::AnthropicMessageProcessor
+  class ToolCallProgressTracker
+    attr_reader :current_key, :current_value, :tool_call
+
+    def initialize(tool_call)
+      @tool_call = tool_call
+      @current_key = nil
+      @current_value = nil
+      @parser = DiscourseAi::Utils::JsonStreamingParser.new
+
+      @parser.key do |k|
+        @current_key = k
+        @current_value = nil
+      end
+      @parser.value do |v|
+        @current_value = v
+
+        if @current_key
+          tool_call.tool_progress.call(
+            { name: tool_call.name, id: tool_call.id, key: @current_key, value: @current_value },
+          )
+        end
+      end
+    end
+
+    def <<(json)
+      # llm could send broken json
+      # in that case just deal with it later
+      # don't stream
+      return if @broken
+
+      begin
+        @parser << json
+      rescue DiscourseAi::Utils::ParserError
+        @broken = true
+        return
+      end
+
+      if @parser.state == :start_string && @current_key
+        # this is is worth notifying
+        tool_call.tool_progress.call(
+          {
+            name: tool_call.name,
+            id: tool_call.id,
+            key: @current_key,
+            value: @parser.buf,
+            done: false,
+          },
+        )
+      end
+    end
+  end
+
   class AnthropicToolCall
-    attr_reader :name, :raw_json, :id
+    attr_reader :name, :raw_json, :id, :tool_progress
 
-    def initialize(name, id)
+    def initialize(name, id, tool_progress)
       @name = name
       @id = id
       @raw_json = +""
+      if tool_progress
+        @tool_progress = tool_progress
+        @tool_call_progress_tracker = ToolCallProgressTracker.new(self)
+      end
     end
 
     def append(json)
       @raw_json << json
+      @tool_call_progress_tracker << json if @tool_progress
     end
   end
 
   attr_reader :tool_calls, :input_tokens, :output_tokens
 
-  def initialize(streaming_mode:)
+  def initialize(streaming_mode:, tool_progress:)
     @streaming_mode = streaming_mode
     @tool_calls = []
+    @tool_progress = tool_progress
   end
 
   def to_xml_tool_calls(function_buffer)
@@ -58,7 +116,7 @@ def process_message(payload)
       if parsed[:type] == "content_block_start" && parsed.dig(:content_block, :type) == "tool_use"
         tool_name = parsed.dig(:content_block, :name)
         tool_id = parsed.dig(:content_block, :id)
-        @tool_calls << AnthropicToolCall.new(tool_name, tool_id) if tool_name
+        @tool_calls << AnthropicToolCall.new(tool_name, tool_id, @tool_progress) if tool_name
       elsif parsed[:type] == "content_block_start" || parsed[:type] == "content_block_delta"
         if @tool_calls.present?
           result = parsed.dig(:delta, :partial_json).to_s
@@ -83,7 +141,7 @@ def process_message(payload)
       if content.is_a?(Array)
         tool_call = content.find { |c| c[:type] == "tool_use" }
         if tool_call
-          @tool_calls << AnthropicToolCall.new(tool_call[:name], tool_call[:id])
+          @tool_calls << AnthropicToolCall.new(tool_call[:name], tool_call[:id], @tool_progress)
           @tool_calls.last.append(tool_call[:input].to_json)
         else
           result = parsed.dig(:content, 0, :text).to_s

lib/completions/endpoints/anthropic.rb

@@ -92,7 +92,10 @@ def prepare_request(payload)
 
         def processor
           @processor ||=
-            DiscourseAi::Completions::AnthropicMessageProcessor.new(streaming_mode: @streaming_mode)
+            DiscourseAi::Completions::AnthropicMessageProcessor.new(
+              streaming_mode: @streaming_mode,
+              tool_progress: @tool_progress,
+            )
         end
 
         def add_to_function_buffer(function_buffer, partial: nil, payload: nil)

lib/completions/endpoints/aws_bedrock.rb

@@ -157,7 +157,10 @@ def final_log_update(log)
 
         def processor
           @processor ||=
-            DiscourseAi::Completions::AnthropicMessageProcessor.new(streaming_mode: @streaming_mode)
+            DiscourseAi::Completions::AnthropicMessageProcessor.new(
+              streaming_mode: @streaming_mode,
+              tool_progress: @tool_progress,
+            )
         end
 
         def add_to_function_buffer(function_buffer, partial: nil, payload: nil)

lib/completions/endpoints/base.rb

@@ -62,12 +62,14 @@ def perform_completion!(
           model_params = {},
           feature_name: nil,
           feature_context: nil,
+          tool_progress: nil,
           &blk
         )
           allow_tools = dialect.prompt.has_tools?
           model_params = normalize_model_params(model_params)
           orig_blk = blk
 
+          @tool_progress = tool_progress
           @streaming_mode = block_given?
           to_strip = xml_tags_to_strip(dialect)
           @xml_stripper =

lib/completions/endpoints/canned_response.rb

@@ -28,7 +28,8 @@ def perform_completion!(
           _user,
           _model_params,
           feature_name: nil,
-          feature_context: nil
+          feature_context: nil,
+          tool_progress: nil
         )
           @dialect = dialect
           response = responses[completions]

lib/completions/endpoints/fake.rb

@@ -116,7 +116,8 @@ def perform_completion!(
           user,
           model_params = {},
           feature_name: nil,
-          feature_context: nil
+          feature_context: nil,
+          tool_progress: nil
         )
           self.class.last_call = { dialect: dialect, user: user, model_params: model_params }
 

lib/completions/endpoints/open_ai.rb

@@ -33,6 +33,7 @@ def perform_completion!(
           model_params = {},
           feature_name: nil,
           feature_context: nil,
+          tool_progress: nil,
           &blk
         )
           if dialect.respond_to?(:is_gpt_o?) && dialect.is_gpt_o? && block_given?

lib/completions/llm.rb

@@ -192,6 +192,7 @@ def generate(
         user:,
         feature_name: nil,
         feature_context: nil,
+        tool_progress: nil,
         &partial_read_blk
       )
         self.class.record_prompt(prompt)
@@ -226,6 +227,7 @@ def generate(
           model_params,
           feature_name: feature_name,
           feature_context: feature_context,
+          tool_progress: tool_progress,
           &partial_read_blk
         )
       end