ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/protocols/responses/token_counting.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Responses
+      # The Responses input-token counting endpoint.
+      module TokenCounting
+        COUNT_TOKENS_KEYS = %i[model input instructions tools tool_choice parallel_tool_calls reasoning text].freeze
+
+        module_function
+
+        def count_tokens_url
+          "#{completion_url}/input_tokens"
+        end
+
+        def render_count_tokens_payload(messages, model:, **options)
+          render_payload(messages, model: model, temperature: nil, **options).slice(*COUNT_TOKENS_KEYS)
+        end
+
+        def parse_count_tokens_response(response)
+          response.body.fetch('input_tokens')
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/responses/tools.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Responses
+      # Tools methods of the OpenAI Responses API. Function definitions are
+      # flat rather than nested under a `function` key.
+      module Tools
+        module_function
+
+        def tool_for(tool)
+          definition = {
+            type: 'function',
+            name: tool.name,
+            description: tool.description,
+            parameters: ChatCompletions::Tools.parameters_schema_for(tool),
+            strict: false
+          }
+
+          return definition if tool.provider_options.empty?
+
+          RubyLLM::Support::Utils.deep_merge(definition, tool.provider_options)
+        end
+
+        def build_tool_choice(tool_choice)
+          case tool_choice
+          when :auto, :none, :required
+            tool_choice
+          else
+            {
+              type: 'function',
+              name: tool_choice
+            }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/responses.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    # The OpenAI Responses API. Overrides the chat surface of Chat Completions;
+    # embeddings, images, moderation, and transcription are inherited. Runs
+    # stateless (store: false) and replays encrypted reasoning so multi-turn
+    # tool calls work without server-side state.
+    class Responses < ChatCompletions
+      include Responses::Approvals
+      include Responses::Chat
+      include Responses::Media
+      include Responses::Streaming
+      include Responses::Tools
+
+      SERVER_TOOL_ALIASES = {
+        web_search: { tool: { type: 'web_search' } },
+        file_search: { tool: { type: 'file_search' } },
+        code_execution: { tool: { type: 'code_interpreter', container: { type: 'auto' } } },
+        code_interpreter: { tool: { type: 'code_interpreter', container: { type: 'auto' } } },
+        image_generation: { tool: { type: 'image_generation' } },
+        mcp: lambda do |options|
+          options = Support::Utils.deep_symbolize_keys(options)
+          options[:server_url] = options.delete(:url) if options.key?(:url)
+          options[:server_label] = options.delete(:name) if options.key?(:name)
+          { tool: { type: 'mcp' }.merge(options) }
+        end
+      }.freeze
+
+      def server_tool_aliases
+        SERVER_TOOL_ALIASES
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/vertexai/batch_prediction.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+require 'stringio'
+
+module RubyLLM
+  module Protocols
+    module VertexAI
+      # Shared Vertex AI batchPredictionJobs plumbing. The input row and output
+      # result shapes belong to each Vertex protocol.
+      module BatchPrediction
+        include RubyLLM::Batch::Helpers
+
+        TERMINAL = %w[
+          JOB_STATE_SUCCEEDED
+          JOB_STATE_FAILED
+          JOB_STATE_CANCELLED
+          JOB_STATE_EXPIRED
+          JOB_STATE_PARTIALLY_SUCCEEDED
+        ].freeze
+        private_constant :TERMINAL
+
+        def create_batch(requests)
+          model = single_batch_model!(requests, 'vertexai')
+          validate_batch_requests!(requests)
+          input_uri, output_uri = vertex_batch_storage_uris
+          @provider.upload_file(
+            StringIO.new(vertex_batch_jsonl(requests)),
+            filename: 'input.jsonl',
+            uri: input_uri,
+            content_type: 'application/jsonl'
+          )
+
+          response = @connection.post("#{@provider.location_path}/batchPredictionJobs",
+                                      vertex_batch_job(model, input_uri, output_uri), idempotent: false)
+
+          parse_batch_response(response.body)
+        end
+
+        def find_batch(id)
+          parse_batch_response @connection.get(vertex_batch_name(id)).body
+        end
+
+        def cancel_batch(id)
+          @connection.post("#{vertex_batch_name(id)}:cancel", {})
+          find_batch(id)
+        end
+
+        def batch_results(id)
+          job = @connection.get(vertex_batch_name(id)).body
+          output_uri = vertex_output_uri(job)
+          unless output_uri
+            status = parse_batch_status(job['state'], completed: TERMINAL.include?(job['state']))
+            return [] if %i[failed cancelled].include?(status)
+
+            raise Error, 'vertexai batch has no GCS output URI yet'
+          end
+
+          rows = @provider.list_file_uris(output_uri).grep(/\.jsonl\z/).flat_map do |uri|
+            @provider.download_file(uri).to_s.each_line.filter_map do |line|
+              next if line.strip.empty?
+
+              JSON.parse(line)
+            end
+          end
+          parse_vertex_batch_results(rows, job:)
+        end
+
+        private
+
+        def vertex_batch_job(model, input_uri, output_uri)
+          {
+            displayName: "ruby_llm_#{SecureRandom.hex(8)}",
+            model: vertex_batch_model_path(model),
+            inputConfig: {
+              instancesFormat: 'jsonl',
+              gcsSource: { uris: [input_uri] }
+            },
+            outputConfig: {
+              predictionsFormat: 'jsonl',
+              gcsDestination: { outputUriPrefix: output_uri }
+            }
+          }
+        end
+
+        def vertex_batch_jsonl(requests)
+          requests.map { |request| JSON.generate(vertex_batch_request(request)) }.join("\n")
+        end
+
+        def vertex_batch_request(_request)
+          raise NotImplementedError
+        end
+
+        def validate_batch_requests!(_requests); end
+
+        def vertex_batch_model_path(model)
+          @provider.model_path(model)
+        end
+
+        def vertex_batch_storage_uris
+          base = @config.vertexai_batch_gcs_uri.to_s.sub(%r{/+\z}, '')
+          if base.empty?
+            raise ConfigurationError, 'Set vertexai_batch_gcs_uri to a gs:// bucket prefix for Vertex AI batches'
+          end
+
+          prefix = "#{base}/ruby_llm_batches/#{SecureRandom.hex(8)}"
+          ["#{prefix}/input.jsonl", "#{prefix}/output"]
+        end
+
+        def vertex_batch_name(id)
+          id.to_s.start_with?('projects/') ? id : "#{@provider.location_path}/batchPredictionJobs/#{id}"
+        end
+
+        def parse_batch_response(data)
+          state = data['state']
+          request_counts = data['completionStats']
+
+          {
+            id: data['name'],
+            raw_status: state,
+            completed: TERMINAL.include?(state),
+            request_counts:,
+            request_count: data.dig('labels', 'ruby_llm_request_count')&.to_i || request_counts&.values&.sum(&:to_i),
+            model: data['model']
+          }
+        end
+
+        def parse_batch_status(raw_status, completed:)
+          return :pending unless completed
+          return :succeeded if %w[JOB_STATE_SUCCEEDED JOB_STATE_PARTIALLY_SUCCEEDED].include?(raw_status)
+          return :cancelled if raw_status == 'JOB_STATE_CANCELLED'
+
+          :failed
+        end
+
+        def vertex_batch_result_index(line, fallback_index)
+          key = line['custom_id'] || line.dig('request', 'labels', 'ruby_llm_batch_id')
+          key ? batch_result_index(key) : fallback_index
+        end
+
+        def parse_vertex_batch_results(rows, **)
+          rows.each_with_index.filter_map { |line, index| parse_vertex_batch_result(line, index) }
+        end
+
+        def parse_vertex_batch_result(_line, _fallback_index)
+          raise NotImplementedError
+        end
+
+        def vertex_output_uri(job)
+          job.dig('outputInfo', 'gcsOutputDirectory') ||
+            job.dig('outputConfig', 'gcsDestination', 'outputUriPrefix')
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/vertexai/embedding_prediction/requests.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module VertexAI
+      class EmbeddingPrediction
+        # Input identity and model-specific batch embedding request formats.
+        module Requests
+          private
+
+          def validate_embedding_request!(request)
+            text = request.fetch(:text)
+            values = text.is_a?(Array) ? text : [text]
+            unless values.any? && values.all? { |value| value.is_a?(String) && !value.empty? }
+              raise ArgumentError, 'Vertex AI embedding batches require nonempty text strings'
+            end
+
+            payload = request.fetch(:payload)
+            validate_embedding_payload_keys!(payload)
+            inputs = embedding_batch_inputs(payload)
+            validate_gemini_embedding_options!(payload, inputs) if GEMINI_MODELS.include?(request.fetch(:model))
+            return if inputs.size == values.size
+
+            raise ArgumentError, 'Vertex AI embedding batch payload does not match the number of texts'
+          end
+
+          def validate_embedding_payload_keys!(payload)
+            supported = if payload.key?(:instances)
+                          %i[instances parameters]
+                        elsif payload.key?(:requests)
+                          [:requests]
+                        else
+                          %i[content outputDimensionality taskType title model]
+                        end
+            unknown = payload.keys - supported
+            return if unknown.empty?
+
+            raise ArgumentError, "Unsupported Vertex AI embedding batch options: #{unknown.join(', ')}"
+          end
+
+          def validate_gemini_embedding_options!(payload, inputs)
+            parameters = payload.fetch(:parameters, {}).keys - [:outputDimensionality]
+            supported = %i[content task_type taskType title outputDimensionality model]
+            unsupported = inputs.flat_map(&:keys).uniq - supported
+            return if parameters.empty? && unsupported.empty?
+
+            raise ArgumentError, 'Vertex AI Gemini embedding batches do not support these options: ' \
+                                 "#{(parameters + unsupported).join(', ')}"
+          end
+
+          def embedding_batch_inputs(payload)
+            return payload.fetch(:instances) if payload.key?(:instances)
+            return payload.fetch(:requests) if payload.key?(:requests)
+            return [payload] if payload.key?(:content)
+
+            raise ArgumentError, 'Vertex AI embedding batches require embedding request payloads'
+          end
+
+          def render_embedding_batch_rows(request)
+            inputs = embedding_batch_inputs(request.fetch(:payload))
+            array = request.fetch(:text).is_a?(Array)
+            inputs.each_with_index.map do |input, index|
+              key = "rllm-#{request.fetch(:custom_id)}-#{index}-#{inputs.size}-#{array ? 'a' : 's'}"
+              if LEGACY_MODELS.include?(request.fetch(:model))
+                input.merge(key:)
+              else
+                { key:, request: render_gemini_embedding_request(input, request.fetch(:payload)) }
+              end
+            end
+          end
+
+          def render_gemini_embedding_request(input, payload)
+            content = input[:content]
+            content = { parts: [{ text: content }] } if content.is_a?(String)
+            config = {
+              output_dimensionality: input[:outputDimensionality] || payload.dig(:parameters, :outputDimensionality),
+              task_type: input[:task_type] || input[:taskType], title: input[:title]
+            }.compact
+            { content:, embed_content_config: config }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/vertexai/embedding_prediction/results.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module VertexAI
+      class EmbeddingPrediction
+        # Regroups provider output rows into the submitted embedding requests.
+        module Results
+          private
+
+          def embedding_row_metadata(row)
+            key = row['key'] || row.dig('instance', 'key')
+            match = /\Arllm-(\d+)-(\d+)-(\d+)-([as])\z/.match(key.to_s)
+            raise Error, "Unknown Vertex AI embedding record key: #{key.inspect}" unless match
+
+            [Integer(match[1]), Integer(match[2]), Integer(match[3]), match[4] == 'a']
+          end
+
+          def parse_embedding_batch_group(index, rows, model:)
+            metadata = rows.map { |row| embedding_row_metadata(row) }
+            error = embedding_group_error(rows, metadata)
+            return [index, nil, batch_failure(index, error)] if error
+            return if rows.size < metadata.first[2]
+
+            ordered = rows.sort_by { |row| embedding_row_metadata(row)[1] }
+            parse_embedding_batch_result(index, ordered, model:, array: metadata.first[3])
+          end
+
+          def parse_embedding_batch_result(index, rows, model:, array:)
+            vectors = embedding_batch_vectors(rows)
+            return [index, nil, batch_failure(index, 'Vertex AI returned no valid embedding')] unless vectors
+
+            counts = rows.map { |row| embedding_batch_tokens(row) }
+            result = Embedding.new(vectors: array ? vectors : vectors.first,
+                                   model:, input_tokens: (counts.sum if counts.all?))
+            [index, result]
+          end
+
+          def embedding_group_error(rows, metadata)
+            unless valid_embedding_metadata?(metadata)
+              return 'Invalid or duplicate Vertex AI embedding record positions'
+            end
+
+            rows.filter_map { |row| batch_error_value(row['error']) || batch_error_value(row['status']) }
+                .find { |error| !error.empty? }
+          end
+
+          def valid_embedding_metadata?(metadata)
+            count = metadata.first[2]
+            positions = metadata.map { |entry| entry[1] }
+            return false unless metadata.first[3] || count == 1
+
+            metadata.map { |entry| entry[2..] }.uniq.one? && count.positive? &&
+              positions.uniq.size == positions.size && positions.max < count
+          end
+
+          def embedding_batch_vectors(rows)
+            vectors = rows.map do |row|
+              row.dig('response', 'embedding', 'values') || row.dig('predictions', 0, 'embeddings', 'values')
+            end
+            vectors if vectors.all? { |vector| vector.is_a?(Array) && vector.any? && vector.all?(Numeric) }
+          end
+
+          def embedding_batch_tokens(row)
+            value = row.dig('response', 'tokenCount') ||
+                    row.dig('response', 'usageMetadata', 'promptTokenCount') ||
+                    row.dig('predictions', 0, 'embeddings', 'statistics', 'token_count')
+            Integer(value) unless value.nil?
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/vertexai/embedding_prediction.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module VertexAI
+      # Text embeddings over Vertex AI's Cloud Storage batch prediction API.
+      class EmbeddingPrediction < Protocol
+        include BatchPrediction
+        include Requests
+        include Results
+
+        LEGACY_MODELS = %w[text-embedding-004 text-embedding-005 text-multilingual-embedding-002].freeze
+        GEMINI_MODELS = %w[gemini-embedding-001 gemini-embedding-2].freeze
+        MODELS = (LEGACY_MODELS + GEMINI_MODELS).freeze
+
+        private
+
+        def validate_batch_requests!(requests)
+          @requests = requests
+          model = requests.first.fetch(:model)
+          unless MODELS.include?(model)
+            raise Error,
+                  "Vertex AI embedding batches are not supported for #{model.inspect}"
+          end
+
+          requests.each { |request| validate_embedding_request!(request) }
+          return unless LEGACY_MODELS.include?(model)
+          return if requests.map { |request| request.fetch(:payload).fetch(:parameters, {}) }.uniq.one?
+
+          raise ArgumentError,
+                'Vertex AI legacy embedding batches require the same dimensions and parameters in every request'
+        end
+
+        def vertex_batch_job(model, input_uri, output_uri)
+          job = super.merge(labels: { ruby_llm_request_count: @requests.size.to_s })
+          return job unless LEGACY_MODELS.include?(model)
+
+          job.merge(instanceConfig: { instanceType: 'object', keyField: 'key' },
+                    modelParameters: @requests.first.fetch(:payload).fetch(:parameters, {}))
+        end
+
+        def vertex_batch_jsonl(requests)
+          rows = requests.flat_map { |request| render_embedding_batch_rows(request) }
+          rows.map { |row| JSON.generate(row) }.join("\n")
+        end
+
+        def parse_vertex_batch_results(rows, job:)
+          model = job.fetch('model').split('/').last
+          rows.group_by { |row| embedding_row_metadata(row).first }.filter_map do |index, group|
+            parse_embedding_batch_group(index, group, model:)
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module VertexAI
+      # Google Cloud Storage-backed files for Vertex AI batch input and output.
+      class Files < Protocols::Files
+        # GCS object names allow spaces and other characters URI() rejects.
+        GCS_URI = %r{\Ags://([^/]+)/?(.*)\z}m
+
+        # 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_name, key = parse_gcs_uri(target_uri)
+
+          with_file_body(attachment) do |body|
+            bucket(bucket_name).create_file(body, key, 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_name, key = parse_gcs_uri(file_id)
+          object = bucket(bucket_name).file(key)
+          raise Error, "GCS object not found: #{file_id}" unless object
+
+          uploaded_file(
+            { 'uri' => file_id },
+            id: file_id,
+            uri: file_id,
+            filename: File.basename(key),
+            byte_size: object.size,
+            created_at: object.created_at,
+            mime_type: object.content_type
+          )
+        end
+
+        def download(file_id)
+          bucket_name, key = parse_gcs_uri(file_id)
+          object = bucket(bucket_name).file(key)
+          raise Error, "GCS object not found: #{file_id}" unless object
+
+          object.download.string
+        end
+
+        def list_uris(prefix_uri)
+          bucket_name, prefix = parse_gcs_uri(prefix_uri)
+          uris = []
+          bucket(bucket_name).files(prefix: prefix).all do |object|
+            uris << "gs://#{bucket_name}/#{object.name}"
+          end
+          uris
+        end
+
+        private
+
+        def storage_uri_for(attachment)
+          base = @config.vertexai_batch_gcs_uri.to_s.sub(%r{/+\z}, '')
+          raise ConfigurationError, 'Set vertexai_batch_gcs_uri to a gs:// bucket prefix' if base.empty?
+
+          "#{base}/ruby_llm_uploads/#{SecureRandom.hex(8)}/#{attachment.filename}"
+        end
+
+        def storage
+          require 'google/cloud/storage'
+
+          options = { project_id: @config.vertexai_project_id }
+          if @config.vertexai_service_account_key
+            options[:credentials] =
+              JSON.parse(@config.vertexai_service_account_key)
+          end
+          ::Google::Cloud::Storage.new(**options)
+        rescue LoadError
+          raise Error, 'The google-cloud-storage gem is required for Vertex AI file uploads. ' \
+                       'Please add it to your Gemfile: gem "google-cloud-storage"'
+        end
+
+        def bucket(name)
+          storage.bucket(name) || raise(Error, "GCS bucket not found: #{name}")
+        end
+
+        def parse_gcs_uri(uri)
+          match = GCS_URI.match(uri.to_s)
+          raise ArgumentError, "Expected a gs:// URI, got: #{uri}" unless match
+
+          [match[1], match[2]]
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/vertexai/ranking.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    module VertexAI
+      # Vertex AI Search's Discovery Engine ranking API.
+      class Ranking < Protocol
+        def initialize(provider, model = nil)
+          super
+          @connection = provider.ranking_connection
+        end
+
+        def rerank_url
+          "#{@provider.ranking_config}:rank"
+        end
+
+        def render_rerank_payload(query, documents, model:, top_n: nil, provider_options: {})
+          validate_ranking_input(query, documents, top_n)
+          records = documents.each_with_index.map { |document, index| { id: index.to_s, content: document } }
+          { model: model, query: query, records: records, topN: top_n }.compact.merge(provider_options)
+        end
+
+        def parse_rerank_response(response, model:, documents: [])
+          records = response.body['records']
+          raise Error.new('Vertex AI Search returned no ranking records', response:) unless records.is_a?(Array)
+
+          seen = []
+          results = records.map do |record|
+            index = ranking_index(record, documents)
+            raise Error.new('Vertex AI Search returned a duplicate document id', response:) if seen.include?(index)
+
+            seen << index
+            Rerank::Result.new(index: index, document: documents[index], score: record.fetch('score'))
+          end
+          Rerank.new(results: results, model: model, raw: response.body)
+        end
+
+        private
+
+        def validate_ranking_input(query, documents, top_n)
+          unless query.is_a?(String) && !query.empty?
+            raise ArgumentError, 'Vertex AI Search reranking requires a nonempty query'
+          end
+          unless valid_ranking_documents?(documents)
+            raise ArgumentError, 'Vertex AI Search reranking accepts between 1 and 1000 nonempty text documents'
+          end
+          return if top_n.nil? || (top_n.is_a?(Integer) && top_n.positive?)
+
+          raise ArgumentError, 'top_n must be a positive integer'
+        end
+
+        def valid_ranking_documents?(documents)
+          documents.is_a?(Array) && (1..1000).cover?(documents.length) &&
+            documents.all? { |document| document.is_a?(String) && !document.empty? }
+        end
+
+        def ranking_index(record, documents)
+          id = record['id'] if record.is_a?(Hash)
+          unless id.is_a?(String) && id.match?(/\A(?:0|[1-9]\d*)\z/) && id.to_i < documents.length &&
+                 record['score'].is_a?(Numeric)
+            raise Error, 'Vertex AI Search returned an invalid document id or score'
+          end
+
+          id.to_i
+        end
+      end
+    end
+  end
+end