lex-bedrock 0.1.3 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 68b06b9530a460c09dfba0bcdf216de10a7f95b17fb2faede82d9fcfe77abc0c
-  data.tar.gz: ae606b05bb7ba03720fb34cc61286372163221efa87d1771ccd87fcbf3a1252f
+  metadata.gz: ee5a83a0ee7afd53d1d8f496bd8382332b1312939b3a49cbdf8c86b2269ff6ab
+  data.tar.gz: 467176f83f6f58de4a7640bea5f2fcd6b75be4e2a8d5a77d4fa8f3fe82854c2b
 SHA512:
-  metadata.gz: bd75d1cc333a34ce86b9cd7ca854288c0ada18945e35b55050368a6ae74f7481072c4dc73a9b81549787caf9de898964200dec11dfdccf683abc504ba6bc3685
-  data.tar.gz: 15a7c72e46e0b97cdda8e2c5c1f31c4fd5b3d418bca5bab1a426944008f3021ea48031b98a1ecfdedeaea7be01b2d05395a194fbefbdc877a2a77e785d21e7c0
+  metadata.gz: 106a987eb8c0d61d3c8fc1e409b0168d95277db70f98fba5ee8ebd0dbe9b74c3770837b052b6cb04b32aae8bfd639511de8b907dfaccee09ee72c563792ae051
+  data.tar.gz: d8c81a605d36fd3e6b572580ad4a092a1a7bf913e5edf51e5e26302facf20cd2248da4311fcbf1867fe6cfc35eea78bea61c71c4c7ca0ff43c91d99b920482f6

data/CHANGELOG.md

@@ -1,5 +1,35 @@
 # Changelog
 
+## [0.2.1] - 2026-03-31
+
+### Added
+- `Helpers::Usage` — standardized usage hash normalization from SDK structs (`normalize`) and parsed JSON (`from_json`)
+- `Runners::Invoke#invoke_model` now returns a `usage:` key with standardized token counts extracted from the response body
+
+### Changed
+- `Runners::Converse#create` returns `usage:` as a plain hash (`{ input_tokens:, output_tokens:, cache_read_tokens:, cache_write_tokens: }`) instead of raw SDK struct
+- `Runners::Converse#create_stream` returns the same standardized usage hash
+- `Runners::Invoke#invoke_model` wrapped in `Helpers::Errors.with_retry` for production reliability
+
+## [0.2.0] - 2026-03-31
+
+### Added
+- `Helpers::Errors` — Bedrock exception classification (`ThrottlingException`, `AccessDeniedException`, `ModelNotReadyException`) with exponential backoff retry (`with_retry`)
+- `Helpers::Credentials` — AWS credential caching with TTL-based refresh (3000s default); falls back to SDK default provider chain when no explicit keys given
+- `Helpers::Thinking` — helpers to build `additional_model_request_fields` for extended thinking (`build_thinking_fields`, `sanitize_inference_config`)
+- `Helpers::ModelRegistry` — canonical-to-Bedrock ID mapping for 11 Anthropic models, `resolve`/`known?`/`all` API
+- `Runners::Converse#create_stream` — streaming Converse via `converse_stream` with per-event block callback and full accumulation fallback
+- `Runners::Converse#create_with_thinking` — convenience method that wires thinking betas and config into a single call
+- `Runners::Tokens#count_tokens` — token counting via `CountTokensCommand` with `anthropic_beta`/`thinking` body params
+- `Runners::Profiles` — inference profile listing, lookup, and canonical model ID resolution via `ListInferenceProfiles`/`GetInferenceProfile`
+
+### Changed
+- `Runners::Converse#create` now accepts `top_p`, `top_k`, `stop_sequences`, `tool_config`, `guardrail_config`, `additional_model_request_fields`
+- `Helpers::Client` — `bedrock_runtime_client`/`bedrock_client` now accept a pre-built `credentials:` kwarg; `access_key_id` is now optional (nil triggers SDK default provider chain)
+- `Helpers::Client` — added `region_for_model` for env-var and settings-driven per-model region routing
+- All runner calls wrapped in `Helpers::Errors.with_retry` for production reliability
+- `Legion::Extensions::Bedrock::Client` now includes `Runners::Tokens` and `Runners::Profiles`
+
 ## [0.1.3] - 2026-03-30
 
 ### Changed

