FEATURE: support for claude opus and sonnet (#508)

This provides new support for messages API from Claude.

It is required for latest model access.

Also corrects implementation of function calls.

* Fix message interleving

* fix broken spec

* add new models to automation

config/locales/client.en.yml

@@ -12,6 +12,8 @@ en:
         gpt_3_5_turbo: GPT 3.5 Turbo
         claude_2: Claude 2
         gemini_pro: Gemini Pro
+        claude_3_opus: Claude 3 Opus
+        claude_3_sonnet: Claude 3 Sonnet
       scriptables:
         llm_report:
           fields:
@@ -227,6 +229,9 @@ en:
           context: "Interactions to share:"
 
         bot_names:
+          fake: "Fake Test Bot"
+          claude-3-opus: "Claude 3 Opus"
+          claude-3-sonnet: "Claude 3 Sonnet"
           gpt-4: "GPT-4"
           gpt-4-turbo: "GPT-4 Turbo"
           gpt-3:

config/settings.yml

@@ -313,6 +313,8 @@ discourse_ai:
       - claude-2
       - gemini-pro
       - mixtral-8x7B-Instruct-V0.1
+      - claude-3-opus
+      - claude-3-sonnet
   ai_bot_add_to_header:
     default: true
     client: true

lib/ai_bot/bot.rb

@@ -171,6 +171,10 @@ module DiscourseAi
           "google:gemini-pro"
         when DiscourseAi::AiBot::EntryPoint::FAKE_ID
           "fake:fake"
+        when DiscourseAi::AiBot::EntryPoint::CLAUDE_3_OPUS_ID
+          "anthropic:claude-3-opus"
+        when DiscourseAi::AiBot::EntryPoint::CLAUDE_3_SONNET_ID
+          "anthropic:claude-3-sonnet"
         else
           nil
         end

lib/ai_bot/entry_point.rb

@@ -12,6 +12,8 @@ module DiscourseAi
       MIXTRAL_ID = -114
       GEMINI_ID = -115
       FAKE_ID = -116 # only used for dev and test
+      CLAUDE_3_OPUS_ID = -117
+      CLAUDE_3_SONNET_ID = -118
 
       BOTS = [
         [GPT4_ID, "gpt4_bot", "gpt-4"],
@@ -21,6 +23,8 @@ module DiscourseAi
         [MIXTRAL_ID, "mixtral_bot", "mixtral-8x7B-Instruct-V0.1"],
         [GEMINI_ID, "gemini_bot", "gemini-pro"],
         [FAKE_ID, "fake_bot", "fake"],
+        [CLAUDE_3_OPUS_ID, "claude_3_opus_bot", "claude-3-opus"],
+        [CLAUDE_3_SONNET_ID, "claude_3_sonnet_bot", "claude-3-sonnet"],
       ]
 
       BOT_USER_IDS = BOTS.map(&:first)
@@ -41,6 +45,10 @@ module DiscourseAi
           GEMINI_ID
         in "fake"
           FAKE_ID
+        in "claude-3-opus"
+          CLAUDE_3_OPUS_ID
+        in "claude-3-sonnet"
+          CLAUDE_3_SONNET_ID
         else
           nil
         end

lib/ai_bot/personas/researcher.rb

@@ -14,11 +14,19 @@ module DiscourseAi
 
         def system_prompt
           <<~PROMPT
-            You are research bot. With access to the internet you can find information for users.
+            You are research bot. With access to Google you can find information for users.
 
-            - You fully understand Discourse Markdown and generate it.
+            - You are conversing with: {participants}
-            - When generating responses you always cite your sources.
+            - You understand **Discourse Markdown** and generate it.
+            - When generating responses you always cite your sources using Markdown footnotes.
             - When possible you also quote the sources.
+
+            Example:
+
+            **This** is a content[^1] with two footnotes[^2].
+
+            [^1]: https://www.example.com
+            [^2]: https://www.example2.com
           PROMPT
         end
       end

lib/ai_bot/playground.rb

@@ -109,7 +109,6 @@ module DiscourseAi
             .pluck(:raw, :username, "post_custom_prompts.custom_prompt")
 
         result = []
-        first = true
 
         context.reverse_each do |raw, username, custom_prompt|
           custom_prompt_translation =
@@ -129,12 +128,7 @@ module DiscourseAi
             end
 
           if custom_prompt.present?
-            if first
+            custom_prompt.each(&custom_prompt_translation)
-              custom_prompt.each(&custom_prompt_translation)
-              first = false
-            else
-              custom_prompt.first(2).each(&custom_prompt_translation)
-            end
           else
             context = {
               content: raw,

lib/automation.rb

@@ -8,6 +8,8 @@ module DiscourseAi
       { id: "gpt-3.5-turbo", name: "discourse_automation.ai_models.gpt_3_5_turbo" },
       { id: "claude-2", name: "discourse_automation.ai_models.claude_2" },
       { id: "gemini-pro", name: "discourse_automation.ai_models.gemini_pro" },
+      { id: "claude-3-sonnet", name: "discourse_automation.ai_models.claude_3_sonnet" },
+      { id: "claude-3-opus", name: "discourse_automation.ai_models.claude_3_opus" },
     ]
   end
 end

lib/automation/report_runner.rb

@@ -211,12 +211,13 @@ Follow the provided writing composition instructions carefully and precisely ste
 
       def translate_model(model)
         return "google:gemini-pro" if model == "gemini-pro"
-        return "open_ai:#{model}" if model != "claude-2"
+        return "open_ai:#{model}" if model.start_with? "gpt"
+        return "anthropic:#{model}" if model.start_with? "claude-3"
 
         if DiscourseAi::Completions::Endpoints::AwsBedrock.correctly_configured?("claude-2")
-          "aws_bedrock:claude-2"
+          "aws_bedrock:#{model}"
         else
-          "anthropic:claude-2"
+          "anthropic:#{model}"
         end
       end
     end

lib/completions/dialects/claude_messages.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Completions
+    module Dialects
+      class ClaudeMessages < Dialect
+        class << self
+          def can_translate?(model_name)
+            # TODO: add haiku not released yet as of 2024-03-05
+            %w[claude-3-sonnet claude-3-opus].include?(model_name)
+          end
+
+          def tokenizer
+            DiscourseAi::Tokenizer::AnthropicTokenizer
+          end
+        end
+
+        class ClaudePrompt
+          attr_reader :system_prompt
+          attr_reader :messages
+
+          def initialize(system_prompt, messages)
+            @system_prompt = system_prompt
+            @messages = messages
+          end
+        end
+
+        def translate
+          messages = prompt.messages
+          system_prompt = +""
+
+          messages =
+            trim_messages(messages)
+              .map do |msg|
+                case msg[:type]
+                when :system
+                  system_prompt << msg[:content]
+                  nil
+                when :tool_call
+                  { role: "assistant", content: tool_call_to_xml(msg) }
+                when :tool
+                  { role: "user", content: tool_result_to_xml(msg) }
+                when :model
+                  { role: "assistant", content: msg[:content] }
+                when :user
+                  content = +""
+                  content << "#{msg[:id]}: " if msg[:id]
+                  content << msg[:content]
+
+                  { role: "user", content: content }
+                end
+              end
+              .compact
+
+          if prompt.tools.present?
+            system_prompt << "\n\n"
+            system_prompt << build_tools_prompt
+          end
+
+          interleving_messages = []
+
+          previous_message = nil
+          messages.each do |message|
+            if previous_message
+              if previous_message[:role] == "user" && message[:role] == "user"
+                interleving_messages << { role: "assistant", content: "OK" }
+              elsif previous_message[:role] == "assistant" && message[:role] == "assistant"
+                interleving_messages << { role: "user", content: "OK" }
+              end
+            end
+            interleving_messages << message
+            previous_message = message
+          end
+
+          ClaudePrompt.new(system_prompt.presence, interleving_messages)
+        end
+
+        def max_prompt_tokens
+          # Longer term it will have over 1 million
+          200_000 # Claude-3 has a 200k context window for now
+        end
+      end
+    end
+  end
+end

lib/completions/dialects/dialect.rb

@@ -17,6 +17,7 @@ module DiscourseAi
               DiscourseAi::Completions::Dialects::OrcaStyle,
               DiscourseAi::Completions::Dialects::Gemini,
               DiscourseAi::Completions::Dialects::Mixtral,
+              DiscourseAi::Completions::Dialects::ClaudeMessages,
             ]
 
             if Rails.env.test? || Rails.env.development?
