diff --git a/.rubocop.yml b/.rubocop.yml index e895ef9af..f0a752c94 100644 --- a/.rubocop.yml +++ b/.rubocop.yml @@ -35,5 +35,7 @@ Performance/UnfreezeString: - spec/**/* RSpec/ExampleLength: Enabled: false +RSpec/LeakyLocalVariable: + Enabled: false RSpec/MultipleExpectations: Enabled: false diff --git a/lib/ruby_llm/chat.rb b/lib/ruby_llm/chat.rb index 6b5b94daa..5e85a740c 100644 --- a/lib/ruby_llm/chat.rb +++ b/lib/ruby_llm/chat.rb @@ -97,6 +97,13 @@ def with_schema(schema) self end + def cache_prompts(**) + # No-op: Bedrock does not support prompt caching. This allows + # ruby_conversations (and other callers) to call cache_prompts + # without needing to know which provider is in use. + self + end + def on_new_message(&block) @on[:new_message] = block self diff --git a/lib/ruby_llm/configuration.rb b/lib/ruby_llm/configuration.rb index 34a842c2f..9dd0bc742 100644 --- a/lib/ruby_llm/configuration.rb +++ b/lib/ruby_llm/configuration.rb @@ -18,6 +18,8 @@ class Configuration :bedrock_secret_key, :bedrock_region, :bedrock_session_token, + :bedrock_guardrail_identifier, + :bedrock_guardrail_version, :openrouter_api_key, :ollama_api_base, :gpustack_api_base, diff --git a/lib/ruby_llm/providers/bedrock.rb b/lib/ruby_llm/providers/bedrock.rb index 50474b27f..57c9f5b6b 100644 --- a/lib/ruby_llm/providers/bedrock.rb +++ b/lib/ruby_llm/providers/bedrock.rb @@ -68,6 +68,19 @@ def build_headers(signature_headers, streaming: false) ) end + def build_guardrail_headers + headers = {} + guardrail_id = @config.bedrock_guardrail_identifier + guardrail_version = @config.bedrock_guardrail_version || 'DRAFT' + + if guardrail_id.present? + headers['X-Amzn-Bedrock-GuardrailIdentifier'] = guardrail_id + headers['X-Amzn-Bedrock-GuardrailVersion'] = guardrail_version + Rails.logger.info "[GUARDRAILS] Adding headers: #{headers.inspect}" if defined?(Rails) + end + headers + end + class << self def capabilities Bedrock::Capabilities diff --git a/lib/ruby_llm/providers/bedrock/chat.rb b/lib/ruby_llm/providers/bedrock/chat.rb index 4abdae3f3..695a3252f 100644 --- a/lib/ruby_llm/providers/bedrock/chat.rb +++ b/lib/ruby_llm/providers/bedrock/chat.rb @@ -9,8 +9,10 @@ module Chat def sync_response(connection, payload, additional_headers = {}) signature = sign_request("#{connection.connection.url_prefix}#{completion_url}", payload:) + guardrail_headers = build_guardrail_headers response = connection.post completion_url, payload do |req| req.headers.merge! build_headers(signature.headers, streaming: block_given?) + req.headers.merge!(guardrail_headers) # Add guardrails AFTER signing req.headers = additional_headers.merge(req.headers) unless additional_headers.empty? end Anthropic::Chat.parse_completion_response response diff --git a/lib/ruby_llm/providers/bedrock/streaming/base.rb b/lib/ruby_llm/providers/bedrock/streaming/base.rb index 751f02017..40b88e872 100644 --- a/lib/ruby_llm/providers/bedrock/streaming/base.rb +++ b/lib/ruby_llm/providers/bedrock/streaming/base.rb @@ -20,9 +20,11 @@ def stream_url def stream_response(connection, payload, additional_headers = {}, &block) signature = sign_request("#{connection.connection.url_prefix}#{stream_url}", payload:) accumulator = StreamAccumulator.new + guardrail_headers = build_guardrail_headers response = connection.post stream_url, payload do |req| req.headers.merge! build_headers(signature.headers, streaming: block_given?) + req.headers.merge!(guardrail_headers) # Add guardrails AFTER signing # Merge additional headers, with existing headers taking precedence req.headers = additional_headers.merge(req.headers) unless additional_headers.empty? req.options.on_data = handle_stream do |chunk|