data/lib/legion/extensions/bedrock/client.rb

@@ -4,6 +4,8 @@ require 'legion/extensions/bedrock/helpers/client'
 require 'legion/extensions/bedrock/runners/models'
 require 'legion/extensions/bedrock/runners/converse'
 require 'legion/extensions/bedrock/runners/invoke'
+require 'legion/extensions/bedrock/runners/tokens'
+require 'legion/extensions/bedrock/runners/profiles'
 
 module Legion
   module Extensions
@@ -12,6 +14,8 @@ module Legion
         include Legion::Extensions::Bedrock::Runners::Models
         include Legion::Extensions::Bedrock::Runners::Converse
         include Legion::Extensions::Bedrock::Runners::Invoke
+        include Legion::Extensions::Bedrock::Runners::Tokens
+        include Legion::Extensions::Bedrock::Runners::Profiles
 
         attr_reader :config
 

data/lib/legion/extensions/bedrock/helpers/client.rb

@@ -12,29 +12,53 @@ module Legion
 
           module_function
 
-          def bedrock_runtime_client(access_key_id:, secret_access_key:, region: DEFAULT_REGION,
-                                     session_token: nil, **)
-            opts = {
-              access_key_id:     access_key_id,
-              secret_access_key: secret_access_key,
-              region:            region
-            }
-            opts[:session_token] = session_token if session_token
-
-            Aws::BedrockRuntime::Client.new(**opts)
+          def bedrock_runtime_client(access_key_id: nil, secret_access_key: nil,
+                                     region: DEFAULT_REGION, session_token: nil,
+                                     credentials: nil, **)
+            Aws::BedrockRuntime::Client.new(
+              region:,
+              credentials: credentials || build_credentials(access_key_id:, secret_access_key:,
+                                                            session_token:)
+            )
           end
 
-          def bedrock_client(access_key_id:, secret_access_key:, region: DEFAULT_REGION,
-                             session_token: nil, **)
-            opts = {
-              access_key_id:     access_key_id,
-              secret_access_key: secret_access_key,
-              region:            region
-            }
-            opts[:session_token] = session_token if session_token
+          def bedrock_client(access_key_id: nil, secret_access_key: nil,
+                             region: DEFAULT_REGION, session_token: nil,
+                             credentials: nil, **)
+            Aws::Bedrock::Client.new(
+              region:,
+              credentials: credentials || build_credentials(access_key_id:, secret_access_key:,
+                                                            session_token:)
+            )
+          end
+
+          def region_for_model(model_id:, region: nil)
+            return region if region
+
+            env_key = "BEDROCK_REGION_#{model_id.upcase.gsub(/[^A-Z0-9]/, '_')}"
+            env_val = ENV.fetch(env_key, nil)
+            return env_val if env_val
+
+            settings_region = begin
+              Legion::Settings[:bedrock]&.dig(:model_regions, model_id.to_sym)
+            rescue StandardError => _e
+              nil
+            end
 
-            Aws::Bedrock::Client.new(**opts)
+            settings_region || DEFAULT_REGION
           end
+
+          def build_credentials(access_key_id:, secret_access_key:, session_token:)
+            return nil if access_key_id.nil?
+
+            if session_token
+              Aws::Credentials.new(access_key_id, secret_access_key, session_token)
+            else
+              Aws::Credentials.new(access_key_id, secret_access_key)
+            end
+          end
+
+          private_class_method :build_credentials
         end
       end
     end

