FIX: allow the correct number of tokens for AWS models

Author: SamSaffron

lib/completions/endpoints/aws_bedrock.rb

@@ -19,7 +19,10 @@ def normalize_model_params(model_params)
         end
 
         def default_options(dialect)
-          options = { max_tokens: 3_000, anthropic_version: "bedrock-2023-05-31" }
+          max_tokens = 4096
+          max_tokens = 8192 if bedrock_model_id.match?(/3.5/)
+
+          options = { max_tokens: max_tokens, anthropic_version: "bedrock-2023-05-31" }
 
           options[:stop_sequences] = ["</function_calls>"] if !dialect.native_tool_support? &&
             dialect.prompt.has_tools?
@@ -40,6 +43,27 @@ def xml_tags_to_strip(dialect)
 
         private
 
+        def bedrock_model_id
+          case llm_model.name
+          when "claude-2"
+            "anthropic.claude-v2:1"
+          when "claude-3-haiku"
+            "anthropic.claude-3-haiku-20240307-v1:0"
+          when "claude-3-sonnet"
+            "anthropic.claude-3-sonnet-20240229-v1:0"
+          when "claude-instant-1"
+            "anthropic.claude-instant-v1"
+          when "claude-3-opus"
+            "anthropic.claude-3-opus-20240229-v1:0"
+          when "claude-3-5-sonnet"
+            "anthropic.claude-3-5-sonnet-20241022-v2:0"
+          when "claude-3-5-haiku"
+            "anthropic.claude-3-5-haiku-20241022-v1:0"
+          else
+            llm_model.name
+          end
+        end
+
         def prompt_size(prompt)
           # approximation
           tokenizer.size(prompt.system_prompt.to_s + " " + prompt.messages.to_s)
@@ -48,24 +72,6 @@ def prompt_size(prompt)
         def model_uri
           region = llm_model.lookup_custom_param("region")
 
-          bedrock_model_id =
-            case llm_model.name
-            when "claude-2"
-              "anthropic.claude-v2:1"
-            when "claude-3-haiku"
-              "anthropic.claude-3-haiku-20240307-v1:0"
-            when "claude-3-sonnet"
-              "anthropic.claude-3-sonnet-20240229-v1:0"
-            when "claude-instant-1"
-              "anthropic.claude-instant-v1"
-            when "claude-3-opus"
-              "anthropic.claude-3-opus-20240229-v1:0"
-            when "claude-3-5-sonnet"
-              "anthropic.claude-3-5-sonnet-20241022-v2:0"
-            else
-              llm_model.name
-            end
-
           if region.blank? || bedrock_model_id.blank?
             raise CompletionFailed.new(I18n.t("discourse_ai.llm_models.bedrock_invalid_url"))
           end

spec/lib/completions/endpoints/aws_bedrock_spec.rb

@@ -26,6 +26,22 @@ def encode_message(message)
     Aws::EventStream::Encoder.new.encode(aws_message)
   end
 
+  it "should provide accurate max token count" do
+    prompt = DiscourseAi::Completions::Prompt.new("hello")
+    dialect = DiscourseAi::Completions::Dialects::Claude.new(prompt, model)
+    endpoint = DiscourseAi::Completions::Endpoints::AwsBedrock.new(model)
+
+    model.name = "claude-2"
+    expect(endpoint.default_options(dialect)[:max_tokens]).to eq(4096)
+
+    model.name = "claude-3-5-sonnet"
+    expect(endpoint.default_options(dialect)[:max_tokens]).to eq(8192)
+
+    model.name = "claude-3-5-haiku"
+    options = endpoint.default_options(dialect)
+    expect(options[:max_tokens]).to eq(8192)
+  end
+
   describe "function calling" do
     it "supports old school xml function calls" do
       model.provider_params["disable_native_tools"] = true
@@ -246,7 +262,7 @@ def encode_message(message)
         expect(response).to eq(expected_response)
 
         expected = {
-          "max_tokens" => 3000,
+          "max_tokens" => 4096,
           "anthropic_version" => "bedrock-2023-05-31",
           "messages" => [{ "role" => "user", "content" => "what is the weather in sydney" }],
           "tools" => [
@@ -305,7 +321,7 @@ def encode_message(message)
       expect(request.headers["X-Amz-Content-Sha256"]).to be_present
 
       expected = {
-        "max_tokens" => 3000,
+        "max_tokens" => 4096,
         "anthropic_version" => "bedrock-2023-05-31",
         "messages" => [{ "role" => "user", "content" => "hello world" }],
         "system" => "You are a helpful bot",
@@ -354,7 +370,7 @@ def encode_message(message)
         expect(request.headers["X-Amz-Content-Sha256"]).to be_present
 
         expected = {
-          "max_tokens" => 3000,
+          "max_tokens" => 4096,
           "anthropic_version" => "bedrock-2023-05-31",
           "messages" => [{ "role" => "user", "content" => "hello world" }],
           "system" => "You are a helpful bot",