FEATURE: anthropic function calling (#654)

Adds support for native tool calling (both streaming and non streaming) for Anthropic.

This improves general tool support on the Anthropic models.

Author: SamSaffron

lib/completions/anthropic_message_processor.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+class DiscourseAi::Completions::AnthropicMessageProcessor
+  class AnthropicToolCall
+    attr_reader :name, :raw_json, :id
+
+    def initialize(name, id)
+      @name = name
+      @id = id
+      @raw_json = +""
+    end
+
+    def append(json)
+      @raw_json << json
+    end
+  end
+
+  attr_reader :tool_calls, :input_tokens, :output_tokens
+
+  def initialize(streaming_mode:)
+    @streaming_mode = streaming_mode
+    @tool_calls = []
+  end
+
+  def to_xml_tool_calls(function_buffer)
+    return function_buffer if @tool_calls.blank?
+
+    function_buffer = Nokogiri::HTML5.fragment(<<~TEXT)
+      <function_calls>
+      </function_calls>
+    TEXT
+
+    @tool_calls.each do |tool_call|
+      node =
+        function_buffer.at("function_calls").add_child(
+          Nokogiri::HTML5::DocumentFragment.parse(
+            DiscourseAi::Completions::Endpoints::Base.noop_function_call_text + "\n",
+          ),
+        )
+
+      params = JSON.parse(tool_call.raw_json, symbolize_names: true)
+      xml = params.map { |name, value| "<#{name}>#{value}</#{name}>" }.join("\n")
+
+      node.at("tool_name").content = tool_call.name
+      node.at("tool_id").content = tool_call.id
+      node.at("parameters").children = Nokogiri::HTML5::DocumentFragment.parse(xml) if xml.present?
+    end
+
+    function_buffer
+  end
+
+  def process_message(payload)
+    result = ""
+    parsed = JSON.parse(payload, symbolize_names: true)
+
+    if @streaming_mode
+      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
+      elsif parsed[:type] == "content_block_start" || parsed[:type] == "content_block_delta"
+        if @tool_calls.present?
+          result = parsed.dig(:delta, :partial_json).to_s
+          @tool_calls.last.append(result)
+        else
+          result = parsed.dig(:delta, :text).to_s
+        end
+      elsif parsed[:type] == "message_start"
+        @input_tokens = parsed.dig(:message, :usage, :input_tokens)
+      elsif parsed[:type] == "message_delta"
+        @output_tokens =
+          parsed.dig(:usage, :output_tokens) || parsed.dig(:delta, :usage, :output_tokens)
+      elsif parsed[:type] == "message_stop"
+        # bedrock has this ...
+        if bedrock_stats = parsed.dig("amazon-bedrock-invocationMetrics".to_sym)
+          @input_tokens = bedrock_stats[:inputTokenCount] || @input_tokens
+          @output_tokens = bedrock_stats[:outputTokenCount] || @output_tokens
+        end
+      end
+    else
+      content = parsed.dig(:content)
+      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.last.append(tool_call[:input].to_json)
+        else
+          result = parsed.dig(:content, 0, :text).to_s
+        end
+      end
+
+      @input_tokens = parsed.dig(:usage, :input_tokens)
+      @output_tokens = parsed.dig(:usage, :output_tokens)
+    end
+
+    result
+  end
+end

lib/completions/dialects/claude.rb

@@ -15,10 +15,12 @@ def can_translate?(model_name)
         class ClaudePrompt
           attr_reader :system_prompt
           attr_reader :messages
+          attr_reader :tools
 
-          def initialize(system_prompt, messages)
+          def initialize(system_prompt, messages, tools)
             @system_prompt = system_prompt
             @messages = messages
+            @tools = tools
           end
         end
 
@@ -46,7 +48,11 @@ def translate
             previous_message = message
           end
 
-          ClaudePrompt.new(system_prompt.presence, interleving_messages)
+          ClaudePrompt.new(
+            system_prompt.presence,
+            interleving_messages,
+            tools_dialect.translated_tools,
+          )
         end
 
         def max_prompt_tokens
@@ -58,6 +64,18 @@ def max_prompt_tokens
 
         private
 
+        def tools_dialect
+          @tools_dialect ||= DiscourseAi::Completions::Dialects::ClaudeTools.new(prompt.tools)
+        end
+
+        def tool_call_msg(msg)
+          tools_dialect.from_raw_tool_call(msg)
+        end
+
+        def tool_msg(msg)
+          tools_dialect.from_raw_tool(msg)
+        end
+
         def model_msg(msg)
           { role: "assistant", content: msg[:content] }
         end

lib/completions/dialects/claude_tools.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Completions
+    module Dialects
+      class ClaudeTools
+        def initialize(tools)
+          @raw_tools = tools
+        end
+
+        def translated_tools
+          # Transform the raw tools into the required Anthropic Claude API format
+          raw_tools.map do |t|
+            properties = {}
+            required = []
+
+            if t[:parameters]
+              properties =
+                t[:parameters].each_with_object({}) do |param, h|
+                  h[param[:name]] = {
+                    type: param[:type],
+                    description: param[:description],
+                  }.tap { |hash| hash[:items] = { type: param[:item_type] } if param[:item_type] }
+                end
+              required =
+                t[:parameters].select { |param| param[:required] }.map { |param| param[:name] }
+            end
+
+            {
+              name: t[:name],
+              description: t[:description],
+              input_schema: {
+                type: "object",
+                properties: properties,
+                required: required,
+              },
+            }
+          end
+        end
+
+        def instructions
+          "" # Noop. Tools are listed separate.
+        end
+
+        def from_raw_tool_call(raw_message)
+          call_details = JSON.parse(raw_message[:content], symbolize_names: true)
+          tool_call_id = raw_message[:id]
+
+          {
+            role: "assistant",
+            content: [
+              {
+                type: "tool_use",
+                id: tool_call_id,
+                name: raw_message[:name],
+                input: call_details[:arguments],
+              },
+            ],
+          }
+        end
+
+        def from_raw_tool(raw_message)
+          {
+            role: "user",
+            content: [
+              {
+                type: "tool_result",
+                tool_use_id: raw_message[:id],
+                content: raw_message[:content],
+              },
+            ],
+          }
+        end
+
+        private
+
+        attr_reader :raw_tools
+      end
+    end
+  end
+end

lib/completions/endpoints/anthropic.rb

@@ -45,10 +45,7 @@ def default_options(dialect)
               raise "Unsupported model: #{model}"
             end
 
-          options = { model: mapped_model, max_tokens: 3_000 }
-
-          options[:stop_sequences] = ["</function_calls>"] if dialect.prompt.has_tools?
-          options
+          { model: mapped_model, max_tokens: 3_000 }
         end
 
         def provider_id
@@ -73,6 +70,7 @@ def prepare_payload(prompt, model_params, dialect)
 
           payload[:system] = prompt.system_prompt if prompt.system_prompt.present?
           payload[:stream] = true if @streaming_mode
+          payload[:tools] = prompt.tools if prompt.tools.present?
 
           payload
         end
@@ -87,30 +85,30 @@ def prepare_request(payload)
           Net::HTTP::Post.new(model_uri, headers).tap { |r| r.body = payload }
         end
 
-        def final_log_update(log)
-          log.request_tokens = @input_tokens if @input_tokens
-          log.response_tokens = @output_tokens if @output_tokens
+        def processor
+          @processor ||=
+            DiscourseAi::Completions::AnthropicMessageProcessor.new(streaming_mode: @streaming_mode)
+        end
+
+        def add_to_function_buffer(function_buffer, partial: nil, payload: nil)
+          processor.to_xml_tool_calls(function_buffer) if !partial
         end
 
         def extract_completion_from(response_raw)
-          result = ""
-          parsed = JSON.parse(response_raw, symbolize_names: true)
-
-          if @streaming_mode
-            if parsed[:type] == "content_block_start" || parsed[:type] == "content_block_delta"
-              result = parsed.dig(:delta, :text).to_s
-            elsif parsed[:type] == "message_start"
-              @input_tokens = parsed.dig(:message, :usage, :input_tokens)
-            elsif parsed[:type] == "message_delta"
-              @output_tokens = parsed.dig(:delta, :usage, :output_tokens)
-            end
-          else
-            result = parsed.dig(:content, 0, :text).to_s
-            @input_tokens = parsed.dig(:usage, :input_tokens)
-            @output_tokens = parsed.dig(:usage, :output_tokens)
-          end
+          processor.process_message(response_raw)
+        end
+
+        def has_tool?(_response_data)
+          processor.tool_calls.present?
+        end
+
+        def final_log_update(log)
+          log.request_tokens = processor.input_tokens if processor.input_tokens
+          log.response_tokens = processor.output_tokens if processor.output_tokens
+        end
 
-          result
+        def native_tool_support?
+          true
         end
 
         def partials_from(decoded_chunk)

lib/completions/endpoints/aws_bedrock.rb

@@ -36,7 +36,6 @@ def normalize_model_params(model_params)
 
         def default_options(dialect)
           options = { max_tokens: 3_000, anthropic_version: "bedrock-2023-05-31" }
-          options[:stop_sequences] = ["</function_calls>"] if dialect.prompt.has_tools?
           options
         end
 
@@ -82,6 +81,8 @@ def model_uri
         def prepare_payload(prompt, model_params, dialect)
           payload = default_options(dialect).merge(model_params).merge(messages: prompt.messages)
           payload[:system] = prompt.system_prompt if prompt.system_prompt.present?
+          payload[:tools] = prompt.tools if prompt.tools.present?
+
           payload
         end
 
@@ -142,35 +143,35 @@ def decode(chunk)
         end
 
         def final_log_update(log)
-          log.request_tokens = @input_tokens if @input_tokens
-          log.response_tokens = @output_tokens if @output_tokens
+          log.request_tokens = processor.input_tokens if processor.input_tokens
+          log.response_tokens = processor.output_tokens if processor.output_tokens
+        end
+
+        def processor
+          @processor ||=
+            DiscourseAi::Completions::AnthropicMessageProcessor.new(streaming_mode: @streaming_mode)
+        end
+
+        def add_to_function_buffer(function_buffer, partial: nil, payload: nil)
+          processor.to_xml_tool_calls(function_buffer) if !partial
         end
 
         def extract_completion_from(response_raw)
-          result = ""
-          parsed = JSON.parse(response_raw, symbolize_names: true)
-
-          if @streaming_mode
-            if parsed[:type] == "content_block_start" || parsed[:type] == "content_block_delta"
-              result = parsed.dig(:delta, :text).to_s
-            elsif parsed[:type] == "message_start"
-              @input_tokens = parsed.dig(:message, :usage, :input_tokens)
-            elsif parsed[:type] == "message_delta"
-              @output_tokens = parsed.dig(:delta, :usage, :output_tokens)
-            end
-          else
-            result = parsed.dig(:content, 0, :text).to_s
-            @input_tokens = parsed.dig(:usage, :input_tokens)
-            @output_tokens = parsed.dig(:usage, :output_tokens)
-          end
+          processor.process_message(response_raw)
+        end
 
-          result
+        def has_tool?(_response_data)
+          processor.tool_calls.present?
         end
 
         def partials_from(decoded_chunks)
           decoded_chunks
         end
 
+        def native_tool_support?
+          true
+        end
+
         def chunk_to_string(chunk)
           joined = +chunk.join("\n")
           joined << "\n" if joined.length > 0