@@ -64,6 +65,38 @@ module DiscourseAi
           raise NotImplemented
         end
 
+        def tool_result_to_xml(message)
+          (<<~TEXT).strip
+            <function_results>
+            <result>
+            <tool_name>#{message[:id]}</tool_name>
+            <json>
+            #{message[:content]}
+            </json>
+            </result>
+            </function_results>
+          TEXT
+        end
+
+        def tool_call_to_xml(message)
+          parsed = JSON.parse(message[:content], symbolize_names: true)
+          parameters = +""
+
+          if parsed[:arguments]
+            parameters << "<parameters>\n"
+            parsed[:arguments].each { |k, v| parameters << "<#{k}>#{v}</#{k}>\n" }
+            parameters << "</parameters>\n"
+          end
+
+          (<<~TEXT).strip
+            <function_calls>
+            <invoke>
+            <tool_name>#{parsed[:name]}</tool_name>
+            #{parameters}</invoke>
+            </function_calls>
+          TEXT
+        end
+
         def tools
           tools = +""
 

lib/completions/endpoints/anthropic_messages.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Completions
+    module Endpoints
+      class AnthropicMessages < Base
+        class << self
+          def can_contact?(endpoint_name, model_name)
+            endpoint_name == "anthropic" && %w[claude-3-opus claude-3-sonnet].include?(model_name)
+          end
+
+          def dependant_setting_names
+            %w[ai_anthropic_api_key]
+          end
+
+          def correctly_configured?(_model_name)
+            SiteSetting.ai_anthropic_api_key.present?
+          end
+
+          def endpoint_name(model_name)
+            "Anthropic - #{model_name}"
+          end
+        end
+
+        def normalize_model_params(model_params)
+          # max_tokens, temperature, stop_sequences are already supported
+          model_params
+        end
+
+        def default_options
+          { model: model + "-20240229", max_tokens: 3_000, stop_sequences: ["</function_calls>"] }
+        end
+
+        def provider_id
+          AiApiAuditLog::Provider::Anthropic
+        end
+
+        private
+
+        # this is an approximation, we will update it later if request goes through
+        def prompt_size(prompt)
+          super(prompt.system_prompt.to_s + " " + prompt.messages.to_s)
+        end
+
+        def model_uri
+          @uri ||= URI("https://api.anthropic.com/v1/messages")
+        end
+
+        def prepare_payload(prompt, model_params, _dialect)
+          payload = default_options.merge(model_params).merge(messages: prompt.messages)
+
+          payload[:system] = prompt.system_prompt if prompt.system_prompt.present?
+          payload[:stream] = true if @streaming_mode
+
+          payload
+        end
+
+        def prepare_request(payload)
+          headers = {
+            "anthropic-version" => "2023-06-01",
+            "x-api-key" => SiteSetting.ai_anthropic_api_key,
+            "content-type" => "application/json",
+          }
+
+          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
+        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
+
+          result
+        end
+
+        def partials_from(decoded_chunk)
+          decoded_chunk.split("\n").map { |line| line.split("data: ", 2)[1] }.compact
+        end
+      end
+    end
+  end
+end

