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

Inherits:
Provider
  • Object
show all
Includes:
Logging::Helper
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 = 
'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 Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Class Attribute Details

.registry_publisher ⇒ Object 



49
50
51
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 49

def registry_publisher
  @registry_publisher ||= RegistryPublisher.new
end

Class Method Details

.capabilities ⇒ Object 



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

def capabilities = Capabilities

.configuration_options ⇒ Object 



34
35
36
37
38
39
40
41
42
43
44
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 34

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 



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

def configuration_requirements = []

.resolve_model_id(model_id, config: nil) ⇒ Object

rubocop:disable Lint/UnusedMethodArgument



53
54
55
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 53

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 



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

def slug = 'bedrock'

Instance Method Details

#api_base ⇒ Object 



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

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 



160
161
162
163
164
165
166
167
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 160

def chat(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {})
  log.info { "bedrock.provider.chat: model=#{model_id(model)} messages=#{messages.size}" }
  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: {}, headers: {}, schema: nil, thinking: nil, tool_prefs: nil) ⇒ Object

rubocop:disable Lint/UnusedMethodArgument



210
211
212
213
214
215
216
217
218
219
220
221
222
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 210

def complete(messages, tools:, temperature:, model:, params: {}, headers: {}, schema: nil, thinking: nil, # rubocop:disable Lint/UnusedMethodArgument
             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 



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

def completion_url = 'Converse'

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



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

def count_tokens(messages, model:, system: nil, params: {})
  log.debug { "bedrock.provider.count_tokens: model=#{model_id(model)}" }
  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 



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

def count_tokens_url = 'CountTokens'

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



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

def discover_offerings(live: false, **filters)
  unless live
    log.debug { 'bedrock.provider.discover_offerings: returning static catalog' }
    return static_offerings(**filters)
  end

  log.info { "bedrock.provider.discover_offerings: listing foundation models (region=#{region})" }
  response = bedrock_client.list_foundation_models(**filters)
  Array(value(response, :model_summaries)).map { |summary| offering_from_summary(summary) }.tap do |offerings|
    log.info { "bedrock.provider.discover_offerings: found #{offerings.size} models" }
    self.class.registry_publisher.publish_offerings_async(offerings, readiness: readiness(live: false))
  end
end

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



192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 192

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

  log.info { "bedrock.provider.embed: model=#{mid}" }
  body = { inputText: text, dimensions: dimensions }.compact
  response = runtime_client.invoke_model(
    model_id: mid,
    content_type: 'application/json',
    accept: 'application/json',
    body: Legion::JSON.generate(body)
  )
  parse_embedding_response(response, model: mid)
end

#embedding_url ⇒ Object 



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

def embedding_url(**) = 'InvokeModel'

#health(live: false) ⇒ Object 



115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 115

def health(live: false)
  baseline = {
    provider: :bedrock,
    region: region,
    configured: true,
    ready: true,
    live: live,
    credentials: credential_source
  }
  unless live
    log.debug { "bedrock.provider.health: offline check (region=#{region})" }
    return baseline.merge(checked: false)
  end

  log.info { "bedrock.provider.health: live check (region=#{region})" }
  bedrock_client.list_foundation_models
  log.info { 'bedrock.provider.health: live check passed' }
  baseline.merge(checked: true)
rescue StandardError => e
  handle_exception(e, level: :warn, handled: true, operation: 'bedrock.provider.health')
  baseline.merge(checked: true, ready: false, error: e.class.name, message: e.message)
end

#list_models ⇒ Object 



146
147
148
149
150
151
152
153
154
155
156
157
158
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 146

def list_models
  log.info { 'bedrock.provider.list_models: fetching live model list' }
  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 



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

def models_url = 'ListFoundationModels'

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



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

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 



138
139
140
141
142
143
144
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 138

def readiness(live: false)
  log.debug { "bedrock.provider.readiness: checking (live=#{live})" }
  health(live: live).merge(local: false, remote: true, api_base: api_base,
                           endpoints: endpoint_manifest).tap do |metadata|
    self.class.registry_publisher.publish_readiness_async(metadata) if live
  end
end

#region ⇒ Object 



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

def region
  config.bedrock_region || DEFAULT_REGION
end

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



169
170
171
172
173
174
175
176
177
 
# File 'lib/legion/extensions/llm/bedrock/provider.rb', line 169

def stream(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {},
           &)
  log.info { "bedrock.provider.stream: model=#{model_id(model)} messages=#{messages.size}" }
  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 



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

def stream_url = 'ConverseStream'