ruby-amazon-bedrock 0.2.0 → 0.2.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5cc2fda0ffade519b17acc4d278b6d4d55136c8df511a2da329dccbc3dc63bd9
-  data.tar.gz: c6afb32b619515b886b5cb34ed01791a62ea1fda119407d8b72ce1455d65f149
+  metadata.gz: e34184a2246b1054a64d6d29bc7bb1144487a508c6d27f717d8f92edb4cb32a2
+  data.tar.gz: ae40e48b1107bf163bbb30e4c058d660784f1bf9cc6fcafa0bfa9843b526dd49
 SHA512:
-  metadata.gz: bdadc9d45b3fac52916462f5e47e62cb91e4f25f949d4af9b7fe3ecca9a7bb151bdcfb0837c2fd511115affd49845441a6a964a0175cc1748583186f053de1e2
-  data.tar.gz: 90b2adc4a2f8f49237803f81663ea3d4a54b0177bbd6075f6a86f8982dad6a23cac6266dc50cf92aa4902cef9241f9bf1f4ec22dbf6bf8b98269daa9d1102834
+  metadata.gz: fcdfeb1017d3bf2cefeb1bf77b68a85bd36bd3a7906651b1a56f8ff2544309ba98eee08fc1f04ce80ed22b738302fd38e7811087f56a976c6eeea1f7cfae0010
+  data.tar.gz: f19b95ef5965317c53eda21c96c69e421abdb34ce441cc850731663aa8ed35bbc67000d02a34ebc2f192ce7fbe81e0d18afe3ab59d0dbea38e32c858625fcc89

data/CHANGELOG.md

@@ -1,3 +1,14 @@
+## [0.2.2] - 2024-01-10
+
+- Fix bug for Amazon Titan Text Lite
+- Allow payload builders to receive and set Model parameters as options
+
+## [0.2.1] - 2023-12-25
+
+- Make gem configurable
+- Fix A121 Labs payloads
+- Update README.md with A121 Labs documentation
+
 ## [0.2.0] - 2023-12-14
 
 - Add Text and Image builders to be able to return text or save images depending on the model use case.

data/Gemfile

@@ -11,6 +11,5 @@ gem "pry", "~> 0.13"
 gem "rake", "~> 13.0"
 gem "rspec", "~> 3.0"
 gem "rubocop", "~> 1.21"
-gem "simplecov", "~> 0.21"
 gem "vcr", "~> 6.0"
 gem "webmock", "~> 3.12"

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    ruby-amazon-bedrock (0.2.0)
+    ruby-amazon-bedrock (0.2.2)
       aws-sdk-bedrockruntime (~> 1.0)
 
 GEM
@@ -29,7 +29,6 @@ GEM
     crack (0.4.5)
       rexml
     diff-lcs (1.5.0)
-    docile (1.4.0)
     dotenv (2.8.1)
     hashdiff (1.0.1)
     jmespath (1.6.2)
@@ -76,12 +75,6 @@ GEM
     rubocop-ast (1.30.0)
       parser (>= 3.2.1.0)
     ruby-progressbar (1.13.0)
-    simplecov (0.22.0)
-      docile (~> 1.1)
-      simplecov-html (~> 0.11)
-      simplecov_json_formatter (~> 0.1)
-    simplecov-html (0.12.3)
-    simplecov_json_formatter (0.1.4)
     thor (1.3.0)
     unicode-display_width (2.5.0)
     vcr (6.2.0)
@@ -101,7 +94,6 @@ DEPENDENCIES
   rspec (~> 3.0)
   rubocop (~> 1.21)
   ruby-amazon-bedrock!
-  simplecov (~> 0.21)
   vcr (~> 6.0)
   webmock (~> 3.12)
 

data/README.md

@@ -1,6 +1,7 @@
 # Ruby Amazon Bedrock
 
