ruby_llm-bedrock_invoke 0.1.0.pre.1

data/lib/ruby_llm/bedrock_invoke/tool_search.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module BedrockInvoke
+    # Payload shaping for Anthropic's tool search tool on Bedrock InvokeModel.
+    #
+    # Bedrock exposes tool search only through InvokeModel (not Converse), as
+    # a beta: the request body needs `anthropic_beta: ["tool-search-tool-2025-10-19"]`
+    # and the *unversioned* `tool_search_tool_regex` tool type (the first-party
+    # API uses versioned types and no beta flag; no BM25 variant exists on
+    # Bedrock). Tools marked deferred still send their full definition every
+    # request with `defer_loading: true` — they are excluded from the cached
+    # prompt prefix until the model discovers them via the search tool.
+    module ToolSearch
+      BETA_FLAG = 'tool-search-tool-2025-10-19'
+      SEARCH_TOOL_TYPE = 'tool_search_tool_regex'
+      SEARCH_TOOL_NAME = 'tool_search_tool_regex'
+      SERVER_BLOCK_TYPES = %w[server_tool_use tool_search_tool_result].freeze
+
+      module_function
+
+      def deferred?(tool)
+        tool.respond_to?(:defer_loading?) && tool.defer_loading?
+      end
+
+      # +tools+ is RubyLLM's name => tool hash; returns name strings.
+      def deferred_names(tools)
+        tools.values.select { |tool| deferred?(tool) }.map { |tool| tool.name.to_s }
+      end
+
+      # Mutates a rendered Anthropic Messages payload: flags deferred tool
+      # definitions, prepends the search tool, and opts into the beta.
+      # Tools may be marked via the Deferred module or by carrying
+      # defer_loading in their own provider params. No-op when nothing is
+      # deferred.
+      def apply!(payload, tools)
+        names = deferred_names(tools)
+        definitions = payload[:tools]
+
+        if definitions.nil?
+          return payload if names.empty?
+
+          raise RubyLLM::BedrockInvoke::Error, 'Deferred tools present but payload has no tools'
+        end
+
+        definitions.each do |definition|
+          name = definition_value(definition, :name).to_s
+          definition[:defer_loading] = true if names.include?(name)
+        end
+
+        return payload unless definitions.any? { |d| deferred_definition?(d) }
+
+        inject_search_tool!(definitions)
+        payload[:anthropic_beta] = Array(payload[:anthropic_beta]) | [BETA_FLAG]
+        payload
+      end
+
+      # Both generations deep-merge user params (with_params /
+      # with_provider_options) over the rendered payload AFTER apply! runs,
+      # and deep_merge REPLACES arrays — silently dropping the beta flag.
+      # Called on the final merged payload right before the request goes out.
+      def ensure_beta!(payload)
+        definitions = payload[:tools] || payload['tools']
+        return payload unless definitions.is_a?(Array)
+        return payload unless definitions.any? { |d| deferred_definition?(d) || search_tool_definition?(d) }
+
+        betas = Array(payload.delete('anthropic_beta')) | Array(payload[:anthropic_beta]) | [BETA_FLAG]
+        payload[:anthropic_beta] = betas
+        payload
+      end
+
+      def server_block?(block)
+        SERVER_BLOCK_TYPES.include?(block_type(block))
+      end
+
+      def contains_server_blocks?(blocks)
+        Array(blocks).any? { |block| server_block?(block) }
+      end
+
+      def block_type(block)
+        return nil unless block.is_a?(Hash)
+
+        (block[:type] || block['type']).to_s
+      end
+
+      def deferred_definition?(definition)
+        return false unless definition.is_a?(Hash)
+
+        !!definition_value(definition, :defer_loading)
+      end
+
+      def search_tool_definition?(definition)
+        return false unless definition.is_a?(Hash)
+
+        definition_value(definition, :type) == SEARCH_TOOL_TYPE
+      end
+
+      def definition_value(definition, key)
+        definition[key] || definition[key.to_s]
+      end
+
+      def inject_search_tool!(definitions)
+        return if definitions.any? { |d| search_tool_definition?(d) }
+
+        collision = definitions.find { |d| definition_value(d, :name).to_s == SEARCH_TOOL_NAME }
+        if collision
+          raise RubyLLM::BedrockInvoke::Error,
+                "A tool named '#{SEARCH_TOOL_NAME}' collides with the injected tool search tool. Rename it."
+        end
+
+        definitions.unshift(type: SEARCH_TOOL_TYPE, name: SEARCH_TOOL_NAME)
+      end
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module BedrockInvoke
+    VERSION = '0.1.0.pre.1'
+  end
+end

