From 6f5f34184b51703ba7339afe1160a75221439583 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Wed, 3 Apr 2024 16:06:27 +1100
Subject: [PATCH] FEATURE: add Claude 3 Haiku bot support (#552)

it is close in performance to GPT 4 at a fraction of the cost,
nice to add it to the mix.

Also improves a test case to simulate streaming, I am hunting for
the "calls" word that is jumping into function calls and can't quite
find it.
---
 config/locales/client.en.yml                      |  3 ++-
 config/settings.yml                               |  1 +
 lib/ai_bot/bot.rb                                 |  6 ++++++
 lib/ai_bot/entry_point.rb                         |  4 ++++
 spec/lib/completions/endpoints/anthropic_spec.rb  | 15 +++++++++++----
 .../completions/endpoints/endpoint_compliance.rb  |  4 ++++
 6 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml
index fd3e6754..2b61b6fb 100644
--- a/config/locales/client.en.yml
+++ b/config/locales/client.en.yml
@@ -164,7 +164,7 @@ en:
           #### Group-Specific Access to AI Personas
 
           Moreover, you can set it up so that certain user groups have access to specific personas. This means you can have different AI behaviors for different sections of your forum, further enhancing the diversity and richness of your community's interactions.
-        
+
         uploads:
           title: "Uploads"
           description: "Your AI persona will be able to search and reference the content of included files. Uploaded files must be formatted as plaintext (.txt)"
@@ -265,6 +265,7 @@ en:
           fake: "Fake Test Bot"
           claude-3-opus: "Claude 3 Opus"
           claude-3-sonnet: "Claude 3 Sonnet"
+          claude-3-haiku: "Claude 3 Haiku"
           gpt-4: "GPT-4"
           gpt-4-turbo: "GPT-4 Turbo"
           gpt-3:
diff --git a/config/settings.yml b/config/settings.yml
index cdb28e14..6df71888 100644
--- a/config/settings.yml
+++ b/config/settings.yml
@@ -334,6 +334,7 @@ discourse_ai:
       - mixtral-8x7B-Instruct-V0.1
       - claude-3-opus
       - claude-3-sonnet
+      - claude-3-haiku
   ai_bot_add_to_header:
     default: true
     client: true
diff --git a/lib/ai_bot/bot.rb b/lib/ai_bot/bot.rb
index 40595b67..0f1a553a 100644
--- a/lib/ai_bot/bot.rb
+++ b/lib/ai_bot/bot.rb
@@ -188,6 +188,12 @@ module DiscourseAi
           else
             "anthropic:claude-3-sonnet"
           end
+        when DiscourseAi::AiBot::EntryPoint::CLAUDE_3_HAIKU_ID
+          if DiscourseAi::Completions::Endpoints::AwsBedrock.correctly_configured?("claude-3-haiku")
+            "aws_bedrock:claude-3-haiku"
+          else
+            "anthropic:claude-3-haiku"
+          end
         else
           nil
         end
diff --git a/lib/ai_bot/entry_point.rb b/lib/ai_bot/entry_point.rb
index 88c7ef02..928753d9 100644
--- a/lib/ai_bot/entry_point.rb
+++ b/lib/ai_bot/entry_point.rb
@@ -16,6 +16,7 @@ module DiscourseAi
       FAKE_ID = -116 # only used for dev and test
       CLAUDE_3_OPUS_ID = -117
       CLAUDE_3_SONNET_ID = -118
+      CLAUDE_3_HAIKU_ID = -119
 
       BOTS = [
         [GPT4_ID, "gpt4_bot", "gpt-4"],
@@ -27,6 +28,7 @@ module DiscourseAi
         [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"],
+        [CLAUDE_3_HAIKU_ID, "claude_3_haiku_bot", "claude-3-haiku"],
       ]
 
       BOT_USER_IDS = BOTS.map(&:first)
@@ -63,6 +65,8 @@ module DiscourseAi
           CLAUDE_3_OPUS_ID
         in "claude-3-sonnet"
           CLAUDE_3_SONNET_ID
+        in "claude-3-haiku"
+          CLAUDE_3_HAIKU_ID
         else
           nil
         end
diff --git a/spec/lib/completions/endpoints/anthropic_spec.rb b/spec/lib/completions/endpoints/anthropic_spec.rb
index 158f9fa4..dffbab95 100644
--- a/spec/lib/completions/endpoints/anthropic_spec.rb
+++ b/spec/lib/completions/endpoints/anthropic_spec.rb
@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+require_relative "endpoint_compliance"
 
 RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do
   let(:llm) { DiscourseAi::Completions::Llm.proxy("anthropic:claude-3-opus") }
@@ -196,11 +197,17 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do
       data: {"type":"message_stop"}
     STRING
 
-    stub_request(:post, "https://api.anthropic.com/v1/messages").to_return(status: 200, body: body)
-
     result = +""
-    llm.generate(prompt_with_google_tool, user: Discourse.system_user) do |partial|
-      result << partial
+    body = body.scan(/.*\n/)
+    EndpointMock.with_chunk_array_support do
+      stub_request(:post, "https://api.anthropic.com/v1/messages").to_return(
+        status: 200,
+        body: body,
+      )
+
+      llm.generate(prompt_with_google_tool, user: Discourse.system_user) do |partial|
+        result << partial
+      end
     end
 
     expected = (<<~TEXT).strip
diff --git a/spec/lib/completions/endpoints/endpoint_compliance.rb b/spec/lib/completions/endpoints/endpoint_compliance.rb
index 31f5e86e..82c63ecf 100644
--- a/spec/lib/completions/endpoints/endpoint_compliance.rb
+++ b/spec/lib/completions/endpoints/endpoint_compliance.rb
@@ -108,6 +108,10 @@ class EndpointMock
     ::FinalDestination.send(:const_set, :HTTP, @original_net_http)
   end
 
+  def self.with_chunk_array_support(&blk)
+    self.new(nil).with_chunk_array_support(&blk)
+  end
+
   protected
 
   # Copied from https://github.com/bblimke/webmock/issues/629