ruby_llm-bedrock_invoke 0.1.0.pre.1

data/lib/ruby_llm/bedrock_invoke/protocol_v2.rb

@@ -0,0 +1,162 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module BedrockInvoke
+    # The Anthropic Messages wire format over Bedrock InvokeModel (RubyLLM 2.0).
+    #
+    # Mirrors the in-tree VertexAI::Anthropic pattern: inherit the entire
+    # Anthropic protocol and override only the endpoints, the Bedrock body
+    # envelope (no top-level model/stream, anthropic_version marker), SigV4
+    # signing, and AWS event-stream decoding for streaming.
+    class Protocol < RubyLLM::Protocols::Anthropic
+      include EventStream
+
+      ANTHROPIC_VERSION = 'bedrock-2023-05-31'
+
+      def completion_url
+        "/model/#{RubyLLM::BedrockInvoke.escape_model_id(@model.id)}/invoke"
+      end
+
+      def stream_url
+        "/model/#{RubyLLM::BedrockInvoke.escape_model_id(@model.id)}/invoke-with-response-stream"
+      end
+
+      def render_payload(messages, **options)
+        payload = super
+        payload.delete(:model)
+        payload.delete(:stream)
+        payload[:anthropic_version] = ANTHROPIC_VERSION
+        ToolSearch.apply!(payload, options[:tools] || {})
+        payload
+      end
+
+      def parse_completion_body(data, raw:)
+        message = super
+        preserve_server_blocks(message, data)
+        message
+      end
+
+      # Replays messages that carried server-managed tool-search blocks
+      # verbatim; everything else renders through the stock Anthropic format.
+      def format_message(msg, **)
+        blocks = RubyLLM::BedrockInvoke.raw_blocks(msg)
+        return super unless blocks
+
+        { role: 'assistant', content: blocks }
+      end
+
+      # 2.0 gates thinking on registry reasoning_options, which assumed models
+      # never carry; Bedrock Claude supports budget thinking, so emit it
+      # directly. Adaptive/effort-only thinking is not available on Bedrock.
+      def build_thinking_payload(thinking, _model)
+        inner = RubyLLM::BedrockInvoke.budget_thinking_payload(thinking)
+        return nil if inner.nil?
+
+        { thinking: inner }
+      end
+
+      # Bedrock has no Anthropic Files API.
+      def supports_provider_file_references?
+        false
+      end
+
+      private
+
+      def sync_response(payload, additional_headers = {})
+        ToolSearch.ensure_beta!(payload)
+        body = JSON.generate(payload)
+        response = @connection.post(completion_url, body) do |req|
+          req.headers['Content-Type'] = 'application/json'
+          req.headers.merge!(additional_headers) unless additional_headers.empty?
+          req.headers.merge!(@provider.sign_headers('POST', completion_url, body))
+        end
+        parse_completion_response(response)
+      end
+
+      def stream_response(payload, additional_headers = {}, &block)
+        ToolSearch.ensure_beta!(payload)
+        state = StreamState.new(self)
+        failure_buffer = +''
+        body = JSON.generate(payload)
+
+        response = @connection.post(stream_url, body) do |req|
+          req.headers['Content-Type'] = 'application/json'
+          req.headers['Accept'] = 'application/vnd.amazon.eventstream'
+          req.headers.merge!(additional_headers) unless additional_headers.empty?
+          req.headers.merge!(@provider.sign_headers('POST', stream_url, body))
+          install_stream_handler(req) do |bytes, env, received|
+            consume_stream_bytes(bytes, env, received, state, failure_buffer, &block)
+          end
+        end
+
+        message = state.accumulator.to_message(response)
+        attach_collected_blocks(message, state.collector)
+        RubyLLM.logger.debug { "Stream completed: #{message.content}" }
+        message
+      end
+
+      def install_stream_handler(req, &handler)
+        if Faraday::VERSION.start_with?('1')
+          req.options[:on_data] = proc { |chunk, size| handler.call(chunk, nil, size) }
+        else
+          req.options.on_data = proc { |chunk, received, env| handler.call(chunk, env, received) }
+        end
+      end
+
+      def consume_stream_bytes(bytes, env, received, state, failure_buffer, &block)
+        return handle_failed_stream_chunk(bytes, env, failure_buffer) if failed_stream_chunk?(bytes, env)
+
+        state.reset_if_new_attempt!(bytes, received)
+        each_event_stream_event(state.decoder, bytes) do |event|
+          state.collector.observe(event)
+          next if state.collector.suppress?(event)
+
+          chunk = build_chunk(event)
+          next unless chunk
+
+          state.accumulator.add(chunk)
+          block.call(chunk)
+        end
+      end
+
+      # Faraday 1 gives on_data no response env, so non-200 JSON error bodies
+      # would otherwise be fed to the event-stream decoder (binary frames
+      # never start with '{' — the first bytes are a frame length prelude).
+      def failed_stream_chunk?(bytes, env)
+        return env.status != 200 if env
+
+        bytes.to_s.lstrip.start_with?('{')
+      end
+
+      # Error bodies can arrive split across reads; accumulate until they parse.
+      def handle_failed_stream_chunk(bytes, env, failure_buffer)
+        failure_buffer << bytes
+        data = JSON.parse(failure_buffer)
+        response = env ? env.merge(body: data) : Struct.new(:body, :status).new(data, 500)
+        ErrorMiddleware.parse_error(provider: @provider, response: response)
+      rescue JSON::ParserError
+        RubyLLM.logger.debug { "bedrock_invoke: accumulating failed stream chunk: #{bytes.inspect}" }
+      end
+
+      def error_middleware_provider
+        @provider
+      end
+
+      # Responses carrying server-managed tool-search blocks must be replayed
+      # verbatim. 2.0 Message content is String-only, so the block array rides
+      # in an ivar that format_message re-injects when rendering history.
+      def preserve_server_blocks(message, data)
+        blocks = data.is_a?(Hash) ? data['content'] : nil
+        return unless blocks && ToolSearch.contains_server_blocks?(blocks)
+
+        message.instance_variable_set(RubyLLM::BedrockInvoke::RAW_BLOCKS_IVAR, blocks)
+      end
+
+      def attach_collected_blocks(message, collector)
+        return unless collector.server_blocks?
+
+        message.instance_variable_set(RubyLLM::BedrockInvoke::RAW_BLOCKS_IVAR, collector.content_blocks)
+      end
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/provider_v1.rb