data/lib/ruby_llm-bedrock_invoke.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+require 'ruby_llm'
+require 'erb'
+
+require_relative 'ruby_llm/bedrock_invoke/version'
+require_relative 'ruby_llm/bedrock_invoke/signing'
+require_relative 'ruby_llm/bedrock_invoke/event_stream'
+require_relative 'ruby_llm/bedrock_invoke/tool_search'
+require_relative 'ruby_llm/bedrock_invoke/deferred'
+require_relative 'ruby_llm/bedrock_invoke/block_collector'
+require_relative 'ruby_llm/bedrock_invoke/stream_state'
+require_relative 'ruby_llm/bedrock_invoke/capabilities'
+
+module RubyLLM
+  # Anthropic Claude on AWS Bedrock via InvokeModel — tool search (deferred
+  # tool loading), structured outputs, prompt caching and streaming inside
+  # the AWS data boundary, authenticated with the standard AWS credential chain.
+  module BedrockInvoke
+    class Error < StandardError; end
+
+    CONFIGURATION_OPTIONS = %i[
+      bedrock_invoke_region
+      bedrock_invoke_api_base
+      bedrock_invoke_credential_provider
+    ].freeze
+
+    # Where a response carried server-managed tool-search blocks
+    # (server_tool_use / tool_search_tool_result), the full content-block
+    # array is preserved for verbatim replay. On 1.x it rides in the message
+    # content as RawContent; on 2.0 (String-only content) it rides in this
+    # ivar and is re-injected by the protocol when rendering history.
+    RAW_BLOCKS_IVAR = :@bedrock_invoke_raw_blocks
+
+    # The preserved Anthropic content blocks for a message, or nil.
+    # Works on both RubyLLM generations.
+    def self.raw_blocks(message)
+      if defined?(RawContent) && message.content.is_a?(RawContent)
+        message.content.value
+      elsif message.instance_variable_defined?(RAW_BLOCKS_IVAR)
+        message.instance_variable_get(RAW_BLOCKS_IVAR)
+      end
+    end
+
+    # Bedrock routes /model/{modelId}/... with the id as a single path label;
+    # ARN-based ids (inference profiles, application inference profiles,
+    # imported models) contain '/' and ':' and must be percent-encoded the
+    # way the official AWS SDK encodes them.
+    def self.escape_model_id(model_id)
+      ERB::Util.url_encode(model_id.to_s)
+    end
+
+    # Bedrock Claude supports budget-based extended thinking only (the
+    # adaptive/effort API is not available on Bedrock). Assumed-exist models
+    # carry no registry reasoning_options, so the providers bypass the
+    # registry gates and emit this payload directly.
+    def self.budget_thinking_payload(thinking)
+      return nil unless thinking.respond_to?(:enabled?) && thinking.enabled?
+
+      effort = thinking.respond_to?(:effort) ? thinking.effort&.to_s : nil
+      return nil if effort == 'none'
+
+      budget = thinking.respond_to?(:budget) ? thinking.budget : thinking
+      unless budget.is_a?(Integer)
+        raise ArgumentError,
+              'Bedrock InvokeModel supports budget-based thinking only; adaptive/effort thinking ' \
+              'is not available on Bedrock. Use with_thinking(budget: N).'
+      end
+
+      { type: 'enabled', budget_tokens: budget }
+    end
+
+    # RubyLLM 2.0 split providers (where/auth) from protocols (wire format);
+    # 1.x providers are monolithic. Detect which world we're in at require time.
+    def self.v2_architecture?
+      defined?(RubyLLM::Protocol) ? true : false
+    end
+
+    def self.register!
+      register_configuration_options!
+      RubyLLM::Provider.register :bedrock_invoke, RubyLLM::Providers::BedrockInvoke
+    end
+
+    # 1.x Configuration#inspect redacts ivars by suffix (_key/_secret/_token...)
+    # but not _credential_provider, so a custom credentials object would print
+    # its secrets in consoles and error reporters. 2.0's redaction already
+    # covers the suffix.
+    module ConfigurationRedaction
+      def instance_variables
+        super.reject { |ivar| ivar.to_s.end_with?('_credential_provider') }
+      end
+    end
+
+    # 1.16+ and 2.0 register configuration_options automatically inside
+    # Provider.register; older 1.x needs the accessors added by hand.
+    def self.register_configuration_options!
+      RubyLLM::Configuration.prepend(ConfigurationRedaction) unless v2_architecture?
+      return if RubyLLM::Configuration.respond_to?(:register_provider_options)
+
+      missing = CONFIGURATION_OPTIONS.reject { |key| RubyLLM::Configuration.method_defined?(key) }
+      RubyLLM::Configuration.attr_accessor(*missing) if missing.any?
+    end
+  end
+end
+
+if RubyLLM::BedrockInvoke.v2_architecture?
+  require_relative 'ruby_llm/bedrock_invoke/protocol_v2'
+  require_relative 'ruby_llm/bedrock_invoke/provider_v2'
+else
+  require_relative 'ruby_llm/bedrock_invoke/raw_content'
+  require_relative 'ruby_llm/bedrock_invoke/provider_v1'
+end
+
+RubyLLM::BedrockInvoke.register!