lib/completions/endpoints/base.rb

@@ -16,6 +16,7 @@ module DiscourseAi
               DiscourseAi::Completions::Endpoints::HuggingFace,
               DiscourseAi::Completions::Endpoints::Gemini,
               DiscourseAi::Completions::Endpoints::Vllm,
+              DiscourseAi::Completions::Endpoints::AnthropicMessages,
             ]
 
             if Rails.env.test? || Rails.env.development?
@@ -165,8 +166,9 @@ module DiscourseAi
 
                     begin
                       partial = extract_completion_from(raw_partial)
-                      next if response_data.empty? && partial.blank?
                       next if partial.nil?
+                      # empty vs blank... we still accept " "
+                      next if response_data.empty? && partial.empty?
                       partials_raw << partial.to_s
 
                       # Stop streaming the response as soon as you find a tool.
@@ -213,6 +215,7 @@ module DiscourseAi
               if log
                 log.raw_response_payload = response_raw
                 log.response_tokens = tokenizer.size(partials_raw)
+                final_log_update(log)
                 log.save!
 
                 if Rails.env.development?
@@ -223,6 +226,10 @@ module DiscourseAi
           end
         end
 
+        def final_log_update(log)
+          # for people that need to override
+        end
+
         def default_options
           raise NotImplementedError
         end