data/lib/legion/extensions/bedrock/helpers/credentials.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Bedrock
+      module Helpers
+        module Credentials
+          CREDENTIAL_TTL = 3000 # seconds — refresh before the typical 1-hour STS TTL
+
+          @credential_cache = {} # rubocop:disable ThreadSafety/MutableClassInstanceVariable
+          @cache_mutex      = ::Mutex.new
+
+          module_function
+
+          def resolve(access_key_id: nil, secret_access_key: nil,
+                      session_token: nil, region: 'us-east-2', **)
+            return default_credentials if access_key_id.nil?
+
+            cache_key = "#{access_key_id}/#{region}"
+            @cache_mutex.synchronize do
+              entry = @credential_cache[cache_key]
+              if entry.nil? || stale?(entry)
+                @credential_cache[cache_key] = {
+                  credentials: build_static(
+                    access_key_id:,
+                    secret_access_key:,
+                    session_token:
+                  ),
+                  fetched_at:  ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+                }
+              end
+              @credential_cache[cache_key][:credentials]
+            end
+          end
+
+          def clear_cache!
+            @cache_mutex.synchronize { @credential_cache.clear }
+          end
+
+          def default_credentials
+            Aws::Credentials.new(nil, nil) # triggers SDK default chain
+          end
+
+          def build_static(access_key_id:, secret_access_key:, session_token:)
+            if session_token
+              Aws::Credentials.new(access_key_id, secret_access_key, session_token)
+            else
+              Aws::Credentials.new(access_key_id, secret_access_key)
+            end
+          end
+
+          def stale?(entry)
+            elapsed = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - entry[:fetched_at]
+            elapsed >= CREDENTIAL_TTL
+          end
+
+          private_class_method :default_credentials, :build_static, :stale?
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/bedrock/helpers/errors.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Bedrock
+      module Helpers
+        module Errors
+          THROTTLING_ERRORS = %w[
+            Aws::BedrockRuntime::Errors::ThrottlingException
+            Aws::BedrockRuntime::Errors::ServiceUnavailableException
+            Aws::Bedrock::Errors::ThrottlingException
+            Aws::Bedrock::Errors::ServiceUnavailableException
+          ].freeze
+
+          ACCESS_ERRORS = %w[
+            Aws::BedrockRuntime::Errors::AccessDeniedException
+            Aws::Bedrock::Errors::AccessDeniedException
+          ].freeze
+
+          MODEL_ERRORS = %w[
+            Aws::BedrockRuntime::Errors::ModelNotReadyException
+            Aws::BedrockRuntime::Errors::ModelTimeoutException
+            Aws::BedrockRuntime::Errors::ModelErrorException
+          ].freeze
+
+          MAX_RETRIES   = 3
+          BASE_DELAY    = 0.5  # seconds
+          MAX_DELAY     = 16.0 # seconds
+
+          module_function
+
+          def throttling_error?(exception)
+            THROTTLING_ERRORS.include?(exception.class.name)
+          end
+
+          def access_error?(exception)
+            ACCESS_ERRORS.include?(exception.class.name)
+          end
+
+          def model_error?(exception)
+            MODEL_ERRORS.include?(exception.class.name)
+          end
+
+          def retryable?(exception)
+            throttling_error?(exception)
+          end
+
+          def with_retry(max_retries: MAX_RETRIES)
+            attempts = 0
+            begin
+              yield
+            rescue StandardError => e
+              attempts += 1
+              raise unless retryable?(e) && attempts <= max_retries
+
+              delay = [BASE_DELAY * (2**(attempts - 1)), MAX_DELAY].min
+              sleep(delay)
+              retry
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/bedrock/helpers/model_registry.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Bedrock
+      module Helpers
+        module ModelRegistry
+          # Maps canonical model names to their AWS Bedrock cross-region inference profile IDs.
+          # Use the us.* cross-region IDs for best availability.
+          MODELS = {
+            'claude-3-5-haiku-20241022'  => 'us.anthropic.claude-3-5-haiku-20241022-v1:0',
+            'claude-haiku-4-5-20251001'  => 'us.anthropic.claude-haiku-4-5-20251001-v1:0',
+            'claude-3-5-sonnet-20241022' => 'anthropic.claude-3-5-sonnet-20241022-v2:0',
+            'claude-3-7-sonnet-20250219' => 'us.anthropic.claude-3-7-sonnet-20250219-v1:0',
+            'claude-sonnet-4-20250514'   => 'us.anthropic.claude-sonnet-4-20250514-v1:0',
+            'claude-sonnet-4-5-20250929' => 'us.anthropic.claude-sonnet-4-5-20250929-v1:0',
+            'claude-sonnet-4-6'          => 'us.anthropic.claude-sonnet-4-6',
+            'claude-opus-4-20250514'     => 'us.anthropic.claude-opus-4-20250514-v1:0',
+            'claude-opus-4-1-20250805'   => 'us.anthropic.claude-opus-4-1-20250805-v1:0',
+            'claude-opus-4-5-20251101'   => 'us.anthropic.claude-opus-4-5-20251101-v1:0',
+            'claude-opus-4-6'            => 'us.anthropic.claude-opus-4-6-v1'
+          }.freeze
+
+          module_function
+
+          def resolve(model_id)
+            return model_id if bedrock_id?(model_id)
+
+            MODELS.fetch(model_id, model_id)
+          end
+
+          def known?(model_id)
+            MODELS.key?(model_id)
+          end
+
+          def all
+            MODELS.dup
+          end
+
+          def bedrock_id?(model_id)
+            model_id.include?('.') || model_id.include?(':')
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/bedrock/helpers/thinking.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Bedrock
+      module Helpers
+        module Thinking
+          THINKING_BETA       = 'interleaved-thinking-2025-05-14'
+          CONTEXT_1M_BETA     = 'context-1m-2025-08-07'
+          TOOL_SEARCH_BETA    = 'tool-search-tool-2025-10-19'
+
+          module_function
+
+          def build_thinking_fields(budget_tokens: nil, adaptive: false, extra_betas: [])
+            betas   = [THINKING_BETA] + Array(extra_betas)
+            fields  = { anthropic_beta: betas }
+
+            thinking = if adaptive || budget_tokens.nil?
+                         { type: 'adaptive' }
+                       else
+                         { type: 'enabled', budget_tokens: }
+                       end
+
+            fields[:thinking] = thinking
+            fields
+          end
+
+          def sanitize_inference_config(inference_config:, thinking_enabled:)
+            return inference_config unless thinking_enabled
+
+            inference_config.except(:temperature)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/bedrock/helpers/usage.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Bedrock
+      module Helpers
+        module Usage
+          EMPTY_USAGE = {
+            input_tokens:       0,
+            output_tokens:      0,
+            cache_read_tokens:  0,
+            cache_write_tokens: 0
+          }.freeze
+
+          module_function
+
+          def normalize(sdk_usage)
+            return EMPTY_USAGE.dup unless sdk_usage
+
+            {
+              input_tokens:       sdk_usage.input_tokens || 0,
+              output_tokens:      sdk_usage.output_tokens || 0,
+              cache_read_tokens:  0,
+              cache_write_tokens: 0
+            }
+          end
+
+          def from_json(parsed)
+            return EMPTY_USAGE.dup unless parsed.is_a?(Hash)
+
+            usage = parsed['usage'] || {}
+            {
+              input_tokens:       usage['input_tokens'] || 0,
+              output_tokens:      usage['output_tokens'] || 0,
+              cache_read_tokens:  0,
+              cache_write_tokens: 0
+            }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/bedrock/runners/converse.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require 'legion/extensions/bedrock/helpers/client'
+require 'legion/extensions/bedrock/helpers/errors'
+require 'legion/extensions/bedrock/helpers/thinking'
+require 'legion/extensions/bedrock/helpers/usage'
 
 module Legion
   module Extensions