metadata

@@ -0,0 +1,124 @@
+--- !ruby/object:Gem::Specification
+name: ruby_llm-bedrock_invoke
+version: !ruby/object:Gem::Version
+  version: 0.1.0.pre.1
+platform: ruby
+authors:
+- Chris Petersen
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: aws-eventstream
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1'
+- !ruby/object:Gem::Dependency
+  name: aws-sdk-core
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3'
+- !ruby/object:Gem::Dependency
+  name: base64
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: ruby_llm
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.13'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '3'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.13'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '3'
+description: Talks the native Anthropic Messages format to Bedrock InvokeModel / InvokeModelWithResponseStream,
+  unlocking Anthropic-specific features that the Converse API cannot express — tool
+  search with deferred tool loading, prompt caching via cache_control, and structured
+  outputs — while keeping data inside the AWS boundary. Authenticates with the standard
+  AWS credential chain (IRSA, Pod Identity, instance profiles, env vars) via aws-sdk-core.
+email:
+- chris@petersen.io
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- LICENSE
+- README.md
+- lib/ruby_llm-bedrock_invoke.rb
+- lib/ruby_llm/bedrock_invoke/block_collector.rb
+- lib/ruby_llm/bedrock_invoke/capabilities.rb
+- lib/ruby_llm/bedrock_invoke/deferred.rb
+- lib/ruby_llm/bedrock_invoke/event_stream.rb
+- lib/ruby_llm/bedrock_invoke/protocol_v2.rb
+- lib/ruby_llm/bedrock_invoke/provider_v1.rb
+- lib/ruby_llm/bedrock_invoke/provider_v2.rb
+- lib/ruby_llm/bedrock_invoke/raw_content.rb
+- lib/ruby_llm/bedrock_invoke/signing.rb
+- lib/ruby_llm/bedrock_invoke/stream_state.rb
+- lib/ruby_llm/bedrock_invoke/tool_search.rb
+- lib/ruby_llm/bedrock_invoke/version.rb
+homepage: https://github.com/scientist-labs/ruby_llm-bedrock_invoke
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/scientist-labs/ruby_llm-bedrock_invoke
+  source_code_uri: https://github.com/scientist-labs/ruby_llm-bedrock_invoke
+  rubygems_mfa_required: 'true'
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.1.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.6.9
+specification_version: 4
+summary: RubyLLM provider for Anthropic Claude via the AWS Bedrock InvokeModel API
+test_files: []