dify_llm 1.9.2 → 1.14.1

data/lib/ruby_llm/providers/bedrock/models.rb

@@ -3,100 +3,139 @@
 module RubyLLM
   module Providers
     class Bedrock
-      # Models methods for the AWS Bedrock API implementation
+      # Models methods for AWS Bedrock.
       module Models
-        def list_models
-          mgmt_api_base = "https://bedrock.#{@config.bedrock_region}.amazonaws.com"
-          full_models_url = "#{mgmt_api_base}/#{models_url}"
-          signature = sign_request(full_models_url, method: :get)
-          response = @connection.get(full_models_url) do |req|
-            req.headers.merge! signature.headers
-          end
+        module_function
 
-          parse_list_models_response(response, slug, capabilities)
-        end
+        REGION_PREFIXES = %w[global us eu ap sa ca me af il].freeze
 
-        module_function
+        def models_api_base
+          "https://bedrock.#{bedrock_region}.amazonaws.com"
+        end
 
         def models_url
-          'foundation-models'
+          '/foundation-models'
         end
 
-        def parse_list_models_response(response, slug, capabilities)
-          models = Array(response.body['modelSummaries'])
-
-          models.select { |m| m['modelId'].include?('claude') }.map do |model_data|
-            model_id = model_data['modelId']
-
-            Model::Info.new(
-              id: model_id_with_region(model_id, model_data),
-              name: model_data['modelName'] || capabilities.format_display_name(model_id),
-              provider: slug,
-              family: capabilities.model_family(model_id),
-              created_at: nil,
-              context_window: capabilities.context_window_for(model_id),
-              max_output_tokens: capabilities.max_tokens_for(model_id),
-              modalities: capabilities.modalities_for(model_id),
-              capabilities: capabilities.capabilities_for(model_id),
-              pricing: capabilities.pricing_for(model_id),
-              metadata: {
-                provider_name: model_data['providerName'],
-                inference_types: model_data['inferenceTypesSupported'] || [],
-                streaming_supported: model_data['responseStreamingSupported'] || false,
-                input_modalities: model_data['inputModalities'] || [],
-                output_modalities: model_data['outputModalities'] || []
-              }
-            )
+        def parse_list_models_response(response, slug, _capabilities)
+          Array(response.body['modelSummaries']).map do |model_data|
+            create_model_info(model_data, slug)
           end
         end
 
-        def create_model_info(model_data, slug, _capabilities)
-          model_id = model_data['modelId']
+        def create_model_info(model_data, slug, _capabilities = nil)
+          model_id = model_id_with_region(model_data['modelId'], model_data)
+          converse_data = model_data['converse'] || {}
 
           Model::Info.new(
-            id: model_id_with_region(model_id, model_data),
-            name: model_data['modelName'] || model_id,
+            id: model_id,
+            name: model_data['modelName'],
             provider: slug,
-            family: 'claude',
+            family: model_data['modelFamily'] || model_data['providerName']&.downcase,
             created_at: nil,
-            context_window: 200_000,
-            max_output_tokens: 4096,
-            modalities: { input: ['text'], output: ['text'] },
-            capabilities: [],
+            context_window: parse_context_window(model_data),
+            max_output_tokens: converse_data['maxTokensDefault'] || converse_data['maxTokensMaximum'],
+            modalities: {
+              input: normalize_modalities(model_data['inputModalities']),
+              output: normalize_modalities(model_data['outputModalities'])
+            },
+            capabilities: parse_capabilities(model_data),
             pricing: {},
-            metadata: {}
+            metadata: {
+              provider_name: model_data['providerName'],
+              model_arn: model_data['modelArn'],
+              inference_types: model_data['inferenceTypesSupported'],
+              converse: converse_data
+            }
           )
         end
 
         def model_id_with_region(model_id, model_data)
-          normalize_inference_profile_id(
-            model_id,
-            model_data['inferenceTypesSupported'],
-            @config.bedrock_region
-          )
+          inference_types = Array(model_data['inferenceTypesSupported'])
+          normalize_inference_profile_id(model_id, inference_types, @config.bedrock_region)
         end
 
-        def region_prefix(region)
-          region = region.to_s
-          return 'us' if region.empty?
+        def normalize_inference_profile_id(model_id, inference_types, region)
+          return model_id unless inference_types.include?('INFERENCE_PROFILE')
+          return model_id if inference_types.include?('ON_DEMAND')
 
-          region[0, 2]
+          with_region_prefix(model_id, region)
         end
 
         def with_region_prefix(model_id, region)
-          desired_prefix = region_prefix(region)
-          return model_id if model_id.start_with?("#{desired_prefix}.")
+          prefix = region_prefix(region)
 
-          clean_model_id = model_id.sub(/^[a-z]{2}\./, '')
-          "#{desired_prefix}.#{clean_model_id}"
+          if region_prefixed?(model_id)
+            model_id.sub(/\A(?:#{REGION_PREFIXES.join('|')})\./, "#{prefix}.")
+          else
+            "#{prefix}.#{model_id}"
+          end
         end
 
-        def normalize_inference_profile_id(model_id, inference_types, region)
-          types = Array(inference_types)
-          return model_id unless types.include?('INFERENCE_PROFILE')
-          return model_id if types.include?('ON_DEMAND')
+        def region_prefix(region)
+          prefix = region.to_s.split('-').first
+          prefix = '' if prefix.nil?
+          prefix.empty? ? 'us' : prefix
+        end
 
-          with_region_prefix(model_id, region)
+        def region_prefixed?(model_id)
+          model_id.match?(/\A(?:#{REGION_PREFIXES.join('|')})\./)
+        end
+
+        def normalize_modalities(modalities)
+          Array(modalities).map do |modality|
+            normalized = modality.to_s.downcase
+            case normalized
+            when 'embedding' then 'embeddings'
+            when 'speech' then 'audio'
+            else normalized
+            end
+          end
+        end
+
+        def parse_capabilities(model_data)
+          capabilities = []
+          capabilities << 'streaming' if model_data['responseStreamingSupported']
+
+          converse = model_data['converse'] || {}
+          capabilities << 'function_calling' if converse.is_a?(Hash)
+          capabilities << 'reasoning' if converse.dig('reasoningSupported', 'embedded')
+          capabilities << 'structured_output' if supports_structured_output?(model_data['modelId'])
+
+          capabilities
+        end
+
+        # Structured output supported on Claude 4.5+ and assumed for future major versions.
+        # Bedrock IDs look like: us.anthropic.claude-haiku-4-5-20251001-v1:0
+        # Must handle optional region prefix (us./eu./global.) and anthropic. prefix.
+        def supports_structured_output?(model_id)
+          return false unless model_id
+
+          normalized = model_id.sub(/\A(?:#{REGION_PREFIXES.join('|')})\./, '').delete_prefix('anthropic.')
+          match = normalized.match(/claude-(?:opus|sonnet|haiku)-(\d+)-(\d{1,2})(?:\b|-)/)
+          return false unless match
+
+          major = match[1].to_i
+          minor = match[2].to_i
+          major > 4 || (major == 4 && minor >= 5)
+        end
+
+        def reasoning_embedded?(model)
+          metadata = RubyLLM::Utils.deep_symbolize_keys(model.metadata || {})
+          converse = metadata[:converse] || {}
+          reasoning_supported = converse[:reasoningSupported] || {}
+          reasoning_supported[:embedded] || false
+        end
+
+        def parse_context_window(model_data)
+          value = model_data.dig('description', 'maxContextWindow')
+          return unless value.is_a?(String)
+
+          if value.match?(/\A\d+[kK]\z/)
+            value.to_i * 1000
+          elsif value.match?(/\A\d+\z/)
+            value.to_i
+          end
         end
       end
     end

data/lib/ruby_llm/providers/bedrock/streaming.rb

@@ -1,17 +1,318 @@
 # frozen_string_literal: true
 
-require_relative 'streaming/base'
-require_relative 'streaming/content_extraction'
-require_relative 'streaming/message_processing'
-require_relative 'streaming/payload_processing'
-require_relative 'streaming/prelude_handling'
-
 module RubyLLM
   module Providers
     class Bedrock
-      # Streaming implementation for the AWS Bedrock API.
+      # Streaming implementation for Bedrock ConverseStream (AWS Event Stream).
       module Streaming
-        include Base
+        private
+
+        def stream_url
+          "/model/#{@model.id}/converse-stream"
+        end
+
+        def stream_response(connection, payload, additional_headers = {}, &block)
+          accumulator = StreamAccumulator.new
+          decoder = event_stream_decoder
+          request_payload = api_payload(payload)
+          body = JSON.generate(request_payload)
+
+          response = connection.post(stream_url, request_payload) do |req|
+            req.headers.merge!(sign_headers('POST', stream_url, body))
+            req.headers.merge!(additional_headers) unless additional_headers.empty?
+            req.headers['Accept'] = 'application/vnd.amazon.eventstream'
+
+            if Faraday::VERSION.start_with?('1')
+              req.options[:on_data] = proc do |chunk, _size|
+                parse_stream_chunk(decoder, chunk, accumulator, &block)
+              end
+            else
+              req.options.on_data = proc do |chunk, _bytes, env|
+                if env&.status == 200
+                  parse_stream_chunk(decoder, chunk, accumulator, &block)
+                else
+                  handle_failed_stream(chunk, env)
+                end
+              end
+            end
+          end
+
+          message = accumulator.to_message(response)
+          RubyLLM.logger.debug { "Stream completed: #{message.content}" }
+          message
+        end
+
+        def event_stream_decoder
+          require 'aws-eventstream'
+          Aws::EventStream::Decoder.new
+        rescue LoadError
+          raise Error,
+                'The aws-eventstream gem is required for Bedrock streaming. ' \
+                'Please add it to your Gemfile: gem "aws-eventstream"'
+        end
+
+        def handle_failed_stream(chunk, env)
+          data = JSON.parse(chunk)
+          error_response = env.merge(body: data)
+          ErrorMiddleware.parse_error(provider: self, response: error_response)
+        rescue JSON::ParserError
+          RubyLLM.logger.debug { "Failed Bedrock stream error chunk: #{chunk}" }
+        end
+
+        def parse_stream_chunk(decoder, raw_chunk, accumulator)
+          handle_non_eventstream_error_chunk(raw_chunk)
+
+          decode_events(decoder, raw_chunk).each do |event|
+            chunk = build_chunk(event)
+            next unless chunk
+
+            accumulator.add(chunk)
+            yield chunk
+          end
+        end
+
+        def handle_non_eventstream_error_chunk(raw_chunk)
+          text = raw_chunk.to_s
+
+          if text.start_with?('event: error')
+            payload = text.lines.find { |line| line.start_with?('data:') }&.delete_prefix('data:')&.strip
+            raise_streaming_chunk_error(payload) if payload
+            return
+          end
+
+          return unless text.lstrip.start_with?('{') && text.include?('"error"')
+
+          raise_streaming_chunk_error(text)
+        end
+
+        def raise_streaming_chunk_error(payload)
+          parsed = JSON.parse(payload)
+          message = parsed.dig('error', 'message') || parsed['message'] || 'Bedrock streaming error'
+          response = Struct.new(:body, :status).new({ 'message' => message }, 500)
+          ErrorMiddleware.parse_error(provider: self, response: response)
+        rescue JSON::ParserError
+          nil
+        end
+
+        def decode_events(decoder, raw_chunk)
+          events = []
+          message, eof = decoder.decode_chunk(raw_chunk)
+
+          while message
+            event = decode_event_payload(message.payload.read)
+            if event && RubyLLM.config.log_stream_debug
+              RubyLLM.logger.debug do
+                "Bedrock stream event keys: #{event.keys}"
+              end
+            end
+            events << event if event
+            break if eof
+
+            message, eof = decoder.decode_chunk
+          end
+
+          events
+        end
+
+        def decode_event_payload(payload)
+          outer = JSON.parse(payload)
+
+          if outer['bytes'].is_a?(String)
+            JSON.parse(Base64.decode64(outer['bytes']))
+          else
+            outer
+          end
+        rescue JSON::ParserError => e
+          RubyLLM.logger.debug { "Failed to decode Bedrock stream event payload: #{e.message}" }
+          nil
+        end
+
+        def build_chunk(event)
+          raise_stream_error(event) if stream_error_event?(event)
+
+          metadata_usage, usage, message_usage = event_usage(event)
+
+          Chunk.new(
+            role: :assistant,
+            model_id: event['modelId'] || event.dig('message', 'model') || @model&.id,
+            content: extract_content_delta(event),
+            thinking: Thinking.build(
+              text: extract_thinking_delta(event),
+              signature: extract_thinking_signature(event)
+            ),
+            tool_calls: extract_tool_calls(event),
+            input_tokens: extract_input_tokens(metadata_usage, usage, message_usage),
+            output_tokens: extract_output_tokens(metadata_usage, usage),
+            cached_tokens: extract_cached_tokens(metadata_usage, usage),
+            cache_creation_tokens: extract_cache_creation_tokens(metadata_usage, usage),
+            thinking_tokens: extract_reasoning_tokens(metadata_usage, usage)
+          )
+        end
+
+        def event_usage(event)
+          [
+            event.dig('metadata', 'usage') || {},
+            event['usage'] || {},
+            event.dig('message', 'usage') || {}
+          ]
+        end
+
+        def extract_input_tokens(metadata_usage, usage, message_usage)
+          metadata_usage['inputTokens'] || usage['inputTokens'] || message_usage['input_tokens']
+        end
+
+        def extract_output_tokens(metadata_usage, usage)
+          metadata_usage['outputTokens'] || usage['outputTokens'] || usage['output_tokens']
+        end
+
+        def extract_cached_tokens(metadata_usage, usage)
+          metadata_usage['cacheReadInputTokens'] || usage['cacheReadInputTokens'] || usage['cache_read_input_tokens']
+        end
+
+        def extract_cache_creation_tokens(metadata_usage, usage)
+          metadata_usage['cacheWriteInputTokens'] || usage['cacheWriteInputTokens'] ||
+            usage['cache_creation_input_tokens']
+        end
+
+        def extract_reasoning_tokens(metadata_usage, usage)
+          metadata_usage['reasoningTokens'] || usage['reasoningTokens'] ||
+            usage.dig('output_tokens_details', 'thinking_tokens')
+        end
+
+        def stream_error_event?(event)
+          event.keys.any? { |key| key.end_with?('Exception') } || event['type'] == 'error'
+        end
+
+        def raise_stream_error(event)
+          if event['type'] == 'error'
+            message = event.dig('error', 'message') || 'Bedrock streaming error'
+            response = Struct.new(:body, :status).new({ 'message' => message }, 500)
+            ErrorMiddleware.parse_error(provider: self, response: response)
+            return
+          end
+
+          key = event.keys.find { |candidate| candidate.end_with?('Exception') }
+          payload = event[key]
+          message = payload['message'] || key
+          status = case key
+                   when 'throttlingException' then 429
+                   when 'validationException' then 400
+                   when 'accessDeniedException', 'unrecognizedClientException' then 401
+                   when 'serviceUnavailableException' then 503
+                   else 500
+                   end
+
+          response = Struct.new(:body, :status).new({ 'message' => message }, status)
+          ErrorMiddleware.parse_error(provider: self, response: response)
+        end
+
+        def extract_content_delta(event)
+          delta = normalized_delta(event)
+          return delta['text'] if delta['text']
+
+          return event.dig('delta', 'text') if event.dig('delta', 'type') == 'text_delta'
+
+          nil
+        end
+
+        def extract_thinking_delta(event)
+          delta = normalized_delta(event)
+          reasoning_content = delta['reasoningContent'] || {}
+
+          reasoning_text = reasoning_content['reasoningText'] || {}
+          return reasoning_text['text'] if reasoning_text['text']
+          return event.dig('delta', 'thinking') if event.dig('delta', 'type') == 'thinking_delta'
+
+          nil
+        end
+
+        def extract_thinking_signature(event)
+          signature = extract_signature_from_delta(event)
+          return signature if signature
+
+          signature = extract_signature_from_start(event)
+          return signature if signature
+
+          nil
+        end
+
+        def extract_signature_from_delta(event)
+          delta = normalized_delta(event)
+          reasoning_content = delta['reasoningContent'] || {}
+          reasoning_text = reasoning_content['reasoningText'] || {}
+          return reasoning_text['signature'] if reasoning_text['signature']
+          return event.dig('delta', 'signature') if event.dig('delta', 'type') == 'signature_delta'
+
+          nil
+        end
+
+        def extract_signature_from_start(event)
+          start = event.dig('contentBlockStart', 'start', 'reasoningContent')
+          return nil unless start
+
+          reasoning_text = start['reasoningText'] || {}
+          return reasoning_text['signature'] if reasoning_text['signature']
+          return start['redactedContent'] if start['redactedContent']
+
+          nil
+        end
+
+        def extract_tool_calls(event)
+          return extract_tool_call_start(event) if tool_call_start_event?(event)
+          return extract_tool_call_delta(event) if tool_call_delta_event?(event)
+
+          nil
+        end
+
+        def tool_call_start_event?(event)
+          event['contentBlockStart'] || event['start'] || event.dig('content_block', 'tool_use')
+        end
+
+        def tool_call_delta_event?(event)
+          event['contentBlockDelta'] || event.dig('delta', 'toolUse') || event.dig('delta', 'tool_use') ||
+            event.dig('delta', 'partial_json')
+        end
+
+        def extract_tool_call_start(event)
+          tool_use = event.dig('contentBlockStart', 'start', 'toolUse')
+          tool_use ||= event.dig('start', 'toolUse')
+          tool_use ||= event.dig('content_block', 'tool_use') if event['type'] == 'content_block_start'
+          return nil unless tool_use
+
+          tool_use_id = tool_use['toolUseId'] || tool_use['id']
+          tool_name = tool_use['name']
+          tool_input = tool_use['input'] || {}
+
+          {
+            tool_use_id => ToolCall.new(
+              id: tool_use_id,
+              name: tool_name,
+              arguments: tool_input
+            )
+          }
+        end
+
+        def extract_tool_call_delta(event)
+          input = normalized_delta(event).dig('toolUse', 'input')
+          input ||= normalized_delta(event).dig('tool_use', 'input')
+          input ||= event.dig('delta', 'partial_json') if event.dig('delta', 'type') == 'input_json_delta'
+          return nil unless input
+
+          { nil => ToolCall.new(id: nil, name: nil, arguments: input) }
+        end
+
+        def normalized_delta(event)
+          delta = event.dig('contentBlockDelta', 'delta') || event['delta'] || {}
+          return delta if delta.is_a?(Hash)
+
+          if delta.is_a?(String) && !delta.empty?
+            JSON.parse(delta)
+          else
+            {}
+          end
+        rescue JSON::ParserError
+          {}
+        end
       end
     end
   end

data/lib/ruby_llm/providers/bedrock.rb

@@ -1,82 +1,99 @@
 # frozen_string_literal: true
 
-require 'openssl'
-require 'time'
-
 module RubyLLM
   module Providers
-    # AWS Bedrock API integration.
+    # AWS Bedrock Converse API integration.
     class Bedrock < Provider
+      include Bedrock::Auth
       include Bedrock::Chat
-      include Bedrock::Streaming
-      include Bedrock::Models
-      include Bedrock::Signing
       include Bedrock::Media
-      include Anthropic::Tools
+      include Bedrock::Models
+      include Bedrock::Streaming
 
       def api_base
-        "https://bedrock-runtime.#{@config.bedrock_region}.amazonaws.com"
+        "https://bedrock-runtime.#{bedrock_region}.amazonaws.com"
       end
 
-      def parse_error(response)
-        return if response.body.empty?
-
-        body = try_parse_json(response.body)
-        case body
-        when Hash
-          body['message']
-        when Array
-          body.map do |part|
-            part['message']
-          end.join('. ')
-        else
-          body
-        end
+      def headers
+        {}
       end
 
-      def sign_request(url, method: :post, payload: nil)
-        signer = create_signer
-        request = build_request(url, method:, payload:)
-        signer.sign_request(request)
-      end
+      # rubocop:disable Metrics/ParameterLists
+      def complete(messages, tools:, temperature:, model:, params: {}, headers: {}, schema: nil, thinking: nil,
+                   tool_prefs: nil, &)
+        normalized_params = normalize_params(params, model:)
 
-      def create_signer
-        Signing::Signer.new({
-                              access_key_id: @config.bedrock_api_key,
-                              secret_access_key: @config.bedrock_secret_key,
-                              session_token: @config.bedrock_session_token,
-                              region: @config.bedrock_region,
-                              service: 'bedrock'
-                            })
+        super(
+          messages,
+          tools: tools,
+          tool_prefs: tool_prefs,
+          temperature: temperature,
+          model: model,
+          params: normalized_params,
+          headers: headers,
+          schema: schema,
+          thinking: thinking,
+          &
+        )
       end
+      # rubocop:enable Metrics/ParameterLists
 
-      def build_request(url, method: :post, payload: nil)
-        {
-          connection: @connection,
-          http_method: method,
-          url: url || completion_url,
-          body: payload ? JSON.generate(payload, ascii_only: false) : nil
-        }
-      end
+      def parse_error(response)
+        return if response.body.nil? || response.body.empty?
 
-      def build_headers(signature_headers, streaming: false)
-        accept_header = streaming ? 'application/vnd.amazon.eventstream' : 'application/json'
+        body = try_parse_json(response.body)
+        return body if body.is_a?(String)
 
-        signature_headers.merge(
-          'Content-Type' => 'application/json',
-          'Accept' => accept_header
-        )
+        body['message'] || body['Message'] || body['error'] || body['__type'] || super
+      end
+
+      def list_models
+        response = signed_get(models_api_base, models_url)
+        parse_list_models_response(response, slug, capabilities)
       end
 
       class << self
-        def capabilities
-          Bedrock::Capabilities
+        def configuration_options
+          %i[bedrock_api_key bedrock_secret_key bedrock_region bedrock_session_token]
         end
 
         def configuration_requirements
           %i[bedrock_api_key bedrock_secret_key bedrock_region]
         end
       end
+
+      private
+
+      def bedrock_region
+        @config.bedrock_region
+      end
+
+      def sync_response(connection, payload, additional_headers = {})
+        signed_post(connection, completion_url, payload, additional_headers)
+      end
+
+      def normalize_params(params, model:)
+        normalized = RubyLLM::Utils.deep_symbolize_keys(params || {})
+        additional_fields = normalized[:additionalModelRequestFields] || {}
+
+        top_k = normalized.delete(:top_k)
+        if !top_k.nil? && model_supports_top_k?(model)
+          additional_fields = RubyLLM::Utils.deep_merge(additional_fields, { top_k: top_k })
+        end
+
+        normalized[:additionalModelRequestFields] = additional_fields unless additional_fields.empty?
+        normalized
+      end
+
+      def model_supports_top_k?(model)
+        Bedrock::Models.reasoning_embedded?(model)
+      end
+
+      def api_payload(payload)
+        cleaned = RubyLLM::Utils.deep_symbolize_keys(RubyLLM::Utils.deep_dup(payload))
+        cleaned.delete(:tools)
+        cleaned
+      end
     end
   end
 end