lex-bedrock 0.1.3 → 0.2.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: 68b06b9530a460c09dfba0bcdf216de10a7f95b17fb2faede82d9fcfe77abc0c
4
- data.tar.gz: ae606b05bb7ba03720fb34cc61286372163221efa87d1771ccd87fcbf3a1252f
3
+ metadata.gz: fba4162f800e8419912e30cc4ce1477b5f9894f0f282bd67d3eafa3b61e74d9a
4
+ data.tar.gz: 2c26974438c1758108c6006c14eba339ede7feb48d6a65c157c1156de2c4dd68
5
5
  SHA512:
6
- metadata.gz: bd75d1cc333a34ce86b9cd7ca854288c0ada18945e35b55050368a6ae74f7481072c4dc73a9b81549787caf9de898964200dec11dfdccf683abc504ba6bc3685
7
- data.tar.gz: 15a7c72e46e0b97cdda8e2c5c1f31c4fd5b3d418bca5bab1a426944008f3021ea48031b98a1ecfdedeaea7be01b2d05395a194fbefbdc877a2a77e785d21e7c0
6
+ metadata.gz: 8e3d7990c6e40d37a7ee94a58f2a792ded501d015e318e6d953589a0fe7592e1648dd631e911dc5b37b3f84d1769f7895d0fa628fc652255acb102a8c091c70b
7
+ data.tar.gz: 4c643c663bbcac6110209a9dcb202e04962689fee2a2268321e91e6543d76476d7bc0d14eddc7efd71e5f76e597d66e8b37f3b97a940e4e41e897f7559bc6b07
data/CHANGELOG.md CHANGED
@@ -1,5 +1,24 @@
1
1
  # Changelog
2
2
 
