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 +4 -4
- data/CHANGELOG.md +19 -0
- data/lib/legion/extensions/bedrock/client.rb +4 -0
- data/lib/legion/extensions/bedrock/helpers/client.rb +43 -19
- data/lib/legion/extensions/bedrock/helpers/credentials.rb +62 -0
- data/lib/legion/extensions/bedrock/helpers/errors.rb +65 -0
- data/lib/legion/extensions/bedrock/helpers/model_registry.rb +47 -0
- data/lib/legion/extensions/bedrock/helpers/thinking.rb +37 -0
- data/lib/legion/extensions/bedrock/runners/converse.rb +120 -13
- data/lib/legion/extensions/bedrock/runners/profiles.rb +63 -0
- data/lib/legion/extensions/bedrock/runners/tokens.rb +56 -0
- data/lib/legion/extensions/bedrock/version.rb +1 -1
- data/lib/legion/extensions/bedrock.rb +6 -0
- metadata +7 -1
checksums.yaml
CHANGED
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
---
|
|
2
2
|
SHA256:
|
|
3
|
-
metadata.gz:
|
|
4
|
-
data.tar.gz:
|
|
3
|
+
metadata.gz: fba4162f800e8419912e30cc4ce1477b5f9894f0f282bd67d3eafa3b61e74d9a
|
|
4
|
+
data.tar.gz: 2c26974438c1758108c6006c14eba339ede7feb48d6a65c157c1156de2c4dd68
|
|
5
5
|
SHA512:
|
|
6
|
-
metadata.gz:
|
|
7
|
-
data.tar.gz:
|
|
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
|
|
16
|
-
session_token: nil,
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
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
|
|
28
|
-
session_token: nil,
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
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
|
-
|
|
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
|
|
15
|
-
secret_access_key
|
|
16
|
-
region:
|
|
19
|
+
access_key_id:,
|
|
20
|
+
secret_access_key:,
|
|
21
|
+
region:
|
|
17
22
|
)
|
|
18
23
|
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
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
|
|
@@ -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.
|
|
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:
|