ruby_llm-bedrock_invoke 0.1.0.pre.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 911a2cdfef791e621f685c7560906cd38bb3c6362303d7a488799617c0a66d65
+  data.tar.gz: 76efc4ffcf2932e7baf423a9498044f9e47eccc9bb478bd56f75adb4f3c1f95f
+SHA512:
+  metadata.gz: b5c1b976369f15fc99ad75f1b6661fda71ec07ceafc2db7595af8d8ad159525c152dd5b9f22187c43b389813b8d14e5b9346e4f8fb21cd818a7f61013635ffe0
+  data.tar.gz: 0c03c773fb9f8ea9505a90734e677db5d3803c6c5aeeabe6e67574d0dc0dc6862f4c6f7fac90b6e27fb80d8bdfd4e0e82cc7c5f2b10886ce658b6b07bbf3f10b

data/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Chris Petersen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/README.md

@@ -0,0 +1,194 @@
+# ruby_llm-bedrock_invoke
+
+A [RubyLLM](https://github.com/crmne/ruby_llm) provider for **Anthropic Claude via the AWS Bedrock InvokeModel API**, registered as `:bedrock_invoke`. Works with released RubyLLM **1.13+** and the unreleased **2.0** architecture from one gem.
+
+## Why this exists
+
+RubyLLM's built-in `:bedrock` provider speaks AWS's unified [Converse API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html). Converse is great for portability, but it is a translation layer — and Anthropic features that don't fit its schema simply don't exist there. As of mid-2026, if you need Claude-specific capabilities *and* your data must stay inside the AWS boundary, the feature matrix looks like this:
+
+| Capability | Bedrock Converse | Bedrock Mantle (`/anthropic/v1/messages`) | **Bedrock InvokeModel** | First-party Anthropic API |
+|---|:---:|:---:|:---:|:---:|
+| Tool search / deferred tool loading | ❌ impossible | ❌ not offered | ✅ beta | ✅ GA |
+| Structured outputs (`output_config`) | ✅ | ❌ rejected (400) | ✅ GA | ✅ |
+| Prompt caching (`cache_control`) | ✅ (`cachePoint`) | ✅ | ✅ GA | ✅ |
+| Streaming | ✅ | ✅ | ✅ | ✅ |
+| Data stays inside AWS | ✅ | ✅ | ✅ | ❌ |
+
+InvokeModel is the only surface with the full set. It takes the native Anthropic Messages body verbatim (plus `anthropic_version`), so this gem inherits RubyLLM's entire Anthropic wire format and swaps only the transport: SigV4 signing against `bedrock-runtime` and AWS event-stream decoding.
+
+### The problem deferred tool loading solves
+
+Agents with large tool catalogs (many MCP servers, dozens of tools) face an ugly trade-off:
+
+- **Load every tool eagerly** → tool schemas consume thousands of prompt tokens on every request.
+- **Load tools dynamically** (intent routing, per-request tool selection) → the tools array is part of the prompt-cache prefix, so *any* change to it invalidates the cache for the whole conversation.
+
+Anthropic's [tool search tool](https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool) is the first-party fix: mark long-tail tools `defer_loading: true` and they are **excluded from the context and the cache prefix** until the model discovers them via a server-side search tool. Discovery happens *inline in the message stream* — the prefix never changes, so the cache survives.
+
+This is measurable. `bin/cache_smoke` runs both strategies live against Bedrock with a ~9k-token cached system prefix:
+
+```
+RUN A — deferred tools + tool search        RUN B — inject tool def mid-conversation
+  response              cache READ  WRITE     response              cache READ  WRITE
+  R1 search + call            9064   9064     R1 tool call                   0   8851
+  R2 final answer             9064      0     R2 final answer             8851      0
+  R3 tool call (turn 2)       9064      0     R3 tool call (def added!)      0   8936  <- cache busted
+  R4 final answer             9064      0     R4 final answer             8936      0
+```
+
+Run A discovers and calls a hidden tool with the cache intact on every subsequent request. Run B "loads" the same tool by adding its definition — and pays a full cache re-write. That contrast is the reason this gem targets InvokeModel: on Bedrock, tool search exists nowhere else.
+
+## Installation
+
+```ruby
+gem 'ruby_llm-bedrock_invoke'
+```
+
+Dependencies: `ruby_llm` (>= 1.13, < 3), `aws-sdk-core` (credential chain + SigV4), `aws-eventstream` (streaming).
+
+## Quick start
+
+```ruby
+require 'ruby_llm-bedrock_invoke'
+
+RubyLLM.configure do |config|
+  config.bedrock_invoke_region = 'us-east-1' # or rely on AWS_REGION
+  # No credential config needed: the standard AWS chain is used automatically.
+end
+
+chat = RubyLLM.chat(
+  model: 'global.anthropic.claude-haiku-4-5-20251001-v1:0',
+  provider: :bedrock_invoke,
+  assume_model_exists: true # Bedrock ids/ARNs aren't in RubyLLM's registry
+)
+
+chat.ask 'Hello from inside the AWS boundary!'
+```
+
+Everything RubyLLM documents for chat works unchanged: streaming (`chat.ask('...') { |chunk| print chunk.content }`), tools, `with_temperature`, `with_params`.
+
+### Authentication: why the default chain
+
+Credentials resolve exactly like any AWS SDK client — env vars, shared config/SSO, ECS/EKS **Pod Identity**, **IRSA** web-identity tokens, IMDS — resolved once per process and cached (the credentials object self-refreshes). A token projected into a pod just works with zero RubyLLM configuration. This is deliberately different from core's `:bedrock` provider, which requires static keys or a hand-built credential provider object. To override: `config.bedrock_invoke_credential_provider = anything_responding_to_credentials`.
+
+## Tool search / deferred tool loading
+
+Mark long-tail tools deferred; keep your 3–5 hottest tools eager. When at least one tool is deferred, the provider automatically flags those definitions `defer_loading: true`, injects the `tool_search_tool_regex` server tool, and opts into the `tool-search-tool-2025-10-19` beta (in the request body — Bedrock has no beta HTTP header).
+
+```ruby
+class DatabaseQuery < RubyLLM::Tool
+  include RubyLLM::BedrockInvoke::Deferred   # class-level: always deferred
+  description 'Run a read-only SQL query against the warehouse'
+  # ...
+end
+
+chat.with_tools(HotTool, DatabaseQuery)
+```
+
+For tool instances you don't control (e.g. built by `ruby_llm-mcp`), defer per-instance:
+
+```ruby
+chat.with_tools(hot_tool, *mcp_tools.map { |t| RubyLLM::BedrockInvoke.defer(t) })
+```
+
+**You must tell the model that hidden tools exist.** In live testing, models look at their visible tool list, conclude the capability is missing, and never search. One paragraph in your (cached) system prompt fixes it:
+
+```ruby
+chat.with_instructions <<~PROMPT
+  Your visible tool list is only a subset of the tools available. Before saying
+  a capability is missing, use tool_search_tool_regex to search the full catalog
+  by keyword. Discovered tools can then be called normally.
+PROMPT
+```
+
+### What happens at runtime
+
+1. The model calls the search tool server-side (`server_tool_use` block) with a regex.
+2. The API returns a `tool_search_tool_result` containing `tool_reference` blocks; matching deferred tools are expanded into context **inline in the message stream** — the cached prefix is untouched.
+3. The model calls the discovered tool like any other; RubyLLM's tool loop executes it normally.
+
+Responses that used tool search carry the full server content blocks, which this gem replays verbatim on subsequent requests (required by the API). Access them on either RubyLLM generation with:
+
+```ruby
+blocks = RubyLLM::BedrockInvoke.raw_blocks(message) # nil for ordinary messages
+```
+
+If you persist conversations (Rails `acts_as_chat`), these blocks are process-local — persist and restore them yourself if tool-search turns must survive a reload.
+
+## Prompt caching
+
+Caching is GA and on by default on InvokeModel — you only place breakpoints. The prefix is cached in order `tools → system → messages`, so one `cache_control` breakpoint at the end of your system prompt covers the tool schemas and instructions:
+
+```ruby
+chat = RubyLLM.chat(model: MODEL, provider: :bedrock_invoke, assume_model_exists: true)
+chat.with_params(system: [                       # with_provider_options on RubyLLM 2.0
+  { type: 'text', text: BIG_SYSTEM_PROMPT, cache_control: { type: 'ephemeral' } }
+])
+```
+
+Note that `with_params(system:)` **replaces** the rendered system prompt: if you also use `with_instructions` (e.g. for the tool-search guidance above), fold that text into the blocks array rather than setting both. On the unreleased RubyLLM 2.0 you can use the native `chat.with_caching` API instead.
+
+Verify it's working by reading usage off the response:
+
+```ruby
+response = chat.ask 'First question'
+response.cache_creation_tokens  # > 0: the prefix was written    (2.0: response.tokens.cache_creation)
+response.cached_tokens          # subsequent asks: > 0 = cache hit (2.0: response.tokens.cached)
+```
+
+Why this matters with deferred tools: deferred definitions are excluded from the cached prefix (a deferred tool cannot even carry `cache_control` — the API 400s), so the catalog can grow or shrink *between conversations* and discovery can happen *within* one, all without invalidating the cache. Mind the per-model minimum cacheable size (4,096 tokens for Haiku 4.5-class models) — below it you'll see neither reads nor writes.
+
+## Structured output, streaming, thinking
+
+```ruby
+# Structured output — GA on InvokeModel, no beta flag:
+chat.with_schema(
+  type: 'object',
+  properties: { answer: { type: 'string' } },
+  required: ['answer'],
+  additionalProperties: false
+).ask('...')
+
+# Streaming — InvokeModelWithResponseStream (AWS event-stream framing, decoded for you):
+chat.ask('Tell me a story') { |chunk| print chunk.content }
+
+# Extended thinking — budget-based only (Bedrock has no adaptive/effort API):
+chat.with_thinking(budget: 8_000).ask('Think hard about this')
+```
+
+## Configuration
+
+| Key | Meaning | Default |
+|---|---|---|
+| `bedrock_invoke_region` | AWS region | `AWS_REGION` / `AWS_DEFAULT_REGION` |
+| `bedrock_invoke_api_base` | Endpoint override (host only, no path) | `https://bedrock-runtime.{region}.amazonaws.com` |
+| `bedrock_invoke_credential_provider` | Any object responding to `#credentials` | AWS default chain, cached per region |
+
+## Notes and limits
+
+- **Model coverage**: Claude Sonnet 5 is not served through the InvokeModel surface; Fable 5, Opus 4.8/4.7, Haiku 4.5 and earlier are. Model IDs are not registry-validated, so ARNs, inference profiles, and `global.`/regional prefixes all work — but resolved models carry generic capability metadata: no pricing, and `max_tokens` defaults to 4096 (raise it with `chat.with_params(max_tokens: 32_000)`).
+- Tool search on Bedrock is a **beta** and only offers the regex variant (`tool_search_tool_regex`); no BM25 variant, unlike the first-party API.
+- **Betas ride in the body**: Bedrock InvokeModel has no `anthropic-beta` HTTP header. Opt into additional betas via `chat.with_params(anthropic_beta: [...])` — the tool-search flag is re-applied automatically if your params replace the array.
+- **Keep deferred tools attached** for the lifetime of a conversation that used tool search: past turns replay `tool_reference` blocks that must still resolve to sent tool definitions.
+- On a retried stream (throttling mid-stream), your streaming block may see chunks from the abandoned attempt; the final returned message contains only the successful attempt.
+- Embeddings, image generation, Batches, and the Files API are out of scope (use the stock `:bedrock` provider or first-party Anthropic).
+
+## Development
+
+```bash
+bundle exec rspec                                                # newest 1.x
+BUNDLE_GEMFILE=gemfiles/ruby_llm_1_13.gemfile bundle exec rspec  # oldest supported 1.x
+BUNDLE_GEMFILE=gemfiles/ruby_llm_2_0.gemfile bundle exec rspec   # 2.0: upstream main (needs Ruby >= 3.4;
+                                                                 #      RUBY_LLM_PATH=... for a local checkout)
+
+bundle exec bin/smoke        # live AWS: completion, streaming, schema, tool search
+bundle exec bin/cache_smoke  # live AWS: the cache-preservation proof shown above
+```
+
+## Why a plugin (and the road upstream)
+
+Upstream RubyLLM removed its original InvokeModel chat path in 1.12.0 in favor of Converse, and the maintainer closed the community tool-search PR ([crmne/ruby_llm#745](https://github.com/crmne/ruby_llm/pull/745)) asking for a provider-agnostic design first — [explicitly noting](https://github.com/crmne/ruby_llm/pull/745#issuecomment-4879003367) that *"Bedrock exposes Claude's version through InvokeModel, though not Converse."* The open feature request is [crmne/ruby_llm#748](https://github.com/crmne/ruby_llm/issues/748). This gem is the working Bedrock-side reference implementation for that design conversation: `#745`'s `defer:`/`deferred` user-facing API plus this provider underneath would cover both the first-party and Bedrock dialects.
+
+## License
+
+MIT

data/lib/ruby_llm/bedrock_invoke/block_collector.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+require 'json'
+
+module RubyLLM
+  module BedrockInvoke
+    # Reassembles Anthropic content blocks from streaming events.
+    #
+    # Two jobs:
+    # 1. Rebuild the assistant's full ordered content-block array so a
+    #    response containing tool-search blocks can be replayed verbatim
+    #    (see RawContent).
+    # 2. Identify events belonging to server-managed blocks
+    #    (server_tool_use / tool_search_tool_result) so they can be kept away
+    #    from RubyLLM's tool-call accumulator — its input_json_delta handling
+    #    keys nil-id fragments to the latest client tool call, which a
+    #    server_tool_use block would silently corrupt.
+    class BlockCollector
+      def initialize
+        @blocks = {}
+        @json_buffers = {}
+        @server_indexes = {}
+        @saw_server_block = false
+      end
+
+      def observe(event)
+        case event['type']
+        when 'content_block_start' then start_block(event)
+        when 'content_block_delta' then apply_delta(event)
+        when 'content_block_stop' then finish_block(event)
+        end
+      end
+
+      # True for content-block events that belong to a server-managed block.
+      def suppress?(event)
+        index = event['index']
+        return false if index.nil?
+
+        @server_indexes.key?(index) && event['type'].to_s.start_with?('content_block')
+      end
+
+      def server_blocks?
+        @saw_server_block
+      end
+
+      def content_blocks
+        @blocks.keys.sort.map { |index| @blocks[index] }
+      end
+
+      private
+
+      def start_block(event)
+        index = event['index']
+        block = event['content_block'] || {}
+        @blocks[index] = block
+
+        if ToolSearch.server_block?(block)
+          @saw_server_block = true
+          @server_indexes[index] = true
+        end
+
+        @json_buffers[index] = +'' if %w[tool_use server_tool_use].include?(block['type'])
+      end
+
+      def apply_delta(event)
+        block = @blocks[event['index']]
+        delta = event['delta'] || {}
+        return if block.nil?
+
+        case delta['type']
+        when 'text_delta'
+          append_string(block, 'text', delta['text'])
+        when 'input_json_delta'
+          buffer = @json_buffers[event['index']]
+          buffer << delta['partial_json'].to_s if buffer
+        when 'thinking_delta'
+          append_string(block, 'thinking', delta['thinking'])
+        when 'signature_delta'
+          block['signature'] = delta['signature']
+        when 'citations_delta'
+          (block['citations'] ||= []) << delta['citation']
+        end
+      end
+
+      # In-place append (amortized O(n) over a stream, unlike interpolation).
+      def append_string(block, key, piece)
+        return if piece.nil?
+
+        existing = block[key]
+        block[key] = existing = +'' unless existing.is_a?(String) && !existing.frozen?
+        existing << piece
+      end
+
+      def finish_block(event)
+        index = event['index']
+        buffer = @json_buffers.delete(index)
+        return if buffer.nil? || buffer.empty?
+
+        block = @blocks[index]
+        return if block.nil?
+
+        begin
+          block['input'] = JSON.parse(buffer)
+        rescue JSON::ParserError
+          RubyLLM.logger.debug { "bedrock_invoke: unparseable tool input for block #{index}" }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/capabilities.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module BedrockInvoke
+    # Delegates capability lookups to the Anthropic capabilities module after
+    # normalizing Bedrock model ids (region prefixes, vendor prefix, version
+    # suffixes) into first-party Anthropic ids, so context windows, pricing
+    # and feature flags resolve correctly for ids like
+    # `us.anthropic.claude-sonnet-4-5-20250929-v1:0`.
+    module Capabilities
+      module_function
+
+      def normalize_model_id(model_id)
+        model_id.to_s
+                .sub(/\A(?:us|eu|apac|global)\./, '')
+                .sub(/\Aanthropic\./, '')
+                .sub(/-v\d+(?::\d+)?\z/, '')
+      end
+
+      def anthropic_capabilities
+        RubyLLM::Providers::Anthropic::Capabilities
+      end
+
+      def method_missing(name, *args, &)
+        target = anthropic_capabilities
+        return super unless target.respond_to?(name)
+
+        args = args.dup
+        args[0] = normalize_model_id(args[0]) if args[0].is_a?(String)
+        target.public_send(name, *args, &)
+      end
+
+      def respond_to_missing?(name, include_private = false)
+        anthropic_capabilities.respond_to?(name) || super
+      end
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/deferred.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module BedrockInvoke
+    # Marks a tool for deferred loading (Anthropic tool search).
+    #
+    #   class DatabaseQuery < RubyLLM::Tool
+    #     include RubyLLM::BedrockInvoke::Deferred
+    #     ...
+    #   end
+    #
+    # or per-instance (e.g. for tools built by ruby_llm-mcp):
+    #
+    #   chat.with_tool(RubyLLM::BedrockInvoke.defer(tool))
+    #
+    # Any tool object responding truthily to +defer_loading?+ is treated as
+    # deferred — including tools that define the method themselves.
+    module Deferred
+      def defer_loading?
+        true
+      end
+    end
+
+    def self.defer(tool)
+      if tool.is_a?(Class)
+        # Chat#with_tools instantiates classes, so extending the class object
+        # would be silently lost. Fail loudly instead.
+        raise ArgumentError,
+              'Pass a tool instance to RubyLLM::BedrockInvoke.defer (e.g. defer(MyTool.new)), ' \
+              'or include RubyLLM::BedrockInvoke::Deferred in the tool class.'
+      end
+
+      tool.extend(Deferred) unless ToolSearch.deferred?(tool)
+      tool
+    end
+  end
+end

data/lib/ruby_llm/bedrock_invoke/event_stream.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+
+require 'base64'
+require 'json'
+
+module RubyLLM
+  module BedrockInvoke
+    # Decodes AWS event-stream framing (application/vnd.amazon.eventstream)
+    # from InvokeModelWithResponseStream into plain Anthropic streaming events.
+    #
+    # Each `chunk` event carries a PayloadPart whose base64 `bytes` field
+    # decodes to one standard Anthropic SSE event (message_start,
+    # content_block_delta, ...). Failures surface three ways, all handled
+    # before the payload is trusted: frames with :message-type 'exception'
+    # (typed via :exception-type, payload may be empty), frames with
+    # :message-type 'error' (:error-code/:error-message headers, empty
+    # payload), and exception-suffixed keys inside chunk payloads.
+    module EventStream
+      EXCEPTION_STATUS = {
+        'throttlingException' => 429,
+        'validationException' => 400,
+        'modelTimeoutException' => 408,
+        'modelStreamErrorException' => 424,
+        'serviceUnavailableException' => 503,
+        'internalServerException' => 500
+      }.freeze
+
+      def event_stream_decoder
+        require 'aws-eventstream'
+        ::Aws::EventStream::Decoder.new
+      end
+
+      # Yields one decoded Anthropic event Hash per event in the raw chunk.
+      # Raises through ErrorMiddleware when the stream carries an exception event.
+      def each_event_stream_event(decoder, raw_chunk, &)
+        message, eof = decoder.decode_chunk(raw_chunk)
+        process_event_stream_message(message, &) if message
+        until eof
+          message, eof = decoder.decode_chunk
+          process_event_stream_message(message, &) if message
+        end
+      end
+
+      private
+
+      def process_event_stream_message(message)
+        payload = parse_event_stream_payload(message)
+        raise_frame_errors!(message, payload)
+        return if payload.nil?
+
+        event = unwrap_payload_bytes(payload)
+        return if event.nil?
+
+        raise_embedded_exception!(event)
+        yield event
+      end
+
+      # Exception frames may carry empty payloads; error frames carry only
+      # :error-code/:error-message headers. Both must raise before the
+      # payload-nil early return or the stream would complete as success.
+      def raise_frame_errors!(message, payload)
+        message_type = event_stream_header(message, ':message-type')
+        exception_type = event_stream_header(message, ':exception-type')
+
+        if exception_type || message_type == 'exception'
+          raise_event_stream_error(exception_type || 'exception', payload)
+        elsif message_type == 'error'
+          code = event_stream_header(message, ':error-code')
+          raise_event_stream_error(code || 'error', event_stream_header(message, ':error-message') || code.to_s)
+        end
+      end
+
+      def raise_embedded_exception!(event)
+        error_key = event.keys.find { |key| key.to_s.end_with?('Exception') }
+        raise_event_stream_error(error_key, event[error_key] || event) if error_key
+      end
+
+      def event_stream_header(message, name)
+        header = message.headers[name]
+        header&.value
+      end
+
+      def parse_event_stream_payload(message)
+        raw = message.payload.read
+        return nil if raw.nil? || raw.empty?
+
+        JSON.parse(raw)
+      rescue JSON::ParserError
+        RubyLLM.logger.debug { "bedrock_invoke: undecodable event payload: #{raw.inspect}" }
+        nil
+      end
+
+      def unwrap_payload_bytes(payload)
+        return payload unless payload.is_a?(Hash) && payload.key?('bytes')
+
+        JSON.parse(Base64.decode64(payload['bytes']))
+      rescue JSON::ParserError
+        RubyLLM.logger.debug { 'bedrock_invoke: undecodable bytes payload in stream event' }
+        nil
+      end
+
+      def raise_event_stream_error(type, detail)
+        status = EXCEPTION_STATUS.fetch(type.to_s, 500)
+        message = detail.is_a?(Hash) ? detail['message'] || detail['Message'] : detail&.to_s
+        message = type.to_s if message.nil? || message.empty?
+        response = Struct.new(:body, :status).new({ 'message' => message }, status)
+        RubyLLM::ErrorMiddleware.parse_error(provider: error_middleware_provider, response: response)
+      end
+
+      # The object ErrorMiddleware treats as the provider (self for 1.x
+      # providers; overridden by 2.0 protocols to hand back the provider).
+      def error_middleware_provider
+        self
+      end
+    end
+  end
+end