3
+ ## [0.2.0] - 2026-03-31
4
+
5
+ ### Added
6
+ - `Helpers::Errors` — Bedrock exception classification (`ThrottlingException`, `AccessDeniedException`, `ModelNotReadyException`) with exponential backoff retry (`with_retry`)
7
+ - `Helpers::Credentials` — AWS credential caching with TTL-based refresh (3000s default); falls back to SDK default provider chain when no explicit keys given
8
+ - `Helpers::Thinking` — helpers to build `additional_model_request_fields` for extended thinking (`build_thinking_fields`, `sanitize_inference_config`)
9
+ - `Helpers::ModelRegistry` — canonical-to-Bedrock ID mapping for 11 Anthropic models, `resolve`/`known?`/`all` API
10
+ - `Runners::Converse#create_stream` — streaming Converse via `converse_stream` with per-event block callback and full accumulation fallback
11
+ - `Runners::Converse#create_with_thinking` — convenience method that wires thinking betas and config into a single call
12
+ - `Runners::Tokens#count_tokens` — token counting via `CountTokensCommand` with `anthropic_beta`/`thinking` body params
13
+ - `Runners::Profiles` — inference profile listing, lookup, and canonical model ID resolution via `ListInferenceProfiles`/`GetInferenceProfile`
14
+
15
+ ### Changed
16
+ - `Runners::Converse#create` now accepts `top_p`, `top_k`, `stop_sequences`, `tool_config`, `guardrail_config`, `additional_model_request_fields`
17
+ - `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)
18
+ - `Helpers::Client` — added `region_for_model` for env-var and settings-driven per-model region routing
19
+ - All runner calls wrapped in `Helpers::Errors.with_retry` for production reliability
20
+ - `Legion::Extensions::Bedrock::Client` now includes `Runners::Tokens` and `Runners::Profiles`
21
+
3
22
  ## [0.1.3] - 2026-03-30
4
23
 
5
24
  ### Changed
@@ -4,6 +4,8 @@ require 'legion/extensions/bedrock/helpers/client'
4
4
  require 'legion/extensions/bedrock/runners/models'
5
5
  require 'legion/extensions/bedrock/runners/converse'
6
6
  require 'legion/extensions/bedrock/runners/invoke'
7
+ require 'legion/extensions/bedrock/runners/tokens'
8
+ require 'legion/extensions/bedrock/runners/profiles'
7
9
 
8
10
  module Legion
9
11
  module Extensions
@@ -12,6 +14,8 @@ module Legion
12
14
  include Legion::Extensions::Bedrock::Runners::Models
13
15
  include Legion::Extensions::Bedrock::Runners::Converse
14
16
  include Legion::Extensions::Bedrock::Runners::Invoke
17
+ include Legion::Extensions::Bedrock::Runners::Tokens
18
+ include Legion::Extensions::Bedrock::Runners::Profiles
15
19
 
16
20
  attr_reader :config
17
21
 
@@ -12,29 +12,53 @@ module Legion
12
12
 
13
13
  module_function
14
14
 
15
- def bedrock_runtime_client(access_key_id:, secret_access_key:, region: DEFAULT_REGION,
16
- session_token: nil, **)
17
- opts = {
18
- access_key_id: access_key_id,
19
- secret_access_key: secret_access_key,
20
- region: region
21
- }
22
- opts[:session_token] = session_token if session_token
23
-
24
- Aws::BedrockRuntime::Client.new(**opts)
15
+ def bedrock_runtime_client(access_key_id: nil, secret_access_key: nil,
16
+ region: DEFAULT_REGION, session_token: nil,
17
+ credentials: nil, **)
18
+ Aws::BedrockRuntime::Client.new(
19
+ region:,
20
+ credentials: credentials || build_credentials(access_key_id:, secret_access_key:,
21
+ session_token:)
22
+ )
25
23
  end
26
24
 
27
- def bedrock_client(access_key_id:, secret_access_key:, region: DEFAULT_REGION,
28
- session_token: nil, **)
29
- opts = {
30
- access_key_id: access_key_id,
31
- secret_access_key: secret_access_key,
32
- region: region
33
- }
34
- opts[:session_token] = session_token if session_token
25
+ def bedrock_client(access_key_id: nil, secret_access_key: nil,
26
+ region: DEFAULT_REGION, session_token: nil,
27
+ credentials: nil, **)
28
+ Aws::Bedrock::Client.new(
29
+ region:,
30
+ credentials: credentials || build_credentials(access_key_id:, secret_access_key:,
31
+ session_token:)
32
+ )
33
+ end
34
+
35
+ def region_for_model(model_id:, region: nil)
36
+ return region if region
37
+
38
+ env_key = "BEDROCK_REGION_#{model_id.upcase.gsub(/[^A-Z0-9]/, '_')}"
39
+ env_val = ENV.fetch(env_key, nil)
40
+ return env_val if env_val
41
+
42
+ settings_region = begin
43
+ Legion::Settings[:bedrock]&.dig(:model_regions, model_id.to_sym)
44
+ rescue StandardError => _e
45
+ nil
46
+ end
35
47
 
36
- Aws::Bedrock::Client.new(**opts)
48
+ settings_region || DEFAULT_REGION
37
49
  end
50
+
51
+ def build_credentials(access_key_id:, secret_access_key:, session_token:)
52
+ return nil if access_key_id.nil?
53
+
54
+ if session_token
55
+ Aws::Credentials.new(access_key_id, secret_access_key, session_token)
56
+ else
57
+ Aws::Credentials.new(access_key_id, secret_access_key)
58
+ end
59
+ end
60
+
61
+ private_class_method :build_credentials
38
62
  end
39
63
  end
40
64
  end
@@ -0,0 +1,62 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Legion
4
+ module Extensions
5
+ module Bedrock
6
+ module Helpers
7
+ module Credentials
8
+ CREDENTIAL_TTL = 3000 # seconds — refresh before the typical 1-hour STS TTL
9
+
10
+ @credential_cache = {} # rubocop:disable ThreadSafety/MutableClassInstanceVariable
11
+ @cache_mutex = ::Mutex.new
12
+
13
+ module_function
14
+
15
+ def resolve(access_key_id: nil, secret_access_key: nil,
16
+ session_token: nil, region: 'us-east-2', **)
17
+ return default_credentials if access_key_id.nil?
18
+
19
+ cache_key = "#{access_key_id}/#{region}"
20
+ @cache_mutex.synchronize do
21
+ entry = @credential_cache[cache_key]
22
+ if entry.nil? || stale?(entry)
23
+ @credential_cache[cache_key] = {
24
+ credentials: build_static(
25
+ access_key_id:,
26
+ secret_access_key:,
27
+ session_token:
28
+ ),
29
+ fetched_at: ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
30
+ }
31
+ end
32
+ @credential_cache[cache_key][:credentials]
33
+ end
34
+ end
35
+
36
+ def clear_cache!
37
+ @cache_mutex.synchronize { @credential_cache.clear }
38
+ end
39
+
40
+ def default_credentials
41
+ Aws::Credentials.new(nil, nil) # triggers SDK default chain
42
+ end
43
+
44
+ def build_static(access_key_id:, secret_access_key:, session_token:)
45
+ if session_token
46
+ Aws::Credentials.new(access_key_id, secret_access_key, session_token)
47
+ else
48
+ Aws::Credentials.new(access_key_id, secret_access_key)
49
+ end
50
+ end
51
+
52
+ def stale?(entry)
53
+ elapsed = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - entry[:fetched_at]
54
+ elapsed >= CREDENTIAL_TTL
55
+ end
56
+
57
+ private_class_method :default_credentials, :build_static, :stale?
58
+ end
59
+ end
60
+ end
61
+ end
62
+ end
@@ -0,0 +1,65 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Legion
4
+ module Extensions
5
+ module Bedrock
6
+ module Helpers
7
+ module Errors
8
+ THROTTLING_ERRORS = %w[
9
+ Aws::BedrockRuntime::Errors::ThrottlingException
10
+ Aws::BedrockRuntime::Errors::ServiceUnavailableException
11
+ Aws::Bedrock::Errors::ThrottlingException
12
+ Aws::Bedrock::Errors::ServiceUnavailableException
13
+ ].freeze
14
+
15
+ ACCESS_ERRORS = %w[
16
+ Aws::BedrockRuntime::Errors::AccessDeniedException
17
+ Aws::Bedrock::Errors::AccessDeniedException
18
+ ].freeze
19
+
20
+ MODEL_ERRORS = %w[
21
+ Aws::BedrockRuntime::Errors::ModelNotReadyException
22
+ Aws::BedrockRuntime::Errors::ModelTimeoutException
23
+ Aws::BedrockRuntime::Errors::ModelErrorException
24
+ ].freeze
25
+
26
+ MAX_RETRIES = 3
27
+ BASE_DELAY = 0.5 # seconds
28
+ MAX_DELAY = 16.0 # seconds
29
+
30
+ module_function
31
+
32
+ def throttling_error?(exception)
33
+ THROTTLING_ERRORS.include?(exception.class.name)
34
+ end
35
+
36
+ def access_error?(exception)
37
+ ACCESS_ERRORS.include?(exception.class.name)
38
+ end
39
+
40
+ def model_error?(exception)
41
+ MODEL_ERRORS.include?(exception.class.name)
42
+ end
43
+
44
+ def retryable?(exception)
45
+ throttling_error?(exception)
46
+ end
47
+
48
+ def with_retry(max_retries: MAX_RETRIES)
49
+ attempts = 0
50
+ begin
51
+ yield
52
+ rescue StandardError => e
53
+ attempts += 1
54
+ raise unless retryable?(e) && attempts <= max_retries
55
+
56
+ delay = [BASE_DELAY * (2**(attempts - 1)), MAX_DELAY].min
57
+ sleep(delay)
58
+ retry
59
+ end
60
+ end
61
+ end
62
+ end
63
+ end
64
+ end
65
+ end
@@ -0,0 +1,47 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Legion
4
+ module Extensions
5
+ module Bedrock
6
+ module Helpers
7
+ module ModelRegistry
8
+ # Maps canonical model names to their AWS Bedrock cross-region inference profile IDs.
9
+ # Use the us.* cross-region IDs for best availability.
10
+ MODELS = {
11
+ 'claude-3-5-haiku-20241022' => 'us.anthropic.claude-3-5-haiku-20241022-v1:0',
12
+ 'claude-haiku-4-5-20251001' => 'us.anthropic.claude-haiku-4-5-20251001-v1:0',
13
+ 'claude-3-5-sonnet-20241022' => 'anthropic.claude-3-5-sonnet-20241022-v2:0',
14
+ 'claude-3-7-sonnet-20250219' => 'us.anthropic.claude-3-7-sonnet-20250219-v1:0',
15
+ 'claude-sonnet-4-20250514' => 'us.anthropic.claude-sonnet-4-20250514-v1:0',
16
+ 'claude-sonnet-4-5-20250929' => 'us.anthropic.claude-sonnet-4-5-20250929-v1:0',
17
+ 'claude-sonnet-4-6' => 'us.anthropic.claude-sonnet-4-6',
18
+ 'claude-opus-4-20250514' => 'us.anthropic.claude-opus-4-20250514-v1:0',
19
+ 'claude-opus-4-1-20250805' => 'us.anthropic.claude-opus-4-1-20250805-v1:0',
20
+ 'claude-opus-4-5-20251101' => 'us.anthropic.claude-opus-4-5-20251101-v1:0',
21
+ 'claude-opus-4-6' => 'us.anthropic.claude-opus-4-6-v1'
22
+ }.freeze
23
+
24
+ module_function
25
+
26
+ def resolve(model_id)
27
+ return model_id if bedrock_id?(model_id)
28
+
29
+ MODELS.fetch(model_id, model_id)
30
+ end
31
+
32
+ def known?(model_id)
33
+ MODELS.key?(model_id)
34
+ end
35
+
36
+ def all
37
+ MODELS.dup
38
+ end
39
+
40
+ def bedrock_id?(model_id)
41
+ model_id.include?('.') || model_id.include?(':')
42
+ end
43
+ end
44
+ end
45
+ end
46
+ end
47
+ end
@@ -0,0 +1,37 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Legion
4
+ module Extensions
5
+ module Bedrock
6
+ module Helpers
7
+ module Thinking
8
+ THINKING_BETA = 'interleaved-thinking-2025-05-14'
9
+ CONTEXT_1M_BETA = 'context-1m-2025-08-07'
10
+ TOOL_SEARCH_BETA = 'tool-search-tool-2025-10-19'
11
+
12
+ module_function
13
+
14
+ def build_thinking_fields(budget_tokens: nil, adaptive: false, extra_betas: [])
15
+ betas = [THINKING_BETA] + Array(extra_betas)
16
+ fields = { anthropic_beta: betas }
17
+
18
+ thinking = if adaptive || budget_tokens.nil?
19
+ { type: 'adaptive' }
20
+ else
21
+ { type: 'enabled', budget_tokens: }
22
+ end
23
+
24
+ fields[:thinking] = thinking
25
+ fields
26
+ end
27
+
28
+ def sanitize_inference_config(inference_config:, thinking_enabled:)
29
+ return inference_config unless thinking_enabled
30
+
31
+ inference_config.except(:temperature)
32
+ end
33
+ end
34
+ end
35
+ end
36
+ end
37
+ end
@@ -1,6 +1,8 @@
1
1
  # frozen_string_literal: true
2
2
 
3
3
  require 'legion/extensions/bedrock/helpers/client'
4
+ require 'legion/extensions/bedrock/helpers/errors'
5
+ require 'legion/extensions/bedrock/helpers/thinking'
4
6
 
5
7
  module Legion
6
8
  module Extensions
@@ -9,24 +11,23 @@ module Legion
9
11
  module Converse
10
12
  def create(model_id:, messages:, access_key_id:, secret_access_key:, # rubocop:disable Metrics/ParameterLists
11
13
  system: nil, max_tokens: 1024, temperature: nil,
14
+ top_p: nil, top_k: nil, stop_sequences: nil,
15
+ tool_config: nil, guardrail_config: nil,
16
+ additional_model_request_fields: nil,
12
17
  region: Helpers::Client::DEFAULT_REGION, **)
13
18
  client = Helpers::Client.bedrock_runtime_client(
14
- access_key_id: access_key_id,
15
- secret_access_key: secret_access_key,
16
- region: region
19
+ access_key_id:,
20
+ secret_access_key:,
21
+ region:
17
22
  )
18
23
 
19
- inference_config = { max_tokens: max_tokens }
20
- inference_config[:temperature] = temperature if temperature
21
-
22
- request = {
23
- model_id: model_id,
24
- messages: messages,
25
- inference_config: inference_config
26
- }
27
- request[:system] = [{ text: system }] if system
24
+ request = build_converse_request(
25
+ model_id:, messages:, system:, max_tokens:, temperature:,
26
+ top_p:, top_k:, stop_sequences:, tool_config:,
27
+ guardrail_config:, additional_model_request_fields:
28
+ )
28
29
 
29
- response = client.converse(**request)
30
+ response = Helpers::Errors.with_retry { client.converse(**request) }
30
31
  {
31
32
  result: response.output,
32
33
  usage: response.usage,
@@ -34,6 +35,112 @@ module Legion
34
35
  }
35
36
  end
36
37
 
38
+ def create_stream(model_id:, messages:, access_key_id:, secret_access_key:, # rubocop:disable Metrics/ParameterLists
39
+ system: nil, max_tokens: 1024, temperature: nil,
40
+ top_p: nil, top_k: nil, stop_sequences: nil,
41
+ tool_config: nil, guardrail_config: nil,
42
+ additional_model_request_fields: nil,
43
+ region: Helpers::Client::DEFAULT_REGION, **,
44
+ &block)
45
+ client = Helpers::Client.bedrock_runtime_client(
46
+ access_key_id:,
47
+ secret_access_key:,
48
+ region:
49
+ )
50
+
51
+ request = build_converse_request(
52
+ model_id:, messages:, system:, max_tokens:, temperature:,
53
+ top_p:, top_k:, stop_sequences:, tool_config:,
54
+ guardrail_config:, additional_model_request_fields:
55
+ )
56
+
57
+ accumulated_text = +''
58
+ final_usage = nil
59
+ final_stop = nil
60
+
61
+ Helpers::Errors.with_retry do
62
+ client.converse_stream(**request) do |stream|
63
+ stream.on_content_block_delta_event do |event|
64
+ delta = event.delta
65
+ text = delta.respond_to?(:text) ? delta.text : nil
66
+ next if text.nil?
67
+
68
+ accumulated_text << text
69
+ block&.call(type: :delta, text:)
70
+ end
71
+
72
+ stream.on_message_stop_event do |event|
73
+ final_stop = event.stop_reason
74
+ block&.call(type: :stop, stop_reason: final_stop)
75
+ end
76
+
77
+ stream.on_metadata_event do |event|
78
+ final_usage = event.usage
79
+ block&.call(type: :usage, usage: final_usage)
80
+ end
81
+ end
82
+ end
83
+
84
+ {
85
+ result: accumulated_text,
86
+ usage: final_usage,
87
+ stop_reason: final_stop
88
+ }
89
+ end
90
+
91
+ def create_with_thinking(model_id:, messages:, access_key_id:, secret_access_key:, # rubocop:disable Metrics/ParameterLists
92
+ budget_tokens: nil, adaptive: false, extra_betas: [],
93
+ system: nil, max_tokens: 16_000,
94
+ region: Helpers::Client::DEFAULT_REGION, **opts)
95
+ thinking_fields = Helpers::Thinking.build_thinking_fields(
96
+ budget_tokens:, adaptive:, extra_betas:
97
+ )
98
+
99
+ # Merge with any caller-supplied additional_model_request_fields
100
+ amrf = opts.delete(:additional_model_request_fields) || {}
101
+ merged_amrf = thinking_fields.merge(amrf) do |_key, thinking_val, caller_val|
102
+ thinking_val.is_a?(Array) ? thinking_val | Array(caller_val) : caller_val
103
+ end
104
+
105
+ create(
106
+ model_id:, messages:, access_key_id:, secret_access_key:,
107
+ system:, max_tokens:, region:,
108
+ additional_model_request_fields: merged_amrf,
109
+ **opts
110
+ )
111
+ end
112
+
113
+ private
114
+
115
+ def build_converse_request(model_id:, messages:, system:, max_tokens:, # rubocop:disable Metrics/ParameterLists
116
+ temperature:, top_p:, top_k:, stop_sequences:,
117
+ tool_config:, guardrail_config:,
118
+ additional_model_request_fields:)
119
+ inference_config = { max_tokens: }
120
+ inference_config[:temperature] = temperature if temperature
121
+ inference_config[:top_p] = top_p if top_p
122
+ inference_config[:stop_sequences] = stop_sequences if stop_sequences
123
+
124
+ request = {
125
+ model_id:,
126
+ messages:,
127
+ inference_config:
128
+ }
129
+ request[:system] = [{ text: system }] if system
130
+ request[:tool_config] = tool_config if tool_config
131
+ request[:guardrail_config] = guardrail_config if guardrail_config
132
+ request[:additional_model_request_fields] = additional_model_request_fields \
133
+ if additional_model_request_fields
134
+
135
+ # top_k goes in additional_model_request_fields for Anthropic models
136
+ if top_k
137
+ request[:additional_model_request_fields] ||= {}
138
+ request[:additional_model_request_fields][:top_k] = top_k
139
+ end
140
+
141
+ request
142
+ end
143
+
37
144
  include Legion::Extensions::Helpers::Lex if Legion::Extensions.const_defined?(:Helpers, false) &&
38
145
  Legion::Extensions::Helpers.const_defined?(:Lex, false)
39
146
  end
@@ -0,0 +1,63 @@
1
+ # frozen_string_literal: true
2
+
3
+ require 'legion/extensions/bedrock/helpers/client'
4
+ require 'legion/extensions/bedrock/helpers/errors'
5
+
6
+ module Legion
7
+ module Extensions
8
+ module Bedrock
9
+ module Runners
10
+ module Profiles
11
+ def list_inference_profiles(access_key_id:, secret_access_key:,
12
+ profile_type: 'SYSTEM_DEFINED',
13
+ region: Helpers::Client::DEFAULT_REGION, **)
14
+ client = Helpers::Client.bedrock_client(
15
+ access_key_id:,
16
+ secret_access_key:,
17
+ region:
18
+ )
19
+
20
+ response = Helpers::Errors.with_retry do
21
+ client.list_inference_profiles(type_equals: profile_type)
22
+ end
23
+
24
+ { inference_profiles: response.inference_profile_summaries }
25
+ end
26
+
27
+ def get_inference_profile(profile_id:, access_key_id:, secret_access_key:,
28
+ region: Helpers::Client::DEFAULT_REGION, **)
29
+ client = Helpers::Client.bedrock_client(
30
+ access_key_id:,
31
+ secret_access_key:,
32
+ region:
33
+ )
34
+
35
+ response = Helpers::Errors.with_retry do
36
+ client.get_inference_profile(inference_profile_identifier: profile_id)
37
+ end
38
+
39
+ { inference_profile: response }
40
+ end
41
+
42
+ def resolve_profile_id(canonical_model_id:, access_key_id:, secret_access_key:,
43
+ region: Helpers::Client::DEFAULT_REGION, **)
44
+ result = list_inference_profiles(
45
+ access_key_id:, secret_access_key:, region:
46
+ )
47
+ profiles = result[:inference_profiles]
48
+
49
+ matched = profiles.find do |p|
50
+ p.respond_to?(:models) &&
51
+ Array(p.models).any? { |m| m.model_arn.to_s.include?(canonical_model_id) }
52
+ end
53
+
54
+ matched ? matched.inference_profile_id : canonical_model_id
55
+ end
56
+
57
+ include Legion::Extensions::Helpers::Lex if Legion::Extensions.const_defined?(:Helpers, false) &&
58
+ Legion::Extensions::Helpers.const_defined?(:Lex, false)
59
+ end
60
+ end
61
+ end
62
+ end
63
+ end
@@ -0,0 +1,56 @@
1
+ # frozen_string_literal: true
2
+
3
+ require 'legion/extensions/bedrock/helpers/client'
4
+ require 'legion/extensions/bedrock/helpers/errors'
5
+
6
+ module Legion
7
+ module Extensions
8
+ module Bedrock
9
+ module Runners
10
+ module Tokens
11
+ def count_tokens(model_id:, messages:, access_key_id:, secret_access_key:, # rubocop:disable Metrics/ParameterLists
12
+ system: nil, tools: nil, anthropic_version: 'bedrock-2023-05-31',
13
+ anthropic_beta: nil, thinking: nil,
14
+ region: Helpers::Client::DEFAULT_REGION, **)
15
+ client = Helpers::Client.bedrock_runtime_client(
16
+ access_key_id:,
17
+ secret_access_key:,
18
+ region:
19
+ )
20
+
21
+ request = build_count_tokens_request(
22
+ model_id:, messages:, system:, tools:,
23
+ anthropic_version:, anthropic_beta:, thinking:
24
+ )
25
+
26
+ response = Helpers::Errors.with_retry { client.count_tokens(**request) }
27
+ {
28
+ input_token_count: response.input_tokens
29
+ }
30
+ end
31
+
32
+ private
33
+
34
+ def build_count_tokens_request(model_id:, messages:, system:, tools:,
35
+ anthropic_version:, anthropic_beta:, thinking:)
36
+ body_fields = { anthropic_version: }
37
+ body_fields[:anthropic_beta] = anthropic_beta if anthropic_beta
38
+ body_fields[:thinking] = thinking if thinking
39
+
40
+ request = {
41
+ model_id:,
42
+ messages:,
43
+ system: system ? [{ text: system }] : nil,
44
+ additional_model_request_fields: body_fields
45
+ }
46
+ request[:tools] = tools if tools
47
+ request.compact
48
+ end
49
+
50
+ include Legion::Extensions::Helpers::Lex if Legion::Extensions.const_defined?(:Helpers, false) &&
51
+ Legion::Extensions::Helpers.const_defined?(:Lex, false)
52
+ end
53
+ end
54
+ end
55
+ end
56
+ end
@@ -3,7 +3,7 @@
3
3
  module Legion
4
4
  module Extensions
5
5
  module Bedrock
6
- VERSION = '0.1.3'
6
+ VERSION = '0.2.0'
7
7
  end
8
8
  end
9
9
  end
@@ -2,9 +2,15 @@
2
2
 
3
3
  require 'legion/extensions/bedrock/version'
4
4
  require 'legion/extensions/bedrock/helpers/client'
5
+ require 'legion/extensions/bedrock/helpers/errors'
6
+ require 'legion/extensions/bedrock/helpers/credentials'
7
+ require 'legion/extensions/bedrock/helpers/thinking'
8
+ require 'legion/extensions/bedrock/helpers/model_registry'
5
9
  require 'legion/extensions/bedrock/runners/models'
6
10
  require 'legion/extensions/bedrock/runners/converse'
7
11
  require 'legion/extensions/bedrock/runners/invoke'
12
+ require 'legion/extensions/bedrock/runners/tokens'
13
+ require 'legion/extensions/bedrock/runners/profiles'
8
14
 
9
15
  module Legion
10
16
  module Extensions
metadata CHANGED
@@ -1,7 +1,7 @@
1
1
  --- !ruby/object:Gem::Specification
2
2
  name: lex-bedrock
3
3
  version: !ruby/object:Gem::Version
4
- version: 0.1.3
4
+ version: 0.2.0
5
5
  platform: ruby
6
6
  authors:
7
7
  - Esity
@@ -157,9 +157,15 @@ files:
157
157
  - lib/legion/extensions/bedrock.rb
158
158
  - lib/legion/extensions/bedrock/client.rb
159
159
  - lib/legion/extensions/bedrock/helpers/client.rb
160
+ - lib/legion/extensions/bedrock/helpers/credentials.rb
161
+ - lib/legion/extensions/bedrock/helpers/errors.rb
162
+ - lib/legion/extensions/bedrock/helpers/model_registry.rb
163
+ - lib/legion/extensions/bedrock/helpers/thinking.rb
160
164
  - lib/legion/extensions/bedrock/runners/converse.rb
161
165
  - lib/legion/extensions/bedrock/runners/invoke.rb
162
166
  - lib/legion/extensions/bedrock/runners/models.rb
167
+ - lib/legion/extensions/bedrock/runners/profiles.rb
168
+ - lib/legion/extensions/bedrock/runners/tokens.rb
163
169
  - lib/legion/extensions/bedrock/version.rb
164
170
  homepage: https://github.com/LegionIO/lex-bedrock
165
171
  licenses: