Implement streaming and generation of thinking tokens

This is required for claude and introduces some new concepts into
prompt

Author: SamSaffron

lib/completions/anthropic_message_processor.rb

@@ -44,13 +44,15 @@ def to_tool_call
     end
   end
 
-  attr_reader :tool_calls, :input_tokens, :output_tokens
+  attr_reader :tool_calls, :input_tokens, :output_tokens, :output_thinking
 
-  def initialize(streaming_mode:, partial_tool_calls: false)
+  def initialize(streaming_mode:, partial_tool_calls: false, output_thinking: false)
     @streaming_mode = streaming_mode
     @tool_calls = []
     @current_tool_call = nil
     @partial_tool_calls = partial_tool_calls
+    @output_thinking = output_thinking
+    @thinking = nil
   end
 
   def to_tool_calls
@@ -69,13 +71,48 @@ def process_streamed_message(parsed)
           tool_id,
           partial_tool_calls: @partial_tool_calls,
         ) if tool_name
+    elsif parsed[:type] == "content_block_start" && parsed.dig(:content_block, :type) == "thinking"
+      if @output_thinking
+        @thinking =
+          DiscourseAi::Completions::Thinking.new(
+            message: +parsed.dig(:content_block, :thinking).to_s,
+            signature: +"",
+            partial: true,
+          )
+        result = @thinking.dup
+      end
+    elsif parsed[:type] == "content_block_delta" && parsed.dig(:delta, :type) == "thinking_delta"
+      if @output_thinking
+        delta = parsed.dig(:delta, :thinking)
+        @thinking.message << delta if @thinking
+        result = DiscourseAi::Completions::Thinking.new(message: delta, partial: true)
+      end
+    elsif parsed[:type] == "content_block_delta" && parsed.dig(:delta, :type) == "signature_delta"
+      if @output_thinking
+        @thinking.signature << parsed.dig(:delta, :signature) if @thinking
+      end
+    elsif parsed[:type] == "content_block_stop" && @thinking
+      @thinking.partial = false
+      result = @thinking
+      @thinking = nil
     elsif parsed[:type] == "content_block_start" || parsed[:type] == "content_block_delta"
       if @current_tool_call
         tool_delta = parsed.dig(:delta, :partial_json).to_s
         @current_tool_call.append(tool_delta)
         result = @current_tool_call.partial_tool_call if @current_tool_call.has_partial?
+      elsif parsed.dig(:content_block, :type) == "redacted_thinking"
+        if @output_thinking
+          result =
+            DiscourseAi::Completions::Thinking.new(
+              message: nil,
+              signature: parsed.dig(:content_block, :data),
+              redacted: true,
+            )
+        end
       else
         result = parsed.dig(:delta, :text).to_s
+        # no need to return empty strings for streaming, no value
+        result = nil if result == ""
       end
     elsif parsed[:type] == "content_block_stop"
       if @current_tool_call
@@ -105,15 +142,32 @@ def process_message(payload)
     content = parsed.dig(:content)
     if content.is_a?(Array)
       result =
-        content.map do |data|
-          if data[:type] == "tool_use"
-            call = AnthropicToolCall.new(data[:name], data[:id])
-            call.append(data[:input].to_json)
-            call.to_tool_call
-          else
-            data[:text]
+        content
+          .map do |data|
+            if data[:type] == "tool_use"
+              call = AnthropicToolCall.new(data[:name], data[:id])
+              call.append(data[:input].to_json)
+              call.to_tool_call
+            elsif data[:type] == "thinking"
+              if @output_thinking
+                DiscourseAi::Completions::Thinking.new(
+                  message: data[:thinking],
+                  signature: data[:signature],
+                )
+              end
+            elsif data[:type] == "redacted_thinking"
+              if @output_thinking
+                DiscourseAi::Completions::Thinking.new(
+                  message: nil,
+                  signature: data[:data],
+                  redacted: true,
+                )
+              end
+            else
+              data[:text]
+            end
           end
-        end
+          .compact
     end
 
     @input_tokens = parsed.dig(:usage, :input_tokens)

lib/completions/dialects/claude.rb

@@ -87,7 +87,30 @@ def tool_msg(msg)
         end
 
         def model_msg(msg)
-          { role: "assistant", content: msg[:content] }
+          if msg[:thinking] || msg[:redacted_thinking_signature]
+            content_array = []
+
+            if msg[:thinking]
+              content_array << {
+                type: "thinking",
+                thinking: msg[:thinking],
+                signature: msg[:thinking_signature],
+              }
+            end
+
+            if msg[:redacted_thinking_signature]
+              content_array << {
+                type: "redacted_thinking",
+                data: msg[:redacted_thinking_signature],
+              }
+            end
+
+            content_array << { type: "text", text: msg[:content] }
+
+            { role: "assistant", content: content_array }
+          else
+            { role: "assistant", content: msg[:content] }
+          end
         end
 
         def system_msg(msg)