@@ -9,31 +12,136 @@ module Legion
         module Converse
           def create(model_id:, messages:, access_key_id:, secret_access_key:, # rubocop:disable Metrics/ParameterLists
                      system: nil, max_tokens: 1024, temperature: nil,
+                     top_p: nil, top_k: nil, stop_sequences: nil,
+                     tool_config: nil, guardrail_config: nil,
+                     additional_model_request_fields: nil,
                      region: Helpers::Client::DEFAULT_REGION, **)
             client = Helpers::Client.bedrock_runtime_client(
-              access_key_id:     access_key_id,
-              secret_access_key: secret_access_key,
-              region:            region
+              access_key_id:,
+              secret_access_key:,
+              region:
             )
 
-            inference_config = { max_tokens: max_tokens }
-            inference_config[:temperature] = temperature if temperature
-
-            request = {
-              model_id:         model_id,
-              messages:         messages,
-              inference_config: inference_config
-            }
-            request[:system] = [{ text: system }] if system
+            request = build_converse_request(
+              model_id:, messages:, system:, max_tokens:, temperature:,
+              top_p:, top_k:, stop_sequences:, tool_config:,
+              guardrail_config:, additional_model_request_fields:
+            )
 
-            response = client.converse(**request)
+            response = Helpers::Errors.with_retry { client.converse(**request) }
             {
               result:      response.output,
-              usage:       response.usage,
+              usage:       Helpers::Usage.normalize(response.usage),
               stop_reason: response.stop_reason
             }
           end
 
