2023-11-23 10:58:54 -05:00
|
|
|
# frozen_string_literal: true
|
|
|
|
|
2023-12-18 14:06:06 -05:00
|
|
|
require "aws-sigv4"
|
|
|
|
|
2023-11-23 10:58:54 -05:00
|
|
|
module DiscourseAi
|
|
|
|
module Completions
|
|
|
|
module Endpoints
|
|
|
|
class AwsBedrock < Base
|
2024-01-29 14:04:25 -05:00
|
|
|
class << self
|
|
|
|
def can_contact?(endpoint_name, model_name)
|
2024-03-18 15:48:46 -04:00
|
|
|
endpoint_name == "aws_bedrock" &&
|
|
|
|
%w[claude-instant-1 claude-2 claude-3-haiku claude-3-sonnet].include?(model_name)
|
2024-01-29 14:04:25 -05:00
|
|
|
end
|
|
|
|
|
|
|
|
def dependant_setting_names
|
|
|
|
%w[ai_bedrock_access_key_id ai_bedrock_secret_access_key ai_bedrock_region]
|
|
|
|
end
|
|
|
|
|
2024-03-18 15:48:46 -04:00
|
|
|
def correctly_configured?(model)
|
2023-11-23 14:33:37 -05:00
|
|
|
SiteSetting.ai_bedrock_access_key_id.present? &&
|
2024-01-29 14:04:25 -05:00
|
|
|
SiteSetting.ai_bedrock_secret_access_key.present? &&
|
2024-03-18 15:48:46 -04:00
|
|
|
SiteSetting.ai_bedrock_region.present? && can_contact?("aws_bedrock", model)
|
2024-01-29 14:04:25 -05:00
|
|
|
end
|
|
|
|
|
|
|
|
def endpoint_name(model_name)
|
|
|
|
"AWS Bedrock - #{model_name}"
|
|
|
|
end
|
2023-11-23 10:58:54 -05:00
|
|
|
end
|
|
|
|
|
2024-01-04 07:53:47 -05:00
|
|
|
def normalize_model_params(model_params)
|
|
|
|
model_params = model_params.dup
|
|
|
|
|
2024-03-18 15:48:46 -04:00
|
|
|
# max_tokens, temperature, stop_sequences, top_p are already supported
|
2024-01-04 07:53:47 -05:00
|
|
|
|
|
|
|
model_params
|
|
|
|
end
|
|
|
|
|
2024-03-18 15:48:46 -04:00
|
|
|
def default_options(dialect)
|
|
|
|
options = { max_tokens: 3_000, anthropic_version: "bedrock-2023-05-31" }
|
|
|
|
options[:stop_sequences] = ["</function_calls>"] if dialect.prompt.has_tools?
|
|
|
|
options
|
2023-11-23 10:58:54 -05:00
|
|
|
end
|
|
|
|
|
|
|
|
def provider_id
|
|
|
|
AiApiAuditLog::Provider::Anthropic
|
|
|
|
end
|
|
|
|
|
|
|
|
private
|
|
|
|
|
2024-03-18 15:48:46 -04:00
|
|
|
def prompt_size(prompt)
|
|
|
|
# approximation
|
|
|
|
super(prompt.system_prompt.to_s + " " + prompt.messages.to_s)
|
|
|
|
end
|
|
|
|
|
2023-11-23 10:58:54 -05:00
|
|
|
def model_uri
|
2024-01-09 12:10:20 -05:00
|
|
|
# See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html
|
2024-03-18 15:48:46 -04:00
|
|
|
#
|
|
|
|
# FYI there is a 2.0 version of Claude, very little need to support it given
|
|
|
|
# haiku/sonnet are better fits anyway, we map to claude-2.1
|
|
|
|
bedrock_model_id =
|
|
|
|
case model
|
|
|
|
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"
|
|
|
|
end
|
2023-11-23 13:49:24 -05:00
|
|
|
|
2023-11-23 10:58:54 -05:00
|
|
|
api_url =
|
2024-03-18 15:48:46 -04:00
|
|
|
"https://bedrock-runtime.#{SiteSetting.ai_bedrock_region}.amazonaws.com/model/#{bedrock_model_id}/invoke"
|
2023-11-23 10:58:54 -05:00
|
|
|
|
|
|
|
api_url = @streaming_mode ? (api_url + "-with-response-stream") : api_url
|
|
|
|
|
|
|
|
URI(api_url)
|
|
|
|
end
|
|
|
|
|
2024-03-18 15:48:46 -04:00
|
|
|
def prepare_payload(prompt, model_params, dialect)
|
|
|
|
payload = default_options(dialect).merge(model_params).merge(messages: prompt.messages)
|
|
|
|
payload[:system] = prompt.system_prompt if prompt.system_prompt.present?
|
|
|
|
payload
|
2023-11-23 10:58:54 -05:00
|
|
|
end
|
|
|
|
|
|
|
|
def prepare_request(payload)
|
|
|
|
headers = { "content-type" => "application/json", "Accept" => "*/*" }
|
|
|
|
|
|
|
|
signer =
|
|
|
|
Aws::Sigv4::Signer.new(
|
|
|
|
access_key_id: SiteSetting.ai_bedrock_access_key_id,
|
|
|
|
region: SiteSetting.ai_bedrock_region,
|
|
|
|
secret_access_key: SiteSetting.ai_bedrock_secret_access_key,
|
|
|
|
service: "bedrock",
|
|
|
|
)
|
|
|
|
|
|
|
|
Net::HTTP::Post
|
2023-11-23 13:49:24 -05:00
|
|
|
.new(model_uri)
|
2023-11-23 10:58:54 -05:00
|
|
|
.tap do |r|
|
|
|
|
r.body = payload
|
|
|
|
|
|
|
|
signed_request =
|
|
|
|
signer.sign_request(req: r, http_method: r.method, url: model_uri, body: r.body)
|
|
|
|
|
|
|
|
r.initialize_http_header(headers.merge(signed_request.headers))
|
|
|
|
end
|
|
|
|
end
|
|
|
|
|
|
|
|
def decode(chunk)
|
2024-02-01 06:01:07 -05:00
|
|
|
parsed =
|
|
|
|
Aws::EventStream::Decoder
|
|
|
|
.new
|
|
|
|
.decode_chunk(chunk)
|
|
|
|
.first
|
|
|
|
.payload
|
|
|
|
.string
|
|
|
|
.then { JSON.parse(_1) }
|
|
|
|
|
|
|
|
bytes = parsed.dig("bytes")
|
|
|
|
|
|
|
|
if !bytes
|
|
|
|
Rails.logger.error("#{self.class.name}: #{parsed.to_s[0..500]}")
|
|
|
|
nil
|
|
|
|
else
|
|
|
|
Base64.decode64(parsed.dig("bytes"))
|
|
|
|
end
|
2023-11-23 10:58:54 -05:00
|
|
|
rescue JSON::ParserError,
|
|
|
|
Aws::EventStream::Errors::MessageChecksumError,
|
|
|
|
Aws::EventStream::Errors::PreludeChecksumError => e
|
|
|
|
Rails.logger.error("#{self.class.name}: #{e.message}")
|
|
|
|
nil
|
|
|
|
end
|
|
|
|
|
2024-03-18 15:48:46 -04:00
|
|
|
def final_log_update(log)
|
|
|
|
log.request_tokens = @input_tokens if @input_tokens
|
|
|
|
log.response_tokens = @output_tokens if @output_tokens
|
|
|
|
end
|
|
|
|
|
2023-11-23 10:58:54 -05:00
|
|
|
def extract_completion_from(response_raw)
|
2024-03-18 15:48:46 -04:00
|
|
|
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
|
2023-11-23 10:58:54 -05:00
|
|
|
end
|
|
|
|
|
|
|
|
def partials_from(decoded_chunk)
|
|
|
|
[decoded_chunk]
|
|
|
|
end
|
|
|
|
end
|
|
|
|
end
|
|
|
|
end
|
|
|
|
end
|