From 501fdf0eb15dd2bafed82466fe7c13fe03c19533 Mon Sep 17 00:00:00 2001 From: Paul Shippy Date: Thu, 27 Feb 2025 13:53:11 -0700 Subject: [PATCH 1/3] Support Amazon Nova models --- lib/langchain/llm/aws_bedrock.rb | 18 +++ .../response/aws_bedrock_amazon_response.rb | 37 +++++ spec/langchain/llm/aws_bedrock_spec.rb | 132 ++++++++++++++++++ 3 files changed, 187 insertions(+) create mode 100644 lib/langchain/llm/response/aws_bedrock_amazon_response.rb diff --git a/lib/langchain/llm/aws_bedrock.rb b/lib/langchain/llm/aws_bedrock.rb index 39f487aca..f7e4860f5 100644 --- a/lib/langchain/llm/aws_bedrock.rb +++ b/lib/langchain/llm/aws_bedrock.rb @@ -25,6 +25,7 @@ class AwsBedrock < Base attr_reader :client, :defaults SUPPORTED_COMPLETION_PROVIDERS = %i[ + amazon anthropic ai21 cohere @@ -32,6 +33,7 @@ class AwsBedrock < Base ].freeze SUPPORTED_CHAT_COMPLETION_PROVIDERS = %i[ + amazon anthropic ai21 mistral @@ -216,6 +218,8 @@ def compose_parameters(params, model_id) params elsif provider_name(model_id) == :mistral params + elsif provider_name(model_id) == :amazon + compose_parameters_amazon(params) end end @@ -238,6 +242,8 @@ def parse_response(response, model_id) Langchain::LLM::AwsBedrockMetaResponse.new(JSON.parse(response.body.string)) elsif provider_name(model_id) == :mistral Langchain::LLM::MistralAIResponse.new(JSON.parse(response.body.string)) + elsif provider_name(model_id) == :amazon + Langchain::LLM::AwsBedrockAmazonResponse.new(JSON.parse(response.body.string)) end end @@ -288,6 +294,18 @@ def compose_parameters_anthropic(params) params.merge(anthropic_version: "bedrock-2023-05-31") end + def compose_parameters_amazon(params) + params = params.merge(inferenceConfig: { + maxTokens: params[:max_tokens], + temperature: params[:temperature], + topP: params[:top_p], + topK: params[:top_k], + stopSequences: params[:stop_sequences] + }.compact) + + params.reject { |k, _| k == :max_tokens || k == :temperature } + end + def response_from_chunks(chunks) raw_response = {} diff --git a/lib/langchain/llm/response/aws_bedrock_amazon_response.rb b/lib/langchain/llm/response/aws_bedrock_amazon_response.rb new file mode 100644 index 000000000..c2c9d5769 --- /dev/null +++ b/lib/langchain/llm/response/aws_bedrock_amazon_response.rb @@ -0,0 +1,37 @@ +# frozen_string_literal: true + +module Langchain::LLM + class AwsBedrockAmazonResponse < BaseResponse + def completion + raw_response.dig("output", "message", "content", 0, "text") + end + + def chat_completion + completion + end + + def chat_completions + completions + end + + def completions + nil + end + + def stop_reason + raw_response.dig("stopReason") + end + + def prompt_tokens + raw_response.dig("usage", "inputTokens").to_i + end + + def completion_tokens + raw_response.dig("usage", "outputTokens").to_i + end + + def total_tokens + raw_response.dig("usage", "totalTokens").to_i + end + end +end diff --git a/spec/langchain/llm/aws_bedrock_spec.rb b/spec/langchain/llm/aws_bedrock_spec.rb index 5e20789c2..70bc64c59 100644 --- a/spec/langchain/llm/aws_bedrock_spec.rb +++ b/spec/langchain/llm/aws_bedrock_spec.rb @@ -119,6 +119,138 @@ end end end + + context "with amazon nova pro provider" do + let(:response) do + { + output: { + message: { + content: [ + {text: "The capital of France is Paris."} + ] + } + }, + usage: {inputTokens: 14, outputTokens: 10} + }.to_json + end + + let(:model_id) { "amazon.nova-pro-v1:0" } + + before do + response_object = double("response_object") + allow(response_object).to receive(:body).and_return(StringIO.new(response)) + allow(subject.client).to receive(:invoke_model) + .with(matching( + model_id:, + body: { + messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], + inferenceConfig: { + maxTokens: 300 + } + }.to_json, + content_type: "application/json", + accept: "application/json" + )) + .and_return(response_object) + end + + it "returns a completion" do + expect( + subject.chat( + messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], + model: model_id + ).chat_completion + ).to eq("The capital of France is Paris.") + end + end + + context "with amazon nova lite provider" do + let(:response) do + { + output: { + message: { + content: [ + {text: "The capital of France is Paris."} + ] + } + }, + usage: {inputTokens: 14, outputTokens: 10} + }.to_json + end + + let(:model_id) { "amazon.nova-lite-v1:0" } + + before do + response_object = double("response_object") + allow(response_object).to receive(:body).and_return(StringIO.new(response)) + allow(subject.client).to receive(:invoke_model) + .with(matching( + model_id:, + body: { + messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], + inferenceConfig: { + maxTokens: 300 + } + }.to_json, + content_type: "application/json", + accept: "application/json" + )) + .and_return(response_object) + end + + it "returns a completion" do + expect( + subject.chat( + messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], + model: model_id + ).chat_completion + ).to eq("The capital of France is Paris.") + end + end + + context "with amazon nova micro provider" do + let(:response) do + { + output: { + message: { + content: [ + {text: "The capital of France is Paris."} + ] + } + }, + usage: {inputTokens: 14, outputTokens: 10} + }.to_json + end + + let(:model_id) { "amazon.nova-micro-v1:0" } + + before do + response_object = double("response_object") + allow(response_object).to receive(:body).and_return(StringIO.new(response)) + allow(subject.client).to receive(:invoke_model) + .with(matching( + model_id:, + body: { + messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], + inferenceConfig: { + maxTokens: 300 + } + }.to_json, + content_type: "application/json", + accept: "application/json" + )) + .and_return(response_object) + end + + it "returns a completion" do + expect( + subject.chat( + messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], + model: model_id + ).chat_completion + ).to eq("The capital of France is Paris.") + end + end end describe "#complete" do From 43739b7765822f795235ed82e286312d1dc66b41 Mon Sep 17 00:00:00 2001 From: Paul Shippy Date: Thu, 27 Feb 2025 13:56:32 -0700 Subject: [PATCH 2/3] Simplify spec to just one model --- spec/langchain/llm/aws_bedrock_spec.rb | 90 +------------------------- 1 file changed, 1 insertion(+), 89 deletions(-) diff --git a/spec/langchain/llm/aws_bedrock_spec.rb b/spec/langchain/llm/aws_bedrock_spec.rb index 70bc64c59..9c81f8212 100644 --- a/spec/langchain/llm/aws_bedrock_spec.rb +++ b/spec/langchain/llm/aws_bedrock_spec.rb @@ -120,7 +120,7 @@ end end - context "with amazon nova pro provider" do + context "with amazon provider" do let(:response) do { output: { @@ -163,94 +163,6 @@ ).to eq("The capital of France is Paris.") end end - - context "with amazon nova lite provider" do - let(:response) do - { - output: { - message: { - content: [ - {text: "The capital of France is Paris."} - ] - } - }, - usage: {inputTokens: 14, outputTokens: 10} - }.to_json - end - - let(:model_id) { "amazon.nova-lite-v1:0" } - - before do - response_object = double("response_object") - allow(response_object).to receive(:body).and_return(StringIO.new(response)) - allow(subject.client).to receive(:invoke_model) - .with(matching( - model_id:, - body: { - messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], - inferenceConfig: { - maxTokens: 300 - } - }.to_json, - content_type: "application/json", - accept: "application/json" - )) - .and_return(response_object) - end - - it "returns a completion" do - expect( - subject.chat( - messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], - model: model_id - ).chat_completion - ).to eq("The capital of France is Paris.") - end - end - - context "with amazon nova micro provider" do - let(:response) do - { - output: { - message: { - content: [ - {text: "The capital of France is Paris."} - ] - } - }, - usage: {inputTokens: 14, outputTokens: 10} - }.to_json - end - - let(:model_id) { "amazon.nova-micro-v1:0" } - - before do - response_object = double("response_object") - allow(response_object).to receive(:body).and_return(StringIO.new(response)) - allow(subject.client).to receive(:invoke_model) - .with(matching( - model_id:, - body: { - messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], - inferenceConfig: { - maxTokens: 300 - } - }.to_json, - content_type: "application/json", - accept: "application/json" - )) - .and_return(response_object) - end - - it "returns a completion" do - expect( - subject.chat( - messages: [{role: "user", content: [{text: "What is the capital of France?"}]}], - model: model_id - ).chat_completion - ).to eq("The capital of France is Paris.") - end - end end describe "#complete" do From e8deacdbcde3a6ac97324f91a6e5be87d20873e3 Mon Sep 17 00:00:00 2001 From: Paul Shippy Date: Thu, 27 Feb 2025 14:01:29 -0700 Subject: [PATCH 3/3] Remove amazon from completion providers - Only tested Chat Completion --- lib/langchain/llm/aws_bedrock.rb | 1 - 1 file changed, 1 deletion(-) diff --git a/lib/langchain/llm/aws_bedrock.rb b/lib/langchain/llm/aws_bedrock.rb index f7e4860f5..1eb056355 100644 --- a/lib/langchain/llm/aws_bedrock.rb +++ b/lib/langchain/llm/aws_bedrock.rb @@ -25,7 +25,6 @@ class AwsBedrock < Base attr_reader :client, :defaults SUPPORTED_COMPLETION_PROVIDERS = %i[ - amazon anthropic ai21 cohere