lib/completions/llm.rb

@@ -24,7 +24,7 @@ module DiscourseAi
           @models_by_provider ||=
             {
               aws_bedrock: %w[claude-instant-1 claude-2],
-              anthropic: %w[claude-instant-1 claude-2],
+              anthropic: %w[claude-instant-1 claude-2 claude-3-sonnet claude-3-opus],
               vllm: %w[
                 mistralai/Mixtral-8x7B-Instruct-v0.1
                 mistralai/Mistral-7B-Instruct-v0.2

spec/lib/completions/dialects/claude_messages_spec.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+RSpec.describe DiscourseAi::Completions::Dialects::ClaudeMessages do
+  describe "#translate" do
+    it "can insert OKs to make stuff interleve properly" do
+      messages = [
+        { type: :user, id: "user1", content: "1" },
+        { type: :model, content: "2" },
+        { type: :user, id: "user1", content: "4" },
+        { type: :user, id: "user1", content: "5" },
+        { type: :model, content: "6" },
+      ]
+
+      prompt = DiscourseAi::Completions::Prompt.new("You are a helpful bot", messages: messages)
+
+      dialectKlass = DiscourseAi::Completions::Dialects::Dialect.dialect_for("claude-3-opus")
+      dialect = dialectKlass.new(prompt, "claude-3-opus")
+      translated = dialect.translate
+
+      expected_messages = [
+        { role: "user", content: "user1: 1" },
+        { role: "assistant", content: "2" },
+        { role: "user", content: "user1: 4" },
+        { role: "assistant", content: "OK" },
+        { role: "user", content: "user1: 5" },
+        { role: "assistant", content: "6" },
+      ]
+
+      expect(translated.messages).to eq(expected_messages)
+    end
+
+    it "can properly translate a prompt" do
+      dialect = DiscourseAi::Completions::Dialects::Dialect.dialect_for("claude-3-opus")
+
+      tools = [
+        {
+          name: "echo",
+          description: "echo a string",
+          parameters: [
+            { name: "text", type: "string", description: "string to echo", required: true },
+          ],
+        },
+      ]
+
+      tool_call_prompt = { name: "echo", arguments: { text: "something" } }
+
+      messages = [
+        { type: :user, id: "user1", content: "echo something" },
+        { type: :tool_call, content: tool_call_prompt.to_json },
+        { type: :tool, id: "tool_id", content: "something".to_json },
+        { type: :model, content: "I did it" },
+        { type: :user, id: "user1", content: "echo something else" },
+      ]
+
+      prompt =
+        DiscourseAi::Completions::Prompt.new(
+          "You are a helpful bot",
+          messages: messages,
+          tools: tools,
+        )
+
+      dialect = dialect.new(prompt, "claude-3-opus")
+      translated = dialect.translate
+
+      expect(translated.system_prompt).to start_with("You are a helpful bot")
+      expect(translated.system_prompt).to include("echo a string")
+
+      expected = [
+        { role: "user", content: "user1: echo something" },
+        {
+          role: "assistant",
+          content:
+            "<function_calls>\n<invoke>\n<tool_name>echo</tool_name>\n<parameters>\n<text>something</text>\n</parameters>\n</invoke>\n</function_calls>",
+        },
+        {
+          role: "user",
+          content:
+            "<function_results>\n<result>\n<tool_name>tool_id</tool_name>\n<json>\n\"something\"\n</json>\n</result>\n</function_results>",
+        },
+        { role: "assistant", content: "I did it" },
+        { role: "user", content: "user1: echo something else" },
+      ]
+
+      expect(translated.messages).to eq(expected)
+    end
+  end
+end

spec/lib/completions/endpoints/anthropic_messages_spec.rb

@@ -0,0 +1,301 @@
+# frozen_string_literal: true
+
+RSpec.describe DiscourseAi::Completions::Endpoints::AnthropicMessages do
+  let(:llm) { DiscourseAi::Completions::Llm.proxy("anthropic:claude-3-opus") }
+
+  let(:prompt) do
+    DiscourseAi::Completions::Prompt.new(
+      "You are hello bot",
+      messages: [type: :user, id: "user1", content: "hello"],
+    )
+  end
+
+  before { SiteSetting.ai_anthropic_api_key = "123" }
+
+  it "does not eat spaces with tool calls" do
+    body = <<~STRING
+      event: message_start
+      data: {"type":"message_start","message":{"id":"msg_019kmW9Q3GqfWmuFJbePJTBR","type":"message","role":"assistant","content":[],"model":"claude-3-opus-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":347,"output_tokens":1}}}
+
+      event: content_block_start
+      data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+
+      event: ping
+      data: {"type": "ping"}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"<function"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"_"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"calls"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"\\n"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"<invoke"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"\\n"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"<tool"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"_"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"name"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"google"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"</tool"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"_"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"name"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"\\n"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"<parameters"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"\\n"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"<query"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"top"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" "}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"10"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" "}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"things"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" to"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" do"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" in"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" japan"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" for"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" tourists"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"</query"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"\\n"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"</parameters"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"\\n"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"</invoke"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":">"}}
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"\\n"}}
+
+      event: content_block_stop
+      data: {"type":"content_block_stop","index":0}
+
+      event: message_delta
+      data: {"type":"message_delta","delta":{"stop_reason":"stop_sequence","stop_sequence":"</function_calls>"},"usage":{"output_tokens":57}}
+
+      event: message_stop
+      data: {"type":"message_stop"}
+    STRING
+
+    stub_request(:post, "https://api.anthropic.com/v1/messages").to_return(status: 200, body: body)
+
+    result = +""
+    llm.generate(prompt, user: Discourse.system_user) { |partial| result << partial }
+
+    expected = (<<~TEXT).strip
+      <function_calls>
+      <invoke>
+      <tool_name>google</tool_name>
+      <tool_id>google</tool_id>
+      <parameters>
+      <query>top 10 things to do in japan for tourists</query>
+      </parameters>
+      </invoke>
+      </function_calls>
+    TEXT
+
+    expect(result.strip).to eq(expected)
+  end
+
+  it "can stream a response" do
+    body = (<<~STRING).strip
+      event: message_start
+      data: {"type": "message_start", "message": {"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY", "type": "message", "role": "assistant", "content": [], "model": "claude-3-opus-20240229", "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 25, "output_tokens": 1}}}
+
+      event: content_block_start
+      data: {"type": "content_block_start", "index":0, "content_block": {"type": "text", "text": ""}}
+
+      event: ping
+      data: {"type": "ping"}
+
+      event: content_block_delta
+      data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}
+
+      event: content_block_delta
+      data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "!"}}
+
+      event: content_block_stop
+      data: {"type": "content_block_stop", "index": 0}
+
+      event: message_delta
+      data: {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence":null, "usage":{"output_tokens": 15}}}
+
+      event: message_stop
+      data: {"type": "message_stop"}
+    STRING
+
+    parsed_body = nil
+
+    stub_request(:post, "https://api.anthropic.com/v1/messages").with(
+      body:
+        proc do |req_body|
+          parsed_body = JSON.parse(req_body, symbolize_names: true)
+          true
+        end,
+      headers: {
+        "Content-Type" => "application/json",
+        "X-Api-Key" => "123",
+        "Anthropic-Version" => "2023-06-01",
+      },
+    ).to_return(status: 200, body: body)
+
+    result = +""
+    llm.generate(prompt, user: Discourse.system_user) { |partial, cancel| result << partial }
+
+    expect(result).to eq("Hello!")
+
+    expected_body = {
+      model: "claude-3-opus-20240229",
+      max_tokens: 3000,
+      stop_sequences: ["</function_calls>"],
+      messages: [{ role: "user", content: "user1: hello" }],
+      system: "You are hello bot",
+      stream: true,
+    }
+    expect(parsed_body).to eq(expected_body)
+
+    log = AiApiAuditLog.order(:id).last
+    expect(log.provider_id).to eq(AiApiAuditLog::Provider::Anthropic)
+    expect(log.request_tokens).to eq(25)
+    expect(log.response_tokens).to eq(15)
+  end
+
+  it "can operate in regular mode" do
+    body = <<~STRING
+      {
+        "content": [
+          {
+            "text": "Hello!",
+            "type": "text"
+          }
+        ],
+        "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF",
+        "model": "claude-3-opus-20240229",
+        "role": "assistant",
+        "stop_reason": "end_turn",
+        "stop_sequence": null,
+        "type": "message",
+        "usage": {
+          "input_tokens": 10,
+          "output_tokens": 25
+        }
+      }
+    STRING
+
+    parsed_body = nil
+    stub_request(:post, "https://api.anthropic.com/v1/messages").with(
+      body:
+        proc do |req_body|
+          parsed_body = JSON.parse(req_body, symbolize_names: true)
+          true
+        end,
+      headers: {
+        "Content-Type" => "application/json",
+        "X-Api-Key" => "123",
+        "Anthropic-Version" => "2023-06-01",
+      },
+    ).to_return(status: 200, body: body)
+
+    result = llm.generate(prompt, user: Discourse.system_user)
+    expect(result).to eq("Hello!")
+
+    expected_body = {
+      model: "claude-3-opus-20240229",
+      max_tokens: 3000,
+      stop_sequences: ["</function_calls>"],
+      messages: [{ role: "user", content: "user1: hello" }],
+      system: "You are hello bot",
+    }
+    expect(parsed_body).to eq(expected_body)
+
+    log = AiApiAuditLog.order(:id).last
+    expect(log.provider_id).to eq(AiApiAuditLog::Provider::Anthropic)
+    expect(log.request_tokens).to eq(10)
+    expect(log.response_tokens).to eq(25)
+  end
+end

spec/lib/completions/endpoints/open_ai_spec.rb

@@ -147,7 +147,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::OpenAi do
     described_class.new("gpt-3.5-turbo", DiscourseAi::Tokenizer::OpenAiTokenizer)
   end
 
-  fab!(:user) { Fabricate(:user) }
+  fab!(:user)
 
   let(:open_ai_mock) { OpenAiMock.new(endpoint) }
 

spec/lib/modules/ai_bot/playground_spec.rb

@@ -412,7 +412,7 @@ RSpec.describe DiscourseAi::AiBot::Playground do
         )
       end
 
-      it "include replies generated from tools only once" do
+      it "include replies generated from tools" do
         custom_prompt = [
           [
             { args: { timezone: "Buenos Aires" }, time: "2023-12-14 17:24:00 -0300" }.to_json,
@@ -424,7 +424,7 @@ RSpec.describe DiscourseAi::AiBot::Playground do
             "time",
             "tool_call",
           ],
-          ["I replied this thanks to the time command", bot_user.username],
+          ["I replied", bot_user.username],
         ]
         PostCustomPrompt.create!(post: second_post, custom_prompt: custom_prompt)
         PostCustomPrompt.create!(post: first_post, custom_prompt: custom_prompt)
@@ -439,6 +439,7 @@ RSpec.describe DiscourseAi::AiBot::Playground do
             { type: :tool, id: "time", content: custom_prompt.first.first },
             { type: :tool_call, content: custom_prompt.second.first, id: "time" },
             { type: :tool, id: "time", content: custom_prompt.first.first },
+            { type: :model, content: "I replied" },
           ],
         )
       end