ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/protocols/bedrock/async_videos.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module Bedrock
+      # Luma Ray 2 video generation through Bedrock asynchronous invocation.
+      class AsyncVideos < Protocol
+        IMAGE_TYPES = %w[image/png image/jpeg].freeze
+
+        def video_url
+          '/async-invoke'
+        end
+
+        def render_video_payload(prompt, model:, with: [], provider_options: {})
+          output_uri = video_output_uri
+          unless prompt.is_a?(String) && (1..5000).cover?(prompt.length)
+            raise ArgumentError, 'Luma Ray 2 requires a prompt between 1 and 5000 characters'
+          end
+
+          options = Support::Utils.deep_symbolize_keys(provider_options)
+          outer = options.slice(:clientRequestToken, :tags)
+          input = { prompt: prompt }.merge(render_keyframes(with)).merge(options.except(*outer.keys))
+          { modelId: model, modelInput: input, clientRequestToken: SecureRandom.uuid,
+            outputDataConfig: { s3OutputDataConfig: { s3Uri: output_uri } } }.merge(outer)
+        end
+
+        def post_video(url, payload)
+          @connection.post(url, payload, idempotent: false) do |request|
+            request.headers.merge!(@provider.sign_headers('POST', url, JSON.generate(payload)))
+          end
+        end
+
+        def parse_video_job(response, model:)
+          id = response.body['invocationArn']
+          raise Error.new('Bedrock did not return a video invocation ARN', response:) unless id
+
+          VideoJob.new(id: id, protocol: self, model: model, raw: response.body)
+        end
+
+        def video_job_url(job)
+          "/async-invoke/#{URI.encode_www_form_component(job.id)}"
+        end
+
+        def refresh_video_job(job)
+          url = video_job_url(job)
+          response = @connection.get(url) do |request|
+            request.headers.merge!(@provider.sign_headers('GET', url, ''))
+          end
+          parse_video_job_status(response, job:)
+        end
+
+        def parse_video_job_status(response, **)
+          body = response.body
+          status = case body['status']
+                   when 'InProgress' then :pending
+                   when 'Completed' then :completed
+                   when 'Failed' then :failed
+                   else raise Error.new("Unknown Bedrock video status: #{body['status'].inspect}", response:)
+                   end
+          { status: status, raw: body, error: body['failureMessage'] }
+        end
+
+        def download_video(job)
+          prefix = job.raw.dig('outputDataConfig', 's3OutputDataConfig', 's3Uri')
+          raise Error, 'Bedrock video job has no output S3 URI' unless prefix
+
+          uris = @provider.list_file_uris("#{prefix.sub(%r{/+\z}, '')}/")
+          videos = uris.select { |uri| uri.downcase.end_with?('.mp4') }
+          raise Error, 'Expected exactly one MP4 in the Bedrock video output' unless videos.one?
+
+          Video.new(data: @provider.download_file(videos.first), mime_type: 'video/mp4', model: job.model, raw: job.raw)
+        end
+
+        private
+
+        def video_output_uri
+          uri = @config.bedrock_video_s3_uri
+          unless uri&.match?(%r{\As3://[a-z0-9][.\-a-z0-9]{1,61}[a-z0-9]/[^?#]+\z})
+            raise ConfigurationError, 'Set bedrock_video_s3_uri to the intended s3://bucket/output-prefix'
+          end
+
+          uri
+        end
+
+        def validate_animate_inputs!(with:)
+          raise ArgumentError, 'Luma Ray 2 accepts at most two reference images' if with.length > 2
+
+          with.each do |attachment|
+            if attachment.provider_file?
+              raise ArgumentError, 'Luma Ray 2 requires image bytes or URLs, not uploaded file ids'
+            end
+
+            raise UnsupportedAttachmentError, attachment.mime_type unless IMAGE_TYPES.include?(attachment.mime_type)
+          end
+        end
+
+        def render_keyframes(attachments)
+          return {} if attachments.empty?
+
+          frames = attachments.each_with_index.to_h do |attachment, index|
+            ["frame#{index}", { type: 'image', source: { type: 'base64', media_type: attachment.mime_type,
+                                                         data: attachment.encoded } }]
+          end
+          { keyframes: frames }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/bedrock/batches.rb

@@ -0,0 +1,129 @@
+# frozen_string_literal: true
+
+require 'digest'
+require 'stringio'
+require 'uri'
+
+module RubyLLM
+  module Protocols
+    module Bedrock
+      # Shared Bedrock Model Invocation Job and S3 result lifecycle.
+      module Batches
+        include RubyLLM::Batch::Helpers
+
+        TERMINAL = %w[Completed PartiallyCompleted Failed Stopped Expired].freeze
+        private_constant :TERMINAL
+
+        def create_batch(requests)
+          model = single_batch_model!(requests, 'bedrock')
+          validate_bedrock_batch_requests!(requests)
+          role_arn = bedrock_batch_role_arn
+          input_uri, output_uri = bedrock_batch_storage_uris
+          @provider.upload_file(
+            StringIO.new(bedrock_batch_jsonl(requests)),
+            filename: 'input.jsonl',
+            uri: input_uri,
+            content_type: 'application/jsonl'
+          )
+
+          response = @provider.signed_post(@provider.control_api_base, '/model-invocation-job', {
+                                             clientRequestToken: Digest::SHA256.hexdigest(input_uri),
+                                             jobName: bedrock_job_name(input_uri, requests:),
+                                             roleArn: role_arn,
+                                             modelId: model,
+                                             modelInvocationType: bedrock_invocation_type,
+                                             inputDataConfig: {
+                                               s3InputDataConfig: { s3Uri: input_uri }
+                                             },
+                                             outputDataConfig: {
+                                               s3OutputDataConfig: { s3Uri: output_uri }
+                                             }
+                                           })
+
+          find_batch(response.body['jobArn'])
+        end
+
+        def find_batch(id)
+          data = @provider.signed_get(@provider.control_api_base, bedrock_job_url(id)).body
+          protocol = @provider.embedding_batch_protocol(data['modelId']) if data['modelInvocationType'] == 'InvokeModel'
+          parser = protocol ? protocol.new(@provider) : self
+          parser.send(:parse_batch_response, data).tap do |result|
+            result[:batch_protocol] = protocol if protocol
+          end
+        end
+
+        def cancel_batch(id)
+          @provider.signed_post(@provider.control_api_base, "#{bedrock_job_url(id)}/stop", {})
+          find_batch(id)
+        end
+
+        def batch_results(id)
+          job = @provider.signed_get(@provider.control_api_base, bedrock_job_url(id)).body
+          output_uri = job.dig('outputDataConfig', 's3OutputDataConfig', 's3Uri')
+          unless output_uri
+            status = parse_batch_status(job['status'], completed: TERMINAL.include?(job['status']))
+            return [] if %i[failed cancelled].include?(status)
+
+            raise Error, 'bedrock batch has no S3 output URI yet'
+          end
+
+          outputs = @provider.list_file_uris(output_uri)
+                             .grep(/\.jsonl\.out\z/)
+                             .reject { |uri| uri.end_with?('/manifest.json.out') }
+                             .map { |uri| @provider.download_file(uri) }
+          parse_bedrock_outputs(outputs, model: job['modelId'])
+        end
+
+        private
+
+        def bedrock_batch_storage_uris
+          base = @config.bedrock_batch_s3_uri.to_s.sub(%r{/+\z}, '')
+          if base.empty?
+            raise ConfigurationError, 'Set bedrock_batch_s3_uri to an s3:// bucket prefix for Bedrock batches'
+          end
+
+          prefix = "#{base}/ruby_llm_batches/#{SecureRandom.hex(8)}"
+          ["#{prefix}/input.jsonl", "#{prefix}/output"]
+        end
+
+        def bedrock_batch_role_arn
+          @config.bedrock_batch_role_arn ||
+            raise(ConfigurationError, 'Set bedrock_batch_role_arn for Bedrock batches')
+        end
+
+        def bedrock_job_name(input_uri, **)
+          "ruby-llm-#{Digest::SHA256.hexdigest(input_uri)[0, 16]}"
+        end
+
+        def bedrock_job_url(id)
+          "/model-invocation-job/#{URI.encode_www_form_component(id)}"
+        end
+
+        def parse_batch_response(data)
+          status = data['status']
+          {
+            id: data['jobArn'] || data['jobIdentifier'],
+            raw_status: status,
+            completed: TERMINAL.include?(status),
+            request_counts: bedrock_request_counts(data)
+          }
+        end
+
+        def parse_batch_status(raw_status, completed:)
+          return :pending unless completed
+          return :succeeded if %w[Completed PartiallyCompleted].include?(raw_status)
+          return :cancelled if raw_status == 'Stopped'
+
+          :failed
+        end
+
+        def bedrock_request_counts(data)
+          {
+            'submitted' => data['submitTime'],
+            'completed' => data['endTime']
+          }.compact
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/bedrock/files.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module Bedrock
+      # S3-backed file storage for Bedrock batch input and output.
+      class Files < Protocols::Files
+        # rubocop:disable-next Lint/UnusedMethodArgument
+        def upload(file, filename: nil, purpose: nil, expires_in: nil, uri: nil, content_type: nil,
+                   provider_options: {})
+          attachment = file_attachment(file, filename:)
+          target_uri = uri || storage_uri_for(attachment)
+          bucket, key = parse_s3_uri(target_uri)
+
+          with_file_body(attachment) do |body|
+            s3_client.put_object(bucket: bucket, key: key, body: body,
+                                 content_type: content_type || file_content_type(attachment))
+          end
+
+          uploaded_file(
+            { 'uri' => target_uri },
+            id: target_uri,
+            uri: target_uri,
+            filename: attachment.filename,
+            byte_size: file_size(attachment),
+            mime_type: content_type || file_content_type(attachment)
+          )
+        end
+
+        def find(file_id)
+          bucket, key = parse_s3_uri(file_id)
+          response = s3_client.head_object(bucket: bucket, key: key)
+
+          uploaded_file(
+            { 'uri' => file_id },
+            id: file_id,
+            uri: file_id,
+            filename: File.basename(key),
+            byte_size: response.content_length,
+            created_at: response.last_modified,
+            mime_type: response.content_type
+          )
+        end
+
+        def download(file_id)
+          bucket, key = parse_s3_uri(file_id)
+          s3_client.get_object(bucket: bucket, key: key).body.read
+        end
+
+        def list_uris(prefix_uri)
+          bucket, prefix = parse_s3_uri(prefix_uri)
+          uris = []
+          token = nil
+
+          loop do
+            options = { bucket: bucket, prefix: prefix }
+            options[:continuation_token] = token if token
+            response = s3_client.list_objects_v2(**options)
+            uris.concat(Array(response.contents).map { |object| "s3://#{bucket}/#{object.key}" })
+            token = response.next_continuation_token
+            break unless token
+          end
+
+          uris
+        end
+
+        private
+
+        def storage_uri_for(attachment)
+          base = @config.bedrock_batch_s3_uri.to_s.sub(%r{/+\z}, '')
+          raise ConfigurationError, 'Set bedrock_batch_s3_uri to an s3:// bucket prefix' if base.empty?
+
+          "#{base}/ruby_llm_uploads/#{SecureRandom.hex(8)}/#{attachment.filename}"
+        end
+
+        def s3_client
+          require 'aws-sdk-s3'
+
+          ::Aws::S3::Client.new(**s3_client_options)
+        rescue LoadError
+          raise Error, 'The aws-sdk-s3 gem is required for Bedrock file uploads. ' \
+                       'Please add it to your Gemfile: gem "aws-sdk-s3"'
+        end
+
+        def s3_client_options
+          options = { region: @config.bedrock_region }
+
+          if @config.bedrock_credential_provider
+            options[:credentials] = @config.bedrock_credential_provider
+          else
+            options.merge!(
+              access_key_id: @config.bedrock_api_key,
+              secret_access_key: @config.bedrock_secret_key,
+              session_token: @config.bedrock_session_token
+            )
+          end
+
+          options
+        end
+
+        def parse_s3_uri(uri)
+          parsed = URI(uri)
+          raise ArgumentError, "Expected an s3:// URI, got: #{uri}" unless parsed.scheme == 's3'
+
+          [parsed.host, parsed.path.delete_prefix('/')]
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/bedrock/guardrails.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module Bedrock
+      # Moderation through an explicitly configured Bedrock guardrail.
+      class Guardrails < Protocol
+        POLICY_FILTERS = {
+          'contentPolicy' => %w[filters],
+          'topicPolicy' => %w[topics],
+          'wordPolicy' => %w[customWords managedWordLists],
+          'sensitiveInformationPolicy' => %w[piiEntities regexes],
+          'contextualGroundingPolicy' => %w[filters]
+        }.freeze
+        private_constant :POLICY_FILTERS
+
+        IMAGE_TYPES = %w[image/png image/jpeg].freeze
+        INTERVENTIONS = %w[BLOCKED ANONYMIZED].freeze
+        private_constant :IMAGE_TYPES, :INTERVENTIONS
+
+        def moderate(input, model:, with: [], provider_options: {})
+          raise ArgumentError, 'Bedrock guardrails do not accept a model' unless model.nil?
+
+          url = @provider.guardrail_url
+          attachments = Attachment.wrap(with, config: @config)
+          validate_moderation_input(input, attachments)
+          payloads = (input.is_a?(Array) ? input : [input]).map do |text|
+            render_moderation_payload(text, attachments, provider_options)
+          end
+          track_usage(:moderation) do
+            responses = payloads.map do |payload|
+              response = post_moderation(url, payload)
+              result = parse_moderation_result(response)
+              @usage_tracker.succeed_attempts(tokens: [Tokens.new])
+              [response, result]
+            end
+            parse_moderation_responses(responses)
+          end
+        end
+
+        private
+
+        def parse_moderation_responses(responses)
+          RubyLLM::Moderation.new(
+            id: responses.one? ? responses.first.first.headers['x-amzn-requestid'] : nil,
+            model: nil,
+            results: responses.map(&:last),
+            raw: responses.one? ? responses.first.first.body : responses.map { |response, _result| response.body }
+          )
+        end
+
+        def validate_moderation_input(input, attachments)
+          unless valid_text_input?(input)
+            raise ArgumentError, 'Bedrock moderation requires text or a nonempty array of texts'
+          end
+          if input.is_a?(Array) && attachments.any?
+            raise ArgumentError, 'Bedrock moderation accepts image attachments with one text input at a time'
+          end
+          raise ArgumentError, 'Bedrock moderation requires text or an image' if input.nil? && attachments.empty?
+
+          attachments.each do |attachment|
+            next if IMAGE_TYPES.include?(attachment.mime_type)
+
+            raise UnsupportedAttachmentError, attachment.mime_type
+          end
+        end
+
+        def valid_text_input?(input)
+          input.nil? || input.is_a?(String) || (input.is_a?(Array) && input.any? && input.all?(String))
+        end
+
+        def render_moderation_payload(input, attachments, provider_options)
+          content = input.nil? ? [] : [{ text: { text: input } }]
+          content.concat(attachments.map do |attachment|
+            { image: { format: attachment.mime_type.delete_prefix('image/'), source: { bytes: attachment.encoded } } }
+          end)
+          { source: 'INPUT', content: }.merge(render_guardrail_options(provider_options))
+        end
+
+        def render_guardrail_options(provider_options)
+          options = provider_options.transform_keys(&:to_sym)
+          unknown = options.keys - %i[source outputScope]
+          unless unknown.empty?
+            raise ArgumentError, "Bedrock moderation options only support source and outputScope: #{unknown.join(', ')}"
+          end
+
+          options
+        end
+
+        def post_moderation(url, payload)
+          @connection.post(url, payload, usage: @usage_tracker) do |request|
+            request.headers.merge!(@provider.sign_headers('POST', url, JSON.generate(payload)))
+          end
+        end
+
+        def parse_moderation_result(response)
+          data = response.body
+          unless data.is_a?(Hash) && %w[NONE GUARDRAIL_INTERVENED].include?(data['action'])
+            raise Error.new('Bedrock moderation returned no recognized verdict', response:)
+          end
+
+          categories = Array(data['assessments']).flat_map { |assessment| parse_moderation_categories(assessment) }
+          RubyLLM::Moderation::Result.new(
+            flagged: data['action'] == 'GUARDRAIL_INTERVENED', categories: categories.uniq, category_scores: {}
+          )
+        end
+
+        def parse_moderation_categories(assessment)
+          POLICY_FILTERS.flat_map do |policy, groups|
+            groups.flat_map do |group|
+              Array(assessment.dig(policy, group)).filter_map do |filter|
+                next unless INTERVENTIONS.include?(filter['action'])
+
+                filter['name'] || filter['type'] || "#{policy}.#{group}"
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/bedrock/rerank.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module Bedrock
+      # Bedrock's document reranking API.
+      class Rerank < Protocol
+        def rerank(query, documents, model:, top_n: nil, provider_options: {})
+          track_usage(:rerank) do
+            payload = render_rerank_payload(query, documents, model:, top_n:, provider_options:)
+            pages = rerank_pages(payload)
+            parse_rerank_pages(pages, model:, documents:)
+          end
+        end
+
+        private
+
+        def render_rerank_payload(query, documents, model:, top_n:, provider_options:)
+          validate_rerank_input(query, documents, top_n)
+          payload = {
+            queries: [{ type: 'TEXT', textQuery: { text: query } }],
+            sources: documents.map { |document| render_rerank_source(document) },
+            rerankingConfiguration: {
+              type: 'BEDROCK_RERANKING_MODEL',
+              bedrockRerankingConfiguration: {
+                modelConfiguration: { modelArn: @provider.rerank_model_arn(model) },
+                numberOfResults: top_n || documents.length
+              }
+            }
+          }
+          Support::Utils.deep_merge(payload, provider_options)
+        end
+
+        def validate_rerank_input(query, documents, top_n)
+          unless query.is_a?(String) && !query.empty?
+            raise ArgumentError, 'Bedrock reranking requires one nonempty text query'
+          end
+          unless documents.is_a?(Array) && (1..1000).cover?(documents.length)
+            raise ArgumentError, 'Bedrock reranking accepts between 1 and 1000 documents'
+          end
+          return if top_n.nil? || (top_n.is_a?(Integer) && (1..1000).cover?(top_n))
+
+          raise ArgumentError, 'Bedrock reranking top_n must be between 1 and 1000'
+        end
+
+        def render_rerank_source(document)
+          content = case document
+                    when String then { type: 'TEXT', textDocument: { text: document } }
+                    when Hash then { type: 'JSON', jsonDocument: document }
+                    else raise ArgumentError, 'Bedrock reranking documents must be text or JSON objects'
+                    end
+          { type: 'INLINE', inlineDocumentSource: content }
+        end
+
+        def rerank_pages(payload)
+          pages = []
+          seen = []
+          loop do
+            response = post_rerank(payload)
+            @usage_tracker.succeed_attempts(tokens: [Tokens.new])
+            pages << response.body
+            token = response.body['nextToken']
+            break if token.nil? || token.empty?
+            raise Error.new('Bedrock reranking returned a repeated pagination token', response:) if seen.include?(token)
+
+            seen << token
+            payload = payload.merge(nextToken: token)
+          end
+          pages
+        end
+
+        def post_rerank(payload)
+          @provider.agent_connection.post('/rerank', payload, usage: @usage_tracker) do |request|
+            request.headers.merge!(@provider.sign_headers('POST', '/rerank', JSON.generate(payload),
+                                                          base_url: @provider.agent_api_base))
+          end
+        end
+
+        def parse_rerank_pages(pages, model:, documents:)
+          results = pages.flat_map do |page|
+            Array(page['results']).map do |item|
+              index = item['index']
+              unless index.is_a?(Integer) && index.between?(0, documents.length - 1)
+                raise Error, 'Bedrock reranking returned an invalid document index'
+              end
+
+              RubyLLM::Rerank::Result.new(index:, document: documents[index], score: item['relevanceScore'])
+            end
+          end
+          RubyLLM::Rerank.new(results:, model:, raw: pages.length == 1 ? pages.first : pages)
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/chat_completions/batches.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class ChatCompletions
+      # OpenAI-compatible file-backed Batch API for Chat Completions.
+      module Batches
+        include Protocols::OpenAI::Batches
+
+        private
+
+        def batch_endpoint
+          '/v1/chat/completions'
+        end
+
+        def validate_batch_requests!(requests)
+          return if requests.all? { |request| chat_completion_payload?(request.fetch(:payload)) }
+
+          raise Error, "#{@provider.slug} batch requests require chat completion payloads"
+        end
+
+        def chat_completion_payload?(payload)
+          payload.key?(:messages) || payload.key?('messages')
+        end
+
+        def parse_batch_completion_response(body)
+          parse_completion_body(body, raw: body)
+        end
+      end
+    end
+  end
+end