@@ -0,0 +1,215 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    # Anthropic Claude on AWS Bedrock via the InvokeModel API (RubyLLM 1.x).
+    #
+    # Inherits the entire Anthropic Messages wire format — message/tool/schema
+    # rendering, structured output via output_config, cache_control, response
+    # and chunk parsing — and swaps the transport underneath: per-request
+    # SigV4 signing against bedrock-runtime and AWS event-stream decoding for
+    # InvokeModelWithResponseStream.
+    class BedrockInvoke < Anthropic
+      include RubyLLM::BedrockInvoke::Signing
+      include RubyLLM::BedrockInvoke::EventStream
+
+      ANTHROPIC_VERSION = 'bedrock-2023-05-31'
+
+      def api_base
+        config_value(:bedrock_invoke_api_base) ||
+          "https://bedrock-runtime.#{bedrock_invoke_region}.amazonaws.com"
+      end
+
+      def headers
+        {}
+      end
+
+      def completion_url
+        "/model/#{RubyLLM::BedrockInvoke.escape_model_id(@model_id)}/invoke"
+      end
+
+      def stream_url
+        "/model/#{RubyLLM::BedrockInvoke.escape_model_id(@model_id)}/invoke-with-response-stream"
+      end
+
+      # rubocop:disable Metrics/ParameterLists
+      def render_payload(messages, tools:, temperature:, model:, stream: false,
+                         schema: nil, thinking: nil, tool_prefs: nil)
+        @model_id = model.id
+        payload = super
+        payload.delete(:model)
+        payload.delete(:stream)
+        payload[:anthropic_version] = ANTHROPIC_VERSION
+        RubyLLM::BedrockInvoke::ToolSearch.apply!(payload, tools)
+        payload
+      end
+      # rubocop:enable Metrics/ParameterLists
+
+      def parse_completion_response(response)
+        message = super
+        preserve_server_blocks(message, response)
+        message
+      end
+
+      def parse_error(response)
+        return if response.body.nil? || response.body.empty?
+
+        body = try_parse_json(response.body)
+        return body if body.is_a?(String)
+
+        body['message'] || body['Message'] || body['error'] || body['__type'] || super
+      end
+
+      def list_models
+        []
+      end
+
+      # Messages preserving server tool-search blocks replay verbatim.
+      # Bypassing the stock formatter also avoids its thinking-block prepend,
+      # which would both duplicate thinking blocks (they are already in the
+      # raw array) and permanently mutate the preserved array via unshift.
+      def format_message(msg, **)
+        blocks = RubyLLM::BedrockInvoke.raw_blocks(msg)
+        return { role: 'assistant', content: blocks } if blocks
+
+        super
+      end
+
+      # 1.16 gates thinking on registry reasoning_options, which assumed
+      # models never carry; Bedrock Claude supports budget thinking, so emit
+      # it directly. Called as (thinking) on 1.13-1.15 (bare payload spliced
+      # into payload[:thinking]) and as (thinking, model) on 1.16 (wrapped
+      # shape unwrapped by add_thinking_fields).
+      def build_thinking_payload(thinking, model = nil)
+        inner = RubyLLM::BedrockInvoke.budget_thinking_payload(thinking)
+        return nil if inner.nil?
+
+        model.nil? ? inner : { thinking: inner }
+      end
+
+      class << self
+        def slug
+          'bedrock_invoke'
+        end
+
+        def capabilities
+          RubyLLM::BedrockInvoke::Capabilities
+        end
+
+        def configuration_requirements
+          []
+        end
+
+        def configuration_options
+          RubyLLM::BedrockInvoke::CONFIGURATION_OPTIONS
+        end
+
+        def configured?(config)
+          region = (config.bedrock_invoke_region if config.respond_to?(:bedrock_invoke_region))
+          !(region || ENV['AWS_REGION'] || ENV.fetch('AWS_DEFAULT_REGION', nil)).nil?
+        end
+
+        def assume_models_exist?
+          true
+        end
+      end
+
+      private
+
+      def sync_response(connection, payload, additional_headers = {})
+        RubyLLM::BedrockInvoke::ToolSearch.ensure_beta!(payload)
+        body = JSON.generate(payload)
+        response = connection.post(completion_url, body) do |req|
+          req.headers['Content-Type'] = 'application/json'
+          req.headers.merge!(additional_headers) unless additional_headers.empty?
+          req.headers.merge!(sign_headers('POST', completion_url, body))
+        end
+        parse_completion_response(response)
+      end
+
+      def stream_response(connection, payload, additional_headers = {}, &block)
+        RubyLLM::BedrockInvoke::ToolSearch.ensure_beta!(payload)
+        state = RubyLLM::BedrockInvoke::StreamState.new(self)
+        failure_buffer = +''
+        body = JSON.generate(payload)
+
+        response = connection.post(stream_url, body) do |req|
+          req.headers['Content-Type'] = 'application/json'
+          req.headers['Accept'] = 'application/vnd.amazon.eventstream'
+          req.headers.merge!(additional_headers) unless additional_headers.empty?
+          req.headers.merge!(sign_headers('POST', stream_url, body))
+          install_stream_handler(req) do |bytes, env, received|
+            consume_stream_bytes(bytes, env, received, state, failure_buffer, &block)
+          end
+        end
+
+        message = state.accumulator.to_message(response)
+        attach_collected_blocks(message, state.collector)
+        RubyLLM.logger.debug { "Stream completed: #{message.content}" }
+        message
+      end
+
+      def install_stream_handler(req, &handler)
+        if Faraday::VERSION.start_with?('1')
+          req.options[:on_data] = proc { |chunk, size| handler.call(chunk, nil, size) }
+        else
+          req.options.on_data = proc { |chunk, received, env| handler.call(chunk, env, received) }
+        end
+      end
+
+      def consume_stream_bytes(bytes, env, received, state, failure_buffer, &block)
+        return handle_failed_stream_chunk(bytes, env, failure_buffer) if failed_stream_chunk?(bytes, env)
+
+        state.reset_if_new_attempt!(bytes, received)
+        each_event_stream_event(state.decoder, bytes) do |event|
+          state.collector.observe(event)
+          next if state.collector.suppress?(event)
+
+          chunk = build_chunk(event)
+          next unless chunk
+
+          state.accumulator.add(chunk)
+          block.call(chunk)
+        end
+      end
+
+      # Faraday 1 gives on_data no response env, so non-200 JSON error bodies
+      # would otherwise be fed to the event-stream decoder (binary frames
+      # never start with '{' — the first bytes are a frame length prelude).
+      def failed_stream_chunk?(bytes, env)
+        return env.status != 200 if env
+
+        bytes.to_s.lstrip.start_with?('{')
+      end
+
+      # Error bodies can arrive split across reads; accumulate until they parse.
+      def handle_failed_stream_chunk(bytes, env, failure_buffer)
+        failure_buffer << bytes
+        data = JSON.parse(failure_buffer)
+        response = env ? env.merge(body: data) : Struct.new(:body, :status).new(data, 500)
+        ErrorMiddleware.parse_error(provider: self, response: response)
+      rescue JSON::ParserError
+        RubyLLM.logger.debug { "bedrock_invoke: accumulating failed stream chunk: #{bytes.inspect}" }
+      end
+
+      # When the response carries server-managed tool-search blocks, the whole
+      # content-block array must go back to the API verbatim on the next
+      # request. Message exposes no writer, so the content ivar is replaced
+      # in place — the safe cross-version seam given 1.13-1.16 constructor drift.
+      def preserve_server_blocks(message, response)
+        blocks = response.body.is_a?(Hash) ? response.body['content'] : nil
+        return unless blocks && RubyLLM::BedrockInvoke::ToolSearch.contains_server_blocks?(blocks)
+
+        raw = RubyLLM::BedrockInvoke::RawContent.new(blocks, text: message.content)
+        message.instance_variable_set(:@content, raw)
+      end
+
+      def attach_collected_blocks(message, collector)
+        return unless collector.server_blocks?
+
+        raw = RubyLLM::BedrockInvoke::RawContent.new(collector.content_blocks, text: message.content)
+        message.instance_variable_set(:@content, raw)
+      end
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/provider_v2.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    # Anthropic Claude on AWS Bedrock via the InvokeModel API (RubyLLM 2.0).
+    #
+    # Where/auth only — the wire format lives in
+    # RubyLLM::BedrockInvoke::Protocol. Credentials come from the standard
+    # AWS chain (env vars, shared config, IRSA/Pod Identity, IMDS) unless a
+    # credential provider is configured explicitly.
+    class BedrockInvoke < Provider
+      include RubyLLM::BedrockInvoke::Signing
+
+      protocol :anthropic, RubyLLM::BedrockInvoke::Protocol
+
+      def api_base
+        config_value(:bedrock_invoke_api_base) ||
+          "https://bedrock-runtime.#{bedrock_invoke_region}.amazonaws.com"
+      end
+
+      def headers
+        {}
+      end
+
+      def ensure_configured!
+        return if self.class.configured?(@config)
+
+        raise ConfigurationError,
+              'No AWS region configured for bedrock_invoke. Set config.bedrock_invoke_region ' \
+              'or the AWS_REGION environment variable.'
+      end
+
+      def parse_error(response)
+        body = response.body
+        body = try_parse_json(body) if respond_to?(:try_parse_json, true)
+        return if body.nil? || (body.respond_to?(:empty?) && body.empty?)
+
+        return body if body.is_a?(String)
+
+        body['message'] || body['Message'] || body['error'] || body['__type'] || super
+      end
+
+      def list_models
+        []
+      end
+
+      class << self
+        def capabilities
+          RubyLLM::BedrockInvoke::Capabilities
+        end
+
+        def configuration_options
+          RubyLLM::BedrockInvoke::CONFIGURATION_OPTIONS
+        end
+
+        def configuration_requirements
+          []
+        end
+
+        def configured?(config)
+          region = (config.bedrock_invoke_region if config.respond_to?(:bedrock_invoke_region))
+          !(region || ENV['AWS_REGION'] || ENV.fetch('AWS_DEFAULT_REGION', nil)).nil?
+        end
+
+        def assume_models_exist?
+          true
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/raw_content.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module BedrockInvoke
+    # Raw Anthropic content blocks that must be replayed verbatim on
+    # subsequent requests (server_tool_use / tool_search_tool_result and their
+    # siblings), while still displaying as the assistant's text.
+    #
+    # RubyLLM renders RubyLLM::Content::Raw message content into the payload
+    # unchanged, which is exactly what the tool search tool requires: its
+    # server-side blocks are passed back untouched and expanded by the API.
+    class RawContent < RubyLLM::Content::Raw
+      def initialize(value, text: nil)
+        super(value)
+        @text = text
+      end
+
+      def to_s
+        @text.to_s
+      end
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/signing.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+require 'aws-sdk-core'
+require 'uri'
+
+module RubyLLM
+  module BedrockInvoke
+    # SigV4 request signing backed by the standard AWS credential chain.
+    #
+    # Credentials resolve exactly like any AWS SDK client: environment
+    # variables, shared config/credentials files, SSO, ECS/EKS container
+    # credentials (Pod Identity), IRSA web-identity tokens, and IMDS. A
+    # token projected into a pod therefore "just works" with no RubyLLM
+    # configuration beyond the region.
+    #
+    # Mixed into objects that expose +@config+ (the RubyLLM configuration)
+    # and +api_base+.
+    module Signing
+      SERVICE = 'bedrock'
+
+      # The default chain is resolved once per process per region and shared:
+      # the returned credentials object self-refreshes (the AWS SDK norm),
+      # and re-resolving per Chat would mean an STS/IMDS network round trip
+      # for every RubyLLM.chat under IRSA or instance profiles.
+      DEFAULT_CHAIN_MUTEX = Mutex.new
+
+      class << self
+        def cached_default_chain_credentials(region)
+          DEFAULT_CHAIN_MUTEX.synchronize do
+            @default_chain_cache ||= {}
+            @default_chain_cache[region] ||= ::Aws::STS::Client.new(region: region).config.credentials
+          end
+        end
+
+        # For tests and credential rotation.
+        def reset_default_chain_cache!
+          DEFAULT_CHAIN_MUTEX.synchronize { @default_chain_cache = {} }
+        end
+      end
+
+      def bedrock_invoke_region
+        region = presence(config_value(:bedrock_invoke_region)) ||
+                 presence(ENV.fetch('AWS_REGION', nil)) ||
+                 presence(ENV.fetch('AWS_DEFAULT_REGION', nil))
+        return region if region
+
+        raise RubyLLM::ConfigurationError,
+              'No AWS region configured for bedrock_invoke. Set config.bedrock_invoke_region ' \
+              'or the AWS_REGION environment variable.'
+      end
+
+      def sigv4_signer
+        @sigv4_signer ||= ::Aws::Sigv4::Signer.new(
+          service: SERVICE,
+          region: bedrock_invoke_region,
+          credentials_provider: bedrock_invoke_credentials
+        )
+      end
+
+      # Returns headers (authorization, x-amz-date, x-amz-content-sha256 and,
+      # for temporary credentials, x-amz-security-token) for one request.
+      # +path+ is relative to +api_base+; +body+ must be the exact bytes sent.
+      # URI.join mirrors Faraday's absolute-path semantics, so a custom
+      # api_base carrying a path component cannot desynchronize the signature
+      # from the transmitted URL.
+      def sign_headers(method, path, body)
+        signature = sigv4_signer.sign_request(
+          http_method: method,
+          url: URI.join(api_base, path).to_s,
+          body: body
+        )
+        signature.headers
+      end
+
+      def bedrock_invoke_credentials
+        explicit = config_value(:bedrock_invoke_credential_provider)
+        return explicit if explicit
+
+        credentials = Signing.cached_default_chain_credentials(bedrock_invoke_region)
+        return credentials if credentials
+
+        raise RubyLLM::ConfigurationError,
+              'No AWS credentials found for bedrock_invoke. Provide credentials via the standard ' \
+              'AWS chain (env vars, shared config, IRSA/Pod Identity, instance profile) or set ' \
+              'config.bedrock_invoke_credential_provider.'
+      end
+
+      private
+
+      def config_value(key)
+        return nil unless @config.respond_to?(key)
+
+        @config.public_send(key)
+      end
+
+      # On <= 1.15 the config accessors are plain attr_accessors, so empty
+      # strings are not coerced to nil the way 1.16+/2.0 writers do.
+      def presence(value)
+        return nil if value.nil?
+        return nil if value.respond_to?(:empty?) && value.to_s.strip.empty?
+
+        value
+      end
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/stream_state.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module BedrockInvoke
+    # Per-attempt streaming state: accumulator, block collector, and
+    # event-stream decoder.
+    #
+    # RubyLLM's Faraday connection installs retry middleware that re-issues
+    # POSTs on retryable failures (throttling raised mid-stream, connection
+    # drops). The on_data callback closes over this holder; resetting it at
+    # the first chunk of each physical attempt prevents a retried stream from
+    # appending to a half-filled accumulator or feeding a decoder stuck on a
+    # partial frame. Detection: Faraday reports cumulative received bytes,
+    # which equal the chunk size exactly when an attempt starts fresh.
+    class StreamState
+      attr_reader :accumulator, :collector, :decoder
+
+      def initialize(event_stream_owner)
+        @owner = event_stream_owner
+        reset!
+      end
+
+      def reset!
+        @accumulator = RubyLLM::StreamAccumulator.new
+        @collector = BlockCollector.new
+        @decoder = @owner.event_stream_decoder
+      end
+
+      def reset_if_new_attempt!(chunk, received_bytes)
+        reset! if received_bytes && received_bytes == chunk.bytesize
+      end
+    end
+  end
+end