+          def create_stream(model_id:, messages:, access_key_id:, secret_access_key:, # rubocop:disable Metrics/ParameterLists
+                            system: nil, max_tokens: 1024, temperature: nil,
+                            top_p: nil, top_k: nil, stop_sequences: nil,
+                            tool_config: nil, guardrail_config: nil,
+                            additional_model_request_fields: nil,
+                            region: Helpers::Client::DEFAULT_REGION, **,
+                            &block)
+            client = Helpers::Client.bedrock_runtime_client(
+              access_key_id:,
+              secret_access_key:,
+              region:
+            )
+
+            request = build_converse_request(
+              model_id:, messages:, system:, max_tokens:, temperature:,
+              top_p:, top_k:, stop_sequences:, tool_config:,
+              guardrail_config:, additional_model_request_fields:
+            )
+
+            accumulated_text = +''
+            final_usage      = nil
+            final_stop       = nil
+
+            Helpers::Errors.with_retry do
+              client.converse_stream(**request) do |stream|
+                stream.on_content_block_delta_event do |event|
+                  delta = event.delta
+                  text  = delta.respond_to?(:text) ? delta.text : nil
+                  next if text.nil?
+
+                  accumulated_text << text
+                  block&.call(type: :delta, text:)
+                end
+
+                stream.on_message_stop_event do |event|
+                  final_stop = event.stop_reason
+                  block&.call(type: :stop, stop_reason: final_stop)
+                end
+
+                stream.on_metadata_event do |event|
+                  final_usage = event.usage
+                  block&.call(type: :usage, usage: final_usage)
+                end
+              end
+            end
+
+            {
+              result:      accumulated_text,
+              usage:       Helpers::Usage.normalize(final_usage),
+              stop_reason: final_stop
+            }
+          end
+
+          def create_with_thinking(model_id:, messages:, access_key_id:, secret_access_key:, # rubocop:disable Metrics/ParameterLists
+                                   budget_tokens: nil, adaptive: false, extra_betas: [],
+                                   system: nil, max_tokens: 16_000,
+                                   region: Helpers::Client::DEFAULT_REGION, **opts)
+            thinking_fields = Helpers::Thinking.build_thinking_fields(
+              budget_tokens:, adaptive:, extra_betas:
+            )
+
+            # Merge with any caller-supplied additional_model_request_fields
+            amrf = opts.delete(:additional_model_request_fields) || {}
+            merged_amrf = thinking_fields.merge(amrf) do |_key, thinking_val, caller_val|
+              thinking_val.is_a?(Array) ? thinking_val | Array(caller_val) : caller_val
+            end
+
+            create(
+              model_id:, messages:, access_key_id:, secret_access_key:,
+              system:, max_tokens:, region:,
+              additional_model_request_fields: merged_amrf,
+              **opts
+            )
+          end
+
+          private
+
+          def build_converse_request(model_id:, messages:, system:, max_tokens:, # rubocop:disable Metrics/ParameterLists
+                                     temperature:, top_p:, top_k:, stop_sequences:,
+                                     tool_config:, guardrail_config:,
+                                     additional_model_request_fields:)
+            inference_config = { max_tokens: }
+            inference_config[:temperature]    = temperature    if temperature
+            inference_config[:top_p]          = top_p          if top_p
+            inference_config[:stop_sequences] = stop_sequences if stop_sequences
+
+            request = {
+              model_id:,
+              messages:,
+              inference_config:
+            }
+            request[:system]                            = [{ text: system }] if system
+            request[:tool_config]                       = tool_config                  if tool_config
+            request[:guardrail_config]                  = guardrail_config             if guardrail_config
+            request[:additional_model_request_fields]   = additional_model_request_fields \
+              if additional_model_request_fields
+
+            # top_k goes in additional_model_request_fields for Anthropic models
+            if top_k
+              request[:additional_model_request_fields] ||= {}
+              request[:additional_model_request_fields][:top_k] = top_k
+            end
+
+            request
+          end
+
           include Legion::Extensions::Helpers::Lex if Legion::Extensions.const_defined?(:Helpers, false) &&
                                                       Legion::Extensions::Helpers.const_defined?(:Lex, false)
         end

