Class: Legion::Extensions::Llm::Bedrock::Provider

Inherits:
Provider
  • Object
show all
Defined in:
lib/legion/extensions/llm/bedrock/provider.rb

Overview

Amazon Bedrock provider implementation for the Legion::Extensions::Llm contract.

Defined Under Namespace

Modules: Capabilities

Constant Summary collapse

DEFAULT_REGION =

rubocop:disable Metrics/ClassLength

'us-east-1'
STATIC_MODELS = 
[
  { model: 'anthropic.claude-3-haiku-20240307-v1:0', alias: 'claude-3-haiku' },
  { model: 'amazon.titan-text-express-v1', alias: 'titan-text-express' },
  { model: 'amazon.titan-embed-text-v2:0', alias: 'titan-embed-text-v2', usage_type: :embedding },
  { model: 'meta.llama3-2-11b-instruct-v1:0', alias: 'llama-3.2-11b-instruct' },
  { model: 'mistral.mistral-large-3-675b-instruct', alias: 'mistral-large-3' }
].freeze
ALIASES = 
STATIC_MODELS.to_h { |entry| [entry.fetch(:alias), entry.fetch(:model)] }.freeze

Class Method Summary collapse

Instance Method Summary collapse

Class Method Details

.capabilities ⇒ Object 



44
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 44

def capabilities = Capabilities

.configuration_options ⇒ Object 



31
32
33
34
35
36
37
38
39
40
41
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 31

def configuration_options
  %i[
    bedrock_region
    bedrock_endpoint
    bedrock_access_key_id
    bedrock_secret_access_key
    bedrock_session_token
    bedrock_profile
    bedrock_stub_responses
  ]
end

.configuration_requirements ⇒ Object 



43
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 43

def configuration_requirements = []

.resolve_model_id(model_id, config: nil) ⇒ Object

rubocop:disable Lint/UnusedMethodArgument



46
47
48
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 46

def resolve_model_id(model_id, config: nil) # rubocop:disable Lint/UnusedMethodArgument
  ALIASES.fetch(model_id.to_s, model_id.to_s)
end

.slug ⇒ Object 



29
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 29

def slug = 'bedrock'

Instance Method Details

#api_base ⇒ Object 



68
69
70
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 68

def api_base
  config.bedrock_endpoint || "https://bedrock-runtime.#{region}.amazonaws.com"
end

#chat(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {}) ⇒ Object 



135
136
137
138
139
140
141
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 135

def chat(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {})
  request = Utils.deep_merge(
    converse_request(messages, model:, temperature:, max_tokens:, tools:, tool_prefs:),
    params
  )
  parse_converse_response(runtime_client.converse(**request), model_id(model))
end

#complete(messages, tools:, temperature:, model:, params: {}, schema: nil, thinking: nil, tool_prefs: nil) ⇒ Object 



181
182
183
184
185
186
187
188
189
190
191
192
193
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 181

def complete(messages, tools:, temperature:, model:, params: {}, schema: nil, thinking: nil, tool_prefs: nil,
             &)
  payload = params.dup
  payload[:additional_model_request_fields] ||= {}
  payload[:additional_model_request_fields][:thinking] = thinking if thinking
  payload[:additional_model_request_fields][:response_format] = schema if schema

  if block_given?
    stream(messages, model:, temperature:, tools:, tool_prefs:, params: payload, &)
  else
    chat(messages, model:, temperature:, tools:, tool_prefs:, params: payload)
  end
end

#completion_url ⇒ Object 



72
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 72

def completion_url = 'Converse'

#count_tokens(messages, model:, system: nil, params: {}) ⇒ Object 



152
153
154
155
156
157
158
159
160
161
162
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 152

def count_tokens(messages, model:, system: nil, params: {})
  request = Utils.deep_merge(
    {
      model_id: model_id(model),
      input: { converse: { messages: format_messages(messages), system: system_blocks(system) }.compact }
    },
    params
  )
  response = runtime_client.count_tokens(**request)
  { input_tokens: value(response, :input_tokens), raw: normalize_response(response) }
end

#count_tokens_url ⇒ Object 



76
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 76

def count_tokens_url = 'CountTokens'

#discover_offerings(live: false, **filters) ⇒ Object 



82
83
84
85
86
87
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 82

def discover_offerings(live: false, **filters)
  return static_offerings(**filters) unless live

  response = bedrock_client.list_foundation_models(**filters)
  Array(value(response, :model_summaries)).map { |summary| offering_from_summary(summary) }
end

#embed(text, model:, dimensions: nil) ⇒ Object 



164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 164

def embed(text, model:, dimensions: nil)
  model_id = model_id(model)
  unless titan_embed?(model_id)
    raise NotImplementedError,
          "Bedrock embedding payload for #{model_id} is not standardized"
  end

  body = { inputText: text, dimensions: dimensions }.compact
  response = runtime_client.invoke_model(
    model_id: model_id,
    content_type: 'application/json',
    accept: 'application/json',
    body: Legion::JSON.generate(body)
  )
  parse_embedding_response(response, model: model_id)
end

#embedding_url ⇒ Object 



75
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 75

def embedding_url(**) = 'InvokeModel'

#health(live: false) ⇒ Object 



101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 101

def health(live: false)
  baseline = {
    provider: :bedrock,
    region: region,
    configured: true,
    ready: true,
    live: live,
    credentials: credential_source
  }
  return baseline.merge(checked: false) unless live

  bedrock_client.list_foundation_models
  baseline.merge(checked: true)
rescue StandardError => e
  baseline.merge(checked: true, ready: false, error: e.class.name, message: e.message)
end

#list_models ⇒ Object 



122
123
124
125
126
127
128
129
130
131
132
133
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 122

def list_models
  discover_offerings(live: true).map do |offering|
    Legion::Extensions::Llm::Model::Info.new(
      id: offering.model,
      name: offering.metadata[:alias] || offering.model,
      provider: :bedrock,
      family: offering.metadata[:model_family],
      capabilities: offering.capabilities.map(&:to_s),
      metadata: offering.to_h
    )
  end
end

#models_url ⇒ Object 



74
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 74

def models_url = 'ListFoundationModels'

#offering_for(model:, model_family: nil, instance_id: :default, **metadata) ⇒ Object 



89
90
91
92
93
94
95
96
97
98
99
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 89

def offering_for(model:, model_family: nil, instance_id: :default, **metadata)
  model_id = self.class.resolve_model_id(model)
  build_offering(
    model: model_id,
    alias_name: alias_for(model_id),
    model_family: model_family || model_family_for(model_id),
    instance_id: instance_id,
    usage_type: metadata.delete(:usage_type) || usage_type_for(model_id),
    metadata: metadata
  )
end

#readiness(live: false) ⇒ Object 



118
119
120
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 118

def readiness(live: false)
  health(live: live).merge(local: false, remote: true, api_base: api_base, endpoints: endpoint_manifest)
end

#region ⇒ Object 



78
79
80
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 78

def region
  config.bedrock_region || DEFAULT_REGION
end

#stream(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {}) ⇒ Object 



143
144
145
146
147
148
149
150
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 143

def stream(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {},
           &)
  request = Utils.deep_merge(
    converse_request(messages, model:, temperature:, max_tokens:, tools:, tool_prefs:),
    params
  )
  stream_converse(request, model_id(model), &)
end

#stream_url ⇒ Object 



73
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 73

def stream_url = 'ConverseStream'