lib/completions/endpoints/anthropic.rb

@@ -123,6 +123,7 @@ def processor
             DiscourseAi::Completions::AnthropicMessageProcessor.new(
               streaming_mode: @streaming_mode,
               partial_tool_calls: partial_tool_calls,
+              output_thinking: output_thinking,
             )
         end
 

lib/completions/endpoints/base.rb

@@ -4,7 +4,7 @@ module DiscourseAi
   module Completions
     module Endpoints
       class Base
-        attr_reader :partial_tool_calls
+        attr_reader :partial_tool_calls, :output_thinking
 
         CompletionFailed = Class.new(StandardError)
         # 6 minutes
@@ -67,12 +67,15 @@ def perform_completion!(
           feature_name: nil,
           feature_context: nil,
           partial_tool_calls: false,
+          output_thinking: false,
           &blk
         )
           LlmQuota.check_quotas!(@llm_model, user)
           start_time = Time.now
 
           @partial_tool_calls = partial_tool_calls
+          @output_thinking = output_thinking
+
           model_params = normalize_model_params(model_params)
           orig_blk = blk
 
@@ -85,6 +88,7 @@ def perform_completion!(
                 feature_name: feature_name,
                 feature_context: feature_context,
                 partial_tool_calls: partial_tool_calls,
+                output_thinking: output_thinking,
               )
 
             wrapped = result

lib/completions/llm.rb

@@ -234,6 +234,7 @@ def initialize(dialect_klass, gateway_klass, llm_model, gateway: nil)
       # @param feature_name { String - Optional } - The feature name to use for the completion.
       # @param feature_context { Hash - Optional } - The feature context to use for the completion.
       # @param partial_tool_calls { Boolean - Optional } - If true, the completion will return partial tool calls.
+      # @param output_thinking { Boolean - Optional } - If true, the completion will return the thinking output for thinking models.
       #
       # @param &on_partial_blk { Block - Optional } - The passed block will get called with the LLM partial response alongside a cancel function.
       #
@@ -250,6 +251,7 @@ def generate(
         feature_name: nil,
         feature_context: nil,
         partial_tool_calls: false,
+        output_thinking: false,
         &partial_read_blk
       )
         self.class.record_prompt(prompt)
@@ -285,6 +287,7 @@ def generate(
           feature_name: feature_name,
           feature_context: feature_context,
           partial_tool_calls: partial_tool_calls,
+          output_thinking: output_thinking,
           &partial_read_blk
         )
       end

lib/completions/prompt.rb

@@ -41,12 +41,26 @@ def initialize(
         @tool_choice = tool_choice
       end
 
-      def push(type:, content:, id: nil, name: nil, upload_ids: nil)
+      def push(
+        type:,
+        content:,
+        id: nil,
+        name: nil,
+        upload_ids: nil,
+        thinking: nil,
+        thinking_signature: nil,
+        redacted_thinking_signature: nil
+      )
         return if type == :system
         new_message = { type: type, content: content }
         new_message[:name] = name.to_s if name
         new_message[:id] = id.to_s if id
         new_message[:upload_ids] = upload_ids if upload_ids
+        new_message[:thinking] = thinking if thinking
+        new_message[:thinking_signature] = thinking_signature if thinking_signature
+        new_message[
+          :redacted_thinking_signature
+        ] = redacted_thinking_signature if redacted_thinking_signature
 
         validate_message(new_message)
         validate_turn(messages.last, new_message)
@@ -73,7 +87,16 @@ def validate_message(message)
           raise ArgumentError, "message type must be one of #{valid_types}"
         end
 
-        valid_keys = %i[type content id name upload_ids]
+        valid_keys = %i[
+          type
+          content
+          id
+          name
+          upload_ids
+          thinking
+          thinking_signature
+          redacted_thinking_signature
+        ]
         if (invalid_keys = message.keys - valid_keys).any?
           raise ArgumentError, "message contains invalid keys: #{invalid_keys}"
         end

lib/completions/thinking.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Completions
+    class Thinking
+      attr_accessor :message, :signature, :redacted, :partial
+
+      def initialize(message:, signature: nil, redacted: false, partial: false)
+        @message = message
+        @signature = signature
+        @redacted = redacted
+        @partial = partial
+      end
+
+      def ==(other)
+        message == other.message && signature == other.signature && redacted == other.redacted &&
+          partial == other.partial
+      end
+
+      def dup
+        Thinking.new(
+          message: message.dup,
+          signature: signature.dup,
+          redacted: redacted,
+          partial: partial,
+        )
+      end
+
+      def to_s
+        "#{message} - #{signature} - #{redacted} - #{partial}"
+      end
+    end
+  end
+end