data/lib/legion/extensions/bedrock/runners/invoke.rb

@@ -2,6 +2,8 @@
 
 require 'json'
 require 'legion/extensions/bedrock/helpers/client'
+require 'legion/extensions/bedrock/helpers/errors'
+require 'legion/extensions/bedrock/helpers/usage'
 
 module Legion
   module Extensions
@@ -12,21 +14,25 @@ module Legion
                            content_type: 'application/json', accept: 'application/json',
                            region: Helpers::Client::DEFAULT_REGION, **)
             client = Helpers::Client.bedrock_runtime_client(
-              access_key_id:     access_key_id,
-              secret_access_key: secret_access_key,
-              region:            region
+              access_key_id:,
+              secret_access_key:,
+              region:
             )
 
-            response = client.invoke_model(
-              model_id:     model_id,
-              body:         ::JSON.dump(body),
-              content_type: content_type,
-              accept:       accept
-            )
+            response = Helpers::Errors.with_retry do
+              client.invoke_model(
+                model_id:,
+                body:         ::JSON.dump(body),
+                content_type:,
+                accept:
+              )
+            end
 
+            parsed = ::JSON.parse(response.body.read)
             {
-              result:       ::JSON.parse(response.body.read),
-              content_type: response.content_type
+              result:       parsed,
+              content_type: response.content_type,
+              usage:        Helpers::Usage.from_json(parsed)
             }
           end
 