-[![CircleCI](https://dl.circleci.com/status-badge/img/circleci/JP4R1PDyZ2Yax5GxVQoMZN/LAwrE89wMVTrgpfzaimon/tree/main.svg?style=svg)](https://dl.circleci.com/status-badge/redirect/circleci/JP4R1PDyZ2Yax5GxVQoMZN/LAwrE89wMVTrgpfzaimon/tree/main)
+[![Gem Version](https://badge.fury.io/rb/ruby-amazon-bedrock.svg)](https://badge.fury.io/rb/ruby-amazon-bedrock)
+[![CircleCI](https://dl.circleci.com/status-badge/img/circleci/JP4R1PDyZ2Yax5GxVQoMZN/LAwrE89wMVTrgpfzaimon/tree/main.svg?style=shield&circle-token=31edf35e4c3f323ce0044c6fceadf9ffb15f5d3c)](https://dl.circleci.com/status-badge/redirect/circleci/JP4R1PDyZ2Yax5GxVQoMZN/LAwrE89wMVTrgpfzaimon/tree/main)
 
 Seamless Integration with Amazon Bedrock API for AI-Powered Text and Image Generation in Ruby 🤖 + 💎. [Amazon Bedrock API](https://aws.amazon.com/es/bedrock/).
 
@@ -13,7 +14,7 @@ Amazon Bedrock is a fully managed service that makes FMs from leading AI startup
 Add the following line to your application's Gemfile:
 
 ```ruby
-gem "ruby-amazon-bedrock"
+gem "ruby-amazon-bedrock", "~> 0.2.2"
 ```
 
 And then execute:
@@ -27,7 +28,7 @@ $ bundle install
 Or install with:
 
 ```bash
-$ gem install ruby-amazon-bedrock
+$ gem install ruby-amazon-bedrock -v 0.2.1
 ```
 
 and require with:
@@ -55,39 +56,116 @@ Be aware that by using `ruby-amazon-bedrock` gem in conjunction with Amazon Bedr
 Instantiate a client by passing your AWS IAM credentials:
 
 ```ruby
-client = RubyAmazonBedrock::Client.new(region: "aws_region", access_key_id: "access_key_id", access_token: "access_token")
+client = RubyAmazonBedrock::Client.new(
+  region: "AWS_REGION",
+  access_key_id: "AWS_ACCESS_KEY_ID",
+  access_token: "AWS_SECRET_ACCESS_KEY"
+)
 ```
 
-## Models Providers
+## With Configuration
+
+```ruby
+RubyAmazonBedrock.configure do |config|
+  config.region = ENV.fetch('AWS_REGION', nil)
+  config.access_key_id = ENV.fetch('AWS_ACCESS_KEY_ID', nil)
+  config.secret_access_key = ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+end
+
+client = RubyAmazonBedrock::Client.new
+```
+
+## Options
+
+The options argument enhances the method's flexibility by allowing additional, model-specific configurations. If not explicitly provided, the method defaults to a set of standard Model parameters. You can see more details about a Model's optional parameters.
+
+## Model Providers
 
 Amazon Bedrock is a fully managed service that makes FMs from leading AI startups and Amazon available via an API, so you can choose from a wide range of FMs to find the model that is best suited for your use case.
 
-<!-- ### AI21 Labs
+### AI21 Labs
 
 Businesses use AI21's Jurassic family of leading LLMs to build generative AI-driven applications and services leveraging existing organizational data. Jurassic supports cross-industry use cases including long and short-form text generation, contextual question answering, summarization, and classification. Designed to follow natural language instructions, Jurassic is trained on a massive corpus of web text and supports six languages in addition to English. [See more ...](https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=ai21.j2-ultra-v1)
 
+Options for these models are:
+
+```ruby
+  {
+    temperature: 0.5,     # float
+    max_tokens: 200,      # integer
+    stop_sequences: [],   # [string]
+    top_p: 0.5,           # float
+    count_penalty: 0,     # integer
+    presence_penalty: 0,  # integer
+    frequency_penalty: 0  # integer
+  }
+```
+
+For more documentation about the parameter data types and values you can [See more details...][https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-jurassic2.html]
+
 How to call AI21 models:
 
-Jurassic-2 Ultra
+_*Jurassic-2 Ultra*_
 
-Supported use cases: Open book question answering, summarization, draft generation, information extraction, ideation
+Supports: Open book question answering, summarization, draft generation, information extraction, ideation
 
 Languages: English, Spanish, French, German, Portuguese, Italian, Dutch
 
 ```ruby
-client.invoke_model('ai21.j2-ultra-v1', 'Your prompt goes here ...')
+client.invoke_model(id: 'ai21.j2-ultra-v1', prompt: "What's natural language processing?", options: {})
+
+# Response
+{:id=>1234,
+ :prompt=>
+  {:text=>"Describe how an LLM works",
+   :tokens=>
+    [{:generatedToken=>{:token=>"▁Describe", :logprob=>-10.674324035644531, :raw_logprob=>-10.674324035644531},
+ :completions=>
+  [{:data=>
+    {:text=>
+       "\nNatural language processing (NLP) is a field of computer science, artificial intelligence, and linguistics concerned with the interactions between computers and human (natural) languages, in particular how to program computers to process and analyze large amounts of natural language data. The goal of NLP is to provide computers with the ability to read, understand, and generate human language.",
+      :tokens=>
+       [{:generatedToken=>{:token=>"<|newline|>", :logprob=>0.0, :raw_logprob=>-0.00046850196667946875},
 ```
 
-Jurassic-2 Mid
+_*Jurassic-2 Mid*_
 
 ```ruby
-client.invoke_model('ai21.j2-mid-v1', 'Your prompt goes here ...')
-``` -->
+client.invoke_model(id: 'ai21.j2-mid-v1', prompt: "What's GenAI?", options: {})
+
+# Response
+{:id=>1234,
+ :prompt=>
+  {:text=>"What's GenAI?",
+   :tokens=>
+    [{:generatedToken=>{:token=>"▁What's", :logprob=>-9.553738594055176, :raw_logprob=>-9.553738594055176},
+      :topTokens=>nil,
+      :textRange=>{:start=>0, :end=>6}},
+ :completions=>
+  [{:data=>
+     {:text=>
+       "\nGenAI is a proposed standard for a generic AI language, which would allow AI systems to communicate and reason with each other in a common language. The goal of GenAI is to create a universal language that can be used by all AI systems, regardless of their specific task or underlying architecture. This would make it easier for AI systems to work together and learn from each other, and it would also make it easier for humans to interact with and understand AI systems.",
+      :tokens=>
+       [{:generatedToken=>{:token=>"<|newline|>", :logprob=>0.0, :raw_logprob=>-0.0009662011871114373},
+```
 
 ### Amazon Titan
 
 Amazon Titan Foundation Models are pre-trained on large datasets, making them powerful, general-purpose models. Use them as is, or customize them by fine tuning the models with your own data for a particular task without annotating large volumes of data. [See more ...](https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=amazon.titan-embed-text-v1)
 
+Options for these models are:
+
+```ruby
+  {
+    temperature: 0.5,   # float
+    top_p: 0.5,         # float
+    max_tokens: 512,    # integer
+    stop_sequences: []  # [string]
+  }
+```
+
+For more documentation about the parameter data types and values you can [See more details...][https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html]
+
 How to call Amazon Titan text models:
 
 _*Titan Text G1 - Lite*_
@@ -97,7 +175,7 @@ _*Titan Text G1 - Lite*_
 Supports: Text generation, Code generation, Rich text formatting, Orchestration (Agents), Fine Tuning.
 
 ```ruby
-client.invoke_model(id: 'amazon.titan-text-lite-v1', input: 'Generate a story about rubies and gems')
+client.invoke_model(id: 'amazon.titan-text-lite-v1', prompt: 'Generate a story about rubies and gems', options: {})
 
 # Response
 {:inputTextTokenCount=>8,
@@ -108,14 +186,12 @@ client.invoke_model(id: 'amazon.titan-text-lite-v1', input: 'Generate a story ab
     :completionReason=>"FINISH"}]}
 ```
 
-<br>
-
 _*Titan Text G1 - Express*_
 
 [https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=amazon.titan-text-express-v1](https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=amazon.titan-text-express-v1)
 
 ```ruby
-client.invoke_model(id: 'amazon.titan-text-express-v1', input: 'Generate a post about cats formatted with HTML tags')
+client.invoke_model(id: 'amazon.titan-text-express-v1', prompt: 'Generate a post about cats formatted with HTML tags', options: {})
 
 # Response
 {:inputTextTokenCount=>9,
@@ -132,6 +208,20 @@ Supports: Text generation, Code generation, Rich text formatting, Orchestration
 
 Anthropic offers the Claude family of large language models purpose built for conversations, summarization, Q&A, workflow automation, coding and more. Early customers report that Claude is much less likely to produce harmful outputs, easier to converse with, and more steerable - so you can get your desired output with less effort. Claude can also take direction on personality, tone, and behavior.
 
+Options for these models are:
+
+```ruby
+  {
+    temperature: 0.3,   # float
+    top_p: 0.5,         # float
+    top_k: 0.5,         # float
+    max_tokens: 1000,   # integer
+    stop_sequences: []  # [string]
+  }
+```
+
+For more documentation about the parameter data types and values you can [See more details...][https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html]
+
 How to call Anthropic models:
 
 _*Claude Instant 1.2*_
@@ -159,7 +249,7 @@ _*Claude 1.3*_
 Supports: Question answering, [information extraction](https://docs.anthropic.com/claude/docs/text-processing), removing PII, content generation, [multiple choice classification](https://docs.anthropic.com/claude/docs/multiple-choice-and-classification), [Roleplay](https://docs.anthropic.com/claude/docs/roleplay-dialogue), comparing text, summarization, [document Q&A with citation](https://docs.anthropic.com/claude/docs/advanced-text-analysis)
 
 ```ruby
-client.invoke_model(id: 'anthropic.claude-instant-v1', input: "You will be acting as a AI customer success agent for a company called Acme Dynamics.")
+client.invoke_model(id: 'anthropic.claude-v1', prompt: "You will be acting as a AI customer success agent for a company called Acme Dynamics.")
 
 # Response
 {:completion=>
@@ -177,7 +267,7 @@ _*Claude 2*_
 Supports: Question answering, [information extraction](https://docs.anthropic.com/claude/docs/text-processing), removing PII, content generation, [multiple choice classification](https://docs.anthropic.com/claude/docs/multiple-choice-and-classification), [Roleplay](https://docs.anthropic.com/claude/docs/roleplay-dialogue), comparing text, summarization, [document Q&A with citation](https://docs.anthropic.com/claude/docs/advanced-text-analysis)
 
 ```ruby
-client.invoke_model(id: 'anthropic.claude-v2', input: "I'm going to provide some text. I want to remove all person
+client.invoke_model(id: 'anthropic.claude-v2', prompt: "I'm going to provide some text. I want to remove all person
 ally identifying information from this text and replace it with XXX. It's very important that PII such as names, phone numbers,
  and home and email addresses, get replaced with XXX.")
 
@@ -192,6 +282,24 @@ ally identifying information from this text and replace it with XXX. It's very i
 
 Cohere models are text generation models for business use cases. Cohere models are trained on data that supports reliable business applications, like text generation, summarization, copywriting, dialogue, extraction, and question answering.
 
+Options for command models are:
+
+```ruby
+  {
+    temperature: 0.3,           # float
+    top_p: 0.5,                 # float
+    top_k: 0.5,                 # float
+    max_tokens: 1000,           # integer
+    stop_sequences: [],         # [string]
+    num_generations: 2,         # integer
+    return_likelihoods: 'ALL',  # string
+    stream: true,               # boolean
+    truncate: 'END'             # srtring
+  }
+```
+
+For more documentation about the parameter data types and values you can [See more details...][https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command.html]
+
 How to call Cohere command models:
 
 _*Command*_
@@ -201,7 +309,7 @@ _*Command*_
 Supports: Summarization, copywriting, dialogue, extraction, and question answering.
 
 ```ruby
-client.invoke_model(id: 'cohere.command-text-v14', input: 'Generate a twit about why Ruby on Rails is a great tool for building a startup. Write a few hashtags')
+client.invoke_model(id: 'cohere.command-text-v14', prompt: 'Generate a twit about why Ruby on Rails is a great tool for building a startup. Write a few hashtags')
 
 # Response
 {:generations=>
@@ -222,7 +330,7 @@ _*Command Light*_
 Supports: Summarization, copywriting, dialogue, extraction, and question answering.
 
 ```ruby
-client.invoke_model(id: 'cohere.command-light-text-v14', input: 'Generate a facebook post about GenAI models available at Amazon Bedrock')
+client.invoke_model(id: 'cohere.command-light-text-v14', prompt: 'Generate a facebook post about GenAI models available at Amazon Bedrock')
 
 # Response
 {:generations=>
@@ -234,6 +342,17 @@ client.invoke_model(id: 'cohere.command-light-text-v14', input: 'Generate a face
  :prompt=>"Generate a facebook post about GenAI models available at Amazon Bedrock:"}
 ```
 
+Options for embed models are:
+
+```ruby
+  {
+    input_type: 'classification', # string
+    truncate: 'END'               # string
+  }
+```
+
+For more documentation about the parameter data types and values you can [See more details...][https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-embed.html]
+
 How to call Cohere embed models:
 
 _*Embed English*_
@@ -244,7 +363,7 @@ Supports: Semantic search, retrieval-augmented generation (RAG), classification,
 
 ```ruby
 # WIP
-# client.invoke_model(id: 'cohere.embed-english-v3', input: 'Your prompt goes here ...')
+# client.invoke_model(id: 'cohere.embed-english-v3', prompt: 'Your prompt goes here ...')
 ```
 
 <br>
@@ -255,13 +374,23 @@ _*Embed Multilingual*_
 
 ```ruby
 # WIP
-# client.invoke_model(id: 'cohere.embed-multilingual-v3', input: 'Your prompt goes here ...')
+# client.invoke_model(id: 'cohere.embed-multilingual-v3', prompt: 'Your prompt goes here ...')
 ```
 
 ### Meta
 
 Meta is looking to unlock the power of large language models. Our latest version of Llama is now accessible to individuals, creators, researchers and businesses of all sizes so that they can experiment, innovate and scale their ideas responsibly.
 
+Options for these models are:
+
+```ruby
+  {
+    max_tokens: 128,  # integer
+    temperature: 0.9, # float
+    top_p: 0.7        # float
+  }
+```
+
 How to call Meta models:
 
 _*Llama 2 Chat 13B*_
@@ -271,7 +400,7 @@ _*Llama 2 Chat 13B*_
 Supports: Llama 2 is intended for commercial and research use in English. Fine-tuned chat models are intended for chat based applications.
 
 ```ruby
-client.invoke_model(id: 'meta.llama2-13b-chat-v1', input: 'Generate an Instagram Post about the Metaverse')
+client.invoke_model(id: 'meta.llama2-13b-chat-v1', prompt: 'Generate an Instagram Post about the Metaverse')
 
 # Resopnse
 {:generation=>
@@ -290,7 +419,7 @@ _*Llama 2 Chat 70B*_
 Supports: Llama 2 is intended for commercial and research use in English. Fine-tuned chat models are intended for chat based applications.
 
 ```ruby
-client.invoke_model(id: 'meta.llama2-70b-chat-v1', input: 'Generate a Facebook add to promote a new website that is selling Ruby on Rails and AI courses')
+client.invoke_model(id: 'meta.llama2-70b-chat-v1', prompt: 'Generate a Facebook add to promote a new website that is selling Ruby on Rails and AI courses')
 
 # Response
 {:generation=>
@@ -304,6 +433,16 @@ client.invoke_model(id: 'meta.llama2-70b-chat-v1', input: 'Generate a Facebook a
 
 Stability AI is the world's leading open-source generative artificial intelligence company, collaborating with public and private sector partners to bring next generation infrastructure to a global audience.
 
+Options for these models are:
+
+```ruby
+  {
+    cfg_scale: 20,  # integer
+    seed: 1,        # integer
+    steps: 10       # integer
+  }
+```
+
 How to call Stability AI models:
 
 _*SDXL 0.8*_
@@ -313,7 +452,7 @@ _*SDXL 0.8*_
 Supports: image generation, image editing
 
 ```ruby
-client.invoke_model(id: 'stability.stable-diffusion-xl-v0', input: 'Generate an image of an orca jumping out of the water', options: { file_path: 'path/to/your/image.jpg' })
+client.invoke_model(id: 'stability.stable-diffusion-xl-v0', prompt: 'Generate an image of an orca jumping out of the water', options: { file_path: 'path/to/your/image.jpg' })
 # NOTE: If file_path is not provided the image will be saved at 'image.jpg'
 
 # Success Response
@@ -338,7 +477,7 @@ Example generated image
 _*SDXL 1.0*_
 
 ```ruby
-client.invoke_model(id: 'stability.stable-diffusion-xl-v1', input: 'Generate an image of a white gold ring with a ruby on it', options: { file_path: 'path/to/your/image.jpg' })
+client.invoke_model(id: 'stability.stable-diffusion-xl-v1', prompt: 'Generate an image of a white gold ring with a ruby on it', options: { file_path: 'path/to/your/image.jpg' })
 # NOTE: If file_path is not provided the image will be saved at 'image.jpg'
 
 # Success Response

data/lib/amazon_bedrock/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RubyAmazonBedrock
-  VERSION = "0.2.0"
+  VERSION = "0.2.2"
 end

data/lib/amazon_bedrock.rb

@@ -3,8 +3,41 @@
 require_relative "amazon_bedrock/version"
 require_relative "bedrock_runtime/client"
 
+# The RubyAmazonBedrock module serves as a namespace for classes that
+# interact with Amazon Bedrock Service. It includes configuration management
+# and custom error classes for handling AWS service interactions.
 module RubyAmazonBedrock
   # Raised when an unknown model identifier is provided to PayloadFactory.
   class UnknownModelError < StandardError; end
   class Error < StandardError; end
+
+  # Configuration class for setting up AWS credentials and region.
+  class Configuration
+    attr_accessor :region, :access_key_id, :secret_access_key
+
+    # Initializes a new Configuration instance, loading values from
+    # environment variables or setting them to nil by default so the
+    # values can be dynamically set.
+    def initialize
+      @region = ENV.fetch('AWS_REGION', nil)
+      @access_key_id = ENV.fetch('AWS_ACCESS_KEY_ID', nil)
+      @secret_access_key = ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+    end
+  end
+
+  class << self
+    # Returns the current configuration instance or creates a new one.
+    #
+    # @return [Configuration] the current configuration instance
+    def configuration
+      @configuration ||= Configuration.new
+    end
+
+    # Provides a block to configure the RubyAmazonBedrock settings.
+    #
+    # @yield [Configuration] Yields the current configuration to a block.
+    def configure
+      yield(configuration)
+    end
+  end
 end

data/lib/bedrock_runtime/client.rb

@@ -4,7 +4,6 @@ require 'base64'
 require 'aws-sdk-bedrockruntime'
 require 'bedrock_runtime/payload_factory'
 require 'bedrock_runtime/response_factory'
-require 'pry'
 
 module RubyAmazonBedrock
   # Client for interfacing with the Amazon Bedrock Runtime.
@@ -15,27 +14,29 @@ module RubyAmazonBedrock
     # Initializes the AWS BedrockRuntime client.
     #
     # @note The AWS credentials and region are fetched from the environment variables.
-    def initialize(region:, access_key_id:, secret_access_key:)
+    def initialize(region: nil, access_key_id: nil, secret_access_key: nil)
+      config = RubyAmazonBedrock.configuration || RubyAmazonBedrock::Configuration.new
+
       @client = Aws::BedrockRuntime::Client.new(
-        region: region,
-        access_key_id: access_key_id,
-        secret_access_key: secret_access_key
+        region: region || config.region,
+        access_key_id: access_key_id || config.access_key_id,
+        secret_access_key: secret_access_key || config.secret_access_key
       )
     end
 
     # Invokes a model using the Bedrock Runtime client.
     #
     # @param id [String] The ID of the model to be invoked.
-    # @param input [String] The input string for what needs to be generated.
+    # @param prompt [String] The prompt string for what needs to be generated.
     # @param options [Hash] Additional options for the model invocation.
     # @return [Aws::BedrockRuntime::Types::InvokeModelOutput] The output from invoking the model.
     # @example Invoke a model with specific ID and input
     #   client = RubyAmazonBedrock::Client.new
     #   client.invoke_model(
-    # 		id: 'model_id', input: 'This is what you want to generate', options: { option_key: 'option_value' }
+    # 		id: 'model_id', prompt: 'This is what you want to generate', options: { option_key: 'option_value' }
     # 	)
-    def invoke_model(id:, input:, options: {})
-      payload_builder_class = RubyAmazonBedrock::PayloadFactory.new(id, input, options).create
+    def invoke_model(id:, prompt:, options: {})
+      payload_builder_class = RubyAmazonBedrock::PayloadFactory.new(id, prompt, options).create
       response = @client.invoke_model(payload_builder_class.build)
 
       response_builder_class = RubyAmazonBedrock::ResponseFactory.new(payload_builder_class.type, response,

data/lib/bedrock_runtime/payload_builders/ai_21_labs/base.rb

@@ -34,14 +34,14 @@ module RubyAmazonBedrock
             content_type: 'application/json',
             accept: '*/*',
             body: {
-              prompt: @input,
-              maxTokens: 200,
-              temperature: 0,
-              topP: 250,
-              stop_sequences: [],
-              countPenalty: { scale: 0 },
-              presencePenalty: { scale: 0 },
-              frequencyPenalty: { scale: 0 }
+              prompt: @prompt,
+              maxTokenCount: parameters[:maxTokenCount],
+              temperature: parameters[:temperature],
+              topP: parameters[:topP],
+              stopSequences: parameters[:stopSequences],
+              countPenalty: { scale: parameters[:countPenalty] },
+              presencePenalty: { scale: parameters[:presencePenalty] },
+              frequencyPenalty: { scale: parameters[:frequencyPenalty] }
             }.to_json
           }
         end
@@ -49,6 +49,18 @@ module RubyAmazonBedrock
         def model_id
           # noop
         end
+
+        def parameters # rubocop:disable Metrics/CyclomaticComplexity
+          {
+            maxTokenCount: @options[:max_tokens] || 200,
+            stopSequences: @options[:stop_sequences] || [],
+            temperature: @options[:temperature] || 0,
+            topP: @options[:top_p] || 1,
+            countPenalty: @options[:count_penalty] || 0,
+            presencePenalty: @options[:presence_penalty] || 0,
+            frequencyPenalty: @options[:frequency_penalty] || 0
+          }
+        end
       end
     end
   end

data/lib/bedrock_runtime/payload_builders/ai_21_labs/j2_mid_v1.rb

@@ -14,7 +14,7 @@ module RubyAmazonBedrock
         #
         # @return [String] 'ai21labs.j2-mid-v1'
         def model_id
-          'ai21labs.j2-mid-v1'
+          'ai21.j2-mid-v1'
         end
       end
     end

data/lib/bedrock_runtime/payload_builders/ai_21_labs/j2_ultra_v1.rb

@@ -14,7 +14,7 @@ module RubyAmazonBedrock
         #
         # @return [String] 'ai21labs.j2-ultra-v1'
         def model_id
-          'ai21labs.j2-ultra-v1'
+          'ai21.j2-ultra-v1'
         end
       end
     end

data/lib/bedrock_runtime/payload_builders/amazon/base.rb

@@ -31,13 +31,8 @@ module RubyAmazonBedrock
             content_type: 'application/json',
             accept: '*/*',
             body: {
-              inputText: @input,
-              textGenerationConfig: {
-                maxTokenCount: 4096,
-                stopSequences: [],
-                temperature: 0,
-                topP: 1
-              }
+              inputText: @prompt,
+              textGenerationConfig: parameters
             }.to_json
           }
         end
@@ -45,6 +40,15 @@ module RubyAmazonBedrock
         def model_id
           # noop
         end
+
+        def parameters
+          {
+            maxTokenCount: @options[:max_tokens] || 4096,
+            stopSequences: @options[:stop_sequences] || [],
+            temperature: @options[:temperature] || 0,
+            topP: @options[:top_p] || 1
+          }
+        end
       end
     end
   end

data/lib/bedrock_runtime/payload_builders/amazon/titan_image_generator_v1.rb

@@ -18,7 +18,7 @@ module RubyAmazonBedrock
             body: {
               taskType: "TEXT_IMAGE",
               textToImageParams: {
-                text: @input,
+                text: @prompt,
                 negativeTtext: @options[:negative_text]
               },
               imageGenerationConfig: {

data/lib/bedrock_runtime/payload_builders/amazon/titan_text_lite_v1.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative 'base'
+
 module RubyAmazonBedrock
   module PayloadBuilders
     module Amazon
@@ -9,21 +11,11 @@ module RubyAmazonBedrock
       # @see https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=amazon.titan-text-lite-v1
       # for more information about the Amazon model.
       class TitanTextLiteV1 < Base
-        def build
-          {
-            model_id: 'amazon.titan-text-lite-v1',
-            content_type: 'application/json',
-            accept: '*/*',
-            body: {
-              inputText: @input,
-              textGenerationConfig: {
-                maxTokenCount: 4096,
-                stopSequences: [],
-                temperature: 0,
-                topP: 1
-              }
-            }.to_json
-          }
+        # Returns the model ID for the Amazon Titan Text G1 - Lite model.
+        #
+        # @return [String] 'amazon.titan-text-lite-v1'
+        def model_id
+          'amazon.titan-text-lite-v1'
         end
       end
     end