data/lib/legion/extensions/bedrock/runners/profiles.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require 'legion/extensions/bedrock/helpers/client'
+require 'legion/extensions/bedrock/helpers/errors'
+
+module Legion
+  module Extensions
+    module Bedrock
+      module Runners
+        module Profiles
+          def list_inference_profiles(access_key_id:, secret_access_key:,
+                                      profile_type: 'SYSTEM_DEFINED',
+                                      region: Helpers::Client::DEFAULT_REGION, **)
+            client = Helpers::Client.bedrock_client(
+              access_key_id:,
+              secret_access_key:,
+              region:
+            )
+
+            response = Helpers::Errors.with_retry do
+              client.list_inference_profiles(type_equals: profile_type)
+            end
+
+            { inference_profiles: response.inference_profile_summaries }
+          end
+
+          def get_inference_profile(profile_id:, access_key_id:, secret_access_key:,
+                                    region: Helpers::Client::DEFAULT_REGION, **)
+            client = Helpers::Client.bedrock_client(
+              access_key_id:,
+              secret_access_key:,
+              region:
+            )
+
+            response = Helpers::Errors.with_retry do
+              client.get_inference_profile(inference_profile_identifier: profile_id)
+            end
+
+            { inference_profile: response }
+          end
+
+          def resolve_profile_id(canonical_model_id:, access_key_id:, secret_access_key:,
+                                 region: Helpers::Client::DEFAULT_REGION, **)
+            result = list_inference_profiles(
+              access_key_id:, secret_access_key:, region:
+            )
+            profiles = result[:inference_profiles]
+
+            matched = profiles.find do |p|
+              p.respond_to?(:models) &&
+                Array(p.models).any? { |m| m.model_arn.to_s.include?(canonical_model_id) }
+            end
+
+            matched ? matched.inference_profile_id : canonical_model_id
+          end
+
+          include Legion::Extensions::Helpers::Lex if Legion::Extensions.const_defined?(:Helpers, false) &&
+                                                      Legion::Extensions::Helpers.const_defined?(:Lex, false)
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/bedrock/runners/tokens.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require 'legion/extensions/bedrock/helpers/client'
+require 'legion/extensions/bedrock/helpers/errors'
+
+module Legion
+  module Extensions
+    module Bedrock
+      module Runners
+        module Tokens
+          def count_tokens(model_id:, messages:, access_key_id:, secret_access_key:, # rubocop:disable Metrics/ParameterLists
+                           system: nil, tools: nil, anthropic_version: 'bedrock-2023-05-31',
+                           anthropic_beta: nil, thinking: nil,
+                           region: Helpers::Client::DEFAULT_REGION, **)
+            client = Helpers::Client.bedrock_runtime_client(
+              access_key_id:,
+              secret_access_key:,
+              region:
+            )
+
+            request = build_count_tokens_request(
+              model_id:, messages:, system:, tools:,
+              anthropic_version:, anthropic_beta:, thinking:
+            )
+
+            response = Helpers::Errors.with_retry { client.count_tokens(**request) }
+            {
+              input_token_count: response.input_tokens
+            }
+          end
+
+          private
+
+          def build_count_tokens_request(model_id:, messages:, system:, tools:,
+                                         anthropic_version:, anthropic_beta:, thinking:)
+            body_fields = { anthropic_version: }
+            body_fields[:anthropic_beta] = anthropic_beta if anthropic_beta
+            body_fields[:thinking]       = thinking if thinking
+
+            request = {
+              model_id:,
+              messages:,
+              system:                          system ? [{ text: system }] : nil,
+              additional_model_request_fields: body_fields
+            }
+            request[:tools] = tools if tools
+            request.compact
+          end
+
+          include Legion::Extensions::Helpers::Lex if Legion::Extensions.const_defined?(:Helpers, false) &&
+                                                      Legion::Extensions::Helpers.const_defined?(:Lex, false)
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/bedrock/version.rb

@@ -3,7 +3,7 @@
 module Legion
   module Extensions
     module Bedrock
-      VERSION = '0.1.3'
+      VERSION = '0.2.1'
     end
   end
 end

data/lib/legion/extensions/bedrock.rb

@@ -2,9 +2,16 @@
 
 require 'legion/extensions/bedrock/version'
 require 'legion/extensions/bedrock/helpers/client'
+require 'legion/extensions/bedrock/helpers/errors'
+require 'legion/extensions/bedrock/helpers/credentials'
+require 'legion/extensions/bedrock/helpers/thinking'
+require 'legion/extensions/bedrock/helpers/model_registry'
+require 'legion/extensions/bedrock/helpers/usage'
 require 'legion/extensions/bedrock/runners/models'
 require 'legion/extensions/bedrock/runners/converse'
 require 'legion/extensions/bedrock/runners/invoke'
+require 'legion/extensions/bedrock/runners/tokens'
+require 'legion/extensions/bedrock/runners/profiles'
 
 module Legion
   module Extensions

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: lex-bedrock
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.2.1
 platform: ruby
 authors:
 - Esity
@@ -157,9 +157,16 @@ files:
 - lib/legion/extensions/bedrock.rb
 - lib/legion/extensions/bedrock/client.rb
 - lib/legion/extensions/bedrock/helpers/client.rb
+- lib/legion/extensions/bedrock/helpers/credentials.rb
+- lib/legion/extensions/bedrock/helpers/errors.rb
+- lib/legion/extensions/bedrock/helpers/model_registry.rb
+- lib/legion/extensions/bedrock/helpers/thinking.rb
+- lib/legion/extensions/bedrock/helpers/usage.rb
 - lib/legion/extensions/bedrock/runners/converse.rb
 - lib/legion/extensions/bedrock/runners/invoke.rb
 - lib/legion/extensions/bedrock/runners/models.rb
+- lib/legion/extensions/bedrock/runners/profiles.rb
+- lib/legion/extensions/bedrock/runners/tokens.rb
 - lib/legion/extensions/bedrock/version.rb
 homepage: https://github.com/LegionIO/lex-bedrock
 licenses: