ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/providers/vertexai/transcription.rb

@@ -3,12 +3,57 @@
 module RubyLLM
   module Providers
     class VertexAI
-      # Vertex AI specific helpers for audio transcription
-      module Transcription
+      class Transcription < VertexAI::Gemini # :nodoc: all
+        def render_transcription_options(timestamps:, **)
+          return {} if timestamps.nil?
+          raise ArgumentError, 'Vertex AI transcription timestamps must be word' unless timestamps == :word
+
+          { generationConfig: { audioTranscriptionConfig: { wordTimestamp: true } } }
+        end
+
+        include Protocols::Gemini::FileTranscription
+
+        def validate_transcription_request(...)
+          super
+          return if @config.vertexai_location == 'global'
+
+          raise ArgumentError, 'Vertex AI dedicated transcription requires vertexai_location = "global"'
+        end
+
+        def render_transcription_payload(attachment, language:, speaker_names:, provider_options:, prompt:, **)
+          config = { languageCodes: language && Array(language), customVocabulary: prompt && Array(prompt),
+                     diarization: speaker_names && true }.compact
+          payload = { contents: [{ role: 'user', parts: [format_audio_part(attachment)] }],
+                      generationConfig: { audioTranscriptionConfig: config } }
+          Support::Utils.deep_merge(payload, provider_options)
+        end
+
         private
 
-        def transcription_url(model)
-          "projects/#{@config.vertexai_project_id}/locations/#{@config.vertexai_location}/publishers/google/models/#{model}:generateContent" # rubocop:disable Layout/LineLength
+        def parse_transcription_response(response, model:)
+          data = response.body
+          parts = data.dig('candidates', 0, 'content', 'parts') || []
+          segments = parts.filter_map { |part| parse_transcription_segment(part) }
+          text = parts.filter_map { |part| part['text'] || part.dig('audioTranscription', 'text') }.join
+          words = segments.flat_map { |segment| segment['words'] }
+          RubyLLM::Transcription.new(text:, model:, segments: segments.empty? ? nil : segments,
+                                     words: words.empty? ? nil : words, **extract_usage(data))
+        end
+
+        def parse_transcription_segment(part)
+          transcription = part['audioTranscription']
+          return unless transcription
+
+          { 'text' => part['text'] || transcription['text'], 'speaker' => transcription['speakerLabel'],
+            'words' => Array(transcription['words']).map do |word|
+              parse_transcription_word(word, transcription)
+            end }.compact
+        end
+
+        def parse_transcription_word(word, transcription)
+          { 'word' => word['word'], 'speaker' => transcription['speakerLabel'],
+            'start' => word['startOffset'] && Float(word['startOffset'].delete_suffix('s')),
+            'end' => word['endOffset'] && Float(word['endOffset'].delete_suffix('s')) }.compact
         end
       end
     end

data/lib/ruby_llm/providers/vertexai/videos.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI
+      # Veo prediction jobs with inline or Cloud Storage video output.
+      module Videos
+        def video_url
+          "#{@provider.model_path(model_id(@model))}:predictLongRunning"
+        end
+
+        def video_job_url(job)
+          "#{job.id.split('/operations/').first}:fetchPredictOperation"
+        end
+
+        def refresh_video_job(job)
+          response = @connection.post video_job_url(job), { operationName: job.id }
+          parse_video_job_status(response, job:)
+        end
+
+        def download_video(job)
+          video = generated_video(job.raw)
+          raise Error, 'Vertex AI returned no video' unless video
+
+          data = if video['bytesBase64Encoded']
+                   Base64.decode64(video['bytesBase64Encoded'])
+                 else
+                   @provider.download_file(video.fetch('gcsUri'))
+                 end
+
+          Video.new(data:, mime_type: video['mimeType'] || 'video/mp4', model: job.model, raw: job.raw)
+        end
+
+        private
+
+        def render_video_extension(source)
+          uri = source.respond_to?(:uri) ? source.uri : source
+          return { gcsUri: uri, mimeType: 'video/mp4' } if uri.is_a?(String) && uri.start_with?('gs://')
+
+          video = video_extension_attachment(source)
+          { bytesBase64Encoded: video.encoded, mimeType: video.mime_type }
+        end
+
+        def render_video_image(image)
+          { bytesBase64Encoded: image.encoded, mimeType: image.mime_type }
+        end
+
+        def generated_video(body)
+          video = body.dig('response', 'videos', 0)
+          video if video && %w[bytesBase64Encoded gcsUri].any? { |key| !video[key].to_s.empty? }
+        end
+
+        def filtered_video_failure(body)
+          error = Array(body.dig('response', 'raiMediaFilteredReasons')).join(' ')
+          error = 'Vertex AI returned no video' if error.empty?
+          { status: :failed, raw: body, error: }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai.rb

@@ -5,47 +5,157 @@ require 'stringio'
 module RubyLLM
   module Providers
     # Google Vertex AI implementation
-    class VertexAI < Gemini
-      include VertexAI::Chat
-      include VertexAI::Streaming
-      include VertexAI::Embeddings
-      include VertexAI::Models
-      include VertexAI::Transcription
+    class VertexAI < Provider
+      protocol :gemini, VertexAI::Gemini, batches: VertexAI::Gemini::Batches
+      protocol :anthropic, VertexAI::Anthropic, batches: VertexAI::Anthropic::Batches
+      protocol :mistral, VertexAI::Mistral
+      protocol :chat_completions, VertexAI::ChatCompletions, batches: VertexAI::ChatCompletions::Batches
+      protocol :embed_content, VertexAI::EmbedContent
+      protocol :embedding_prediction, Protocols::VertexAI::EmbeddingPrediction
+      protocol :transcription, VertexAI::Transcription
+      protocol :live_transcription, VertexAI::LiveTranscription
+      protocol :ranking, Protocols::VertexAI::Ranking
+      protocol :research, Protocols::VertexAI::Research
+      protocol :files, Protocols::VertexAI::Files
 
       SCOPES = [
         'https://www.googleapis.com/auth/cloud-platform',
         'https://www.googleapis.com/auth/generative-language.retriever'
       ].freeze
 
+      class << self
+        def capabilities
+          VertexAI::Capabilities
+        end
+
+        def models_dev_alias(...)
+          VertexAI::Models.models_dev_alias(...)
+        end
+
+        # models.dev pins Vertex AI models to a version (claude-haiku-4-5@20251001);
+        # Vertex AI serves them by bare name.
+        def models_dev_model_id(id)
+          id&.split('@')&.first
+        end
+      end
+
+      # Vertex AI hosts models from several publishers, each speaking its
+      # native protocol. Publisher-prefixed ids are MaaS models served
+      # through the OpenAI-compatible endpoint.
+      def protocol_for(model, operation: nil, **)
+        return protocols[:ranking] if operation == :rerank
+
+        transcription = transcription_protocol_for(model.id) if operation == :transcribe
+        return transcription if transcription
+
+        if operation == :embed && %w[gemini-embedding-2 gemini-embedding-2-preview].include?(model.id)
+          return protocols[:embed_content]
+        end
+
+        case model.id
+        when %r{/} then protocols[:chat_completions]
+        when /\Aclaude/ then protocols[:anthropic]
+        when VertexAI::Mistral::MODELS then protocols[:mistral]
+        else super
+        end
+      end
+
+      def location_path
+        "projects/#{@config.vertexai_project_id}/locations/#{@config.vertexai_location}"
+      end
+
+      def model_path(model, publisher: 'google')
+        "#{location_path}/publishers/#{publisher}/models/#{model}"
+      end
+
       def initialize(config)
         super
         @authorizer = nil
       end
 
+      def batch_protocol
+        batch_protocol_for_name(:gemini)
+      end
+
+      def batch_protocol_for(requests)
+        kinds = requests.map { |request| request.key?(:text) }.uniq
+        raise Error, 'Vertex AI batches take chat or embedding requests, not both' unless kinds.size == 1
+        return protocols[:embedding_prediction] if kinds.first
+
+        models = requests.map { |request| request.fetch(:model) }.uniq
+        raise Error, 'vertexai batch requests must use one model per submission' unless models.one?
+
+        protocol_name = batch_protocol_name_for(models.first)
+        protocol = batch_protocol_for_name(protocol_name)
+        return protocol if protocol
+
+        raise Error, 'vertexai batch requests currently support Gemini, Anthropic, and MaaS chat models'
+      end
+      private :batch_protocol, :batch_protocol_for
+
+      def find_batch(id)
+        batch = super
+        protocol = batch_protocol_for_model_path(batch[:model])
+
+        protocol ? batch.merge(batch_protocol: protocol) : batch
+      end
+
+      def batch_cost_multiplier(model:, component:)
+        return if model.id.include?('/')
+        return 1 if !model.id.start_with?('claude') && %i[cache_read cache_write].include?(component)
+
+        0.5
+      end
+
       def api_base
+        api_base_for(@config.vertexai_location)
+      end
+
+      def api_base_for(location)
         return @config.vertexai_api_base if @config.vertexai_api_base
 
-        if @config.vertexai_location.to_s == 'global'
+        if location.to_s == 'global'
           'https://aiplatform.googleapis.com/v1beta1'
         else
-          "https://#{@config.vertexai_location}-aiplatform.googleapis.com/v1beta1"
+          "https://#{location}-aiplatform.googleapis.com/v1beta1"
         end
       end
 
-      def headers
-        if defined?(VCR) && !VCR.current_cassette.recording?
-          { 'Authorization' => 'Bearer test-token' }
-        else
-          initialize_authorizer unless @authorizer
-          @authorizer.apply({})
+      def ranking_config # :nodoc:
+        @config.vertexai_ranking_config ||
+          "projects/#{@config.vertexai_project_id}/locations/global/rankingConfigs/default_ranking_config"
+      end
+
+      def ranking_connection # :nodoc:
+        base = @config.vertexai_ranking_api_base || 'https://discoveryengine.googleapis.com/v1'
+        @ranking_connection ||= Transport::Connection.new(self, @config, api_base: base).tap do |connection|
+          connection.connection.headers['X-Goog-User-Project'] = @config.vertexai_project_id
         end
-      rescue Google::Auth::AuthorizationError => e
-        raise UnauthorizedError.new(nil, "Invalid Google Cloud credentials for Vertex AI: #{e.message}")
+      end
+
+      # The rescue can't name Google::Auth::AuthorizationError directly:
+      # when googleauth is missing, evaluating the constant would replace
+      # the helpful install error with a NameError.
+      def headers
+        initialize_authorizer unless @authorizer
+        @authorizer.apply({})
+      rescue StandardError => e
+        raise unless defined?(Google::Auth::AuthorizationError) && e.is_a?(Google::Auth::AuthorizationError)
+
+        raise UnauthorizedError, "Invalid Google Cloud credentials for Vertex AI: #{e.message}"
       end
 
       class << self
         def configuration_options
-          %i[vertexai_project_id vertexai_location vertexai_service_account_key vertexai_api_base]
+          %i[
+            vertexai_project_id
+            vertexai_location
+            vertexai_service_account_key
+            vertexai_api_base
+            vertexai_batch_gcs_uri
+            vertexai_ranking_api_base
+            vertexai_ranking_config
+          ]
         end
 
         def configuration_requirements
@@ -55,6 +165,13 @@ module RubyLLM
 
       private
 
+      def transcription_protocol_for(id)
+        case id
+        when 'gemini-3.5-transcribe-preview' then protocols[:transcription]
+        when 'gemini-3.5-transcribe-live-preview' then protocols[:live_transcription]
+        end
+      end
+
       def initialize_authorizer
         require 'googleauth'
         @authorizer =
@@ -70,6 +187,32 @@ module RubyLLM
         raise Error,
               'The googleauth gem ~> 1.15 is required for Vertex AI. Please add it to your Gemfile: gem "googleauth"'
       end
+
+      def batch_protocol_name_for(model)
+        case model
+        when %r{/} then :chat_completions
+        when /\Aclaude/ then :anthropic
+        when VertexAI::Mistral::MODELS then :mistral
+        else :gemini
+        end
+      end
+
+      def batch_protocol_for_model_path(model_path)
+        if Protocols::VertexAI::EmbeddingPrediction::MODELS.include?(model_path.to_s.split('/').last)
+          return protocols[:embedding_prediction]
+        end
+
+        case model_path.to_s
+        when %r{/publishers/google/models/}, %r{\Apublishers/google/models/}
+          batch_protocol_for_name(:gemini)
+        when %r{/publishers/anthropic/models/}, %r{\Apublishers/anthropic/models/}
+          batch_protocol_for_name(:anthropic)
+        when %r{/publishers/mistralai/models/}, %r{\Apublishers/mistralai/models/}
+          nil
+        when %r{/publishers/[^/]+/models/}, %r{\Apublishers/[^/]+/models/}
+          batch_protocol_for_name(:chat_completions)
+        end
+      end
     end
   end
 end

data/lib/ruby_llm/providers/xai/capabilities.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      # Feature capability gaps not represented in upstream model catalogs.
+      module Capabilities
+        def self.augment(capabilities, model_id:, modalities:)
+          return capabilities unless modalities[:output].include?('text')
+
+          additions = ['streaming']
+          additions.push('tool_choice', 'parallel_tool_calls') if model_id == 'grok-4.3'
+          capabilities | additions
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/xai/chat.rb

@@ -10,9 +10,10 @@ module RubyLLM
           role.to_s
         end
 
-        def format_content(content)
-          OpenAI::Media.format_content(
+        def format_content(content, attachments = [])
+          Protocols::ChatCompletions::Media.format_content(
             content,
+            attachments,
             document_attachments: :none,
             image_attachments: true,
             audio_attachments: false

data/lib/ruby_llm/providers/xai/chat_completions/batches.rb

@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      class ChatCompletions
+        # xAI native batch containers. Requests submit in Responses or Chat
+        # Completions shape; results always come back as chat completions.
+        module Batches
+          include RubyLLM::Batch::Helpers
+
+          def create_batch(requests)
+            batch = @connection.post('batches', { name: "ruby_llm_#{SecureRandom.hex(8)}" },
+                                     idempotent: false).body
+            id = batch['batch_id'] || batch['id']
+            @connection.post("batches/#{id}/requests", {
+                               batch_requests: requests.map { |request| xai_batch_request(request) }
+                             }, idempotent: false)
+
+            find_batch(id)
+          end
+
+          def find_batch(id)
+            parse_batch_response @connection.get("batches/#{id}").body
+          end
+
+          def cancel_batch(id)
+            parse_batch_response @connection.post("batches/#{id}:cancel", {}).body
+          end
+
+          def batch_results(id)
+            results = []
+            token = nil
+
+            loop do
+              response = @connection.get("batches/#{id}/results") do |request|
+                request.params[:limit] = 100
+                request.params[:pagination_token] = token if token
+              end.body
+
+              page_results = Array(response['results'] || response['batch_results'])
+              results.concat(page_results.map { |result| parse_batch_result(result) })
+              token = response['pagination_token'] || response['next_page_token']
+              break unless token
+            end
+
+            results
+          end
+
+          private
+
+          def xai_batch_request(request)
+            payload = batch_payload(request)
+            {
+              batch_request_id: request[:custom_id],
+              batch_request: { batch_request_type(payload) => payload }
+            }
+          end
+
+          def batch_request_type(payload)
+            payload.key?(:input) || payload.key?('input') ? :responses : :chat_get_completion
+          end
+
+          def parse_batch_response(data)
+            state = data['state'] || {}
+            completed = completed_batch_state?(state)
+            {
+              id: data['batch_id'] || data['id'],
+              raw_status: data['state'] ? xai_batch_status(state, completed:) : data['status'],
+              completed:,
+              request_count: state['num_requests'],
+              request_counts: state
+            }
+          end
+
+          def parse_batch_status(raw_status, completed:)
+            return :pending unless completed
+
+            raw_status == 'failed' ? :failed : :succeeded
+          end
+
+          def xai_batch_status(state, completed:)
+            return 'failed' unless state['error'].to_s.empty?
+
+            completed ? 'completed' : 'processing'
+          end
+
+          def completed_batch_state?(state)
+            state['num_requests'].to_i.positive? && state['num_pending'].to_i.zero?
+          end
+
+          def parse_batch_result(result)
+            request_id = result['batch_request_id'] || result['custom_id']
+            index = batch_result_index(request_id)
+            body = result.dig('batch_result', 'response', 'chat_get_completion') ||
+                   result.dig('response', 'chat_get_completion')
+
+            if body
+              [index, parse_completion_body(body, raw: body)]
+            else
+              [index, nil, batch_failure(request_id, batch_error_message(result))]
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/xai/chat_completions.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      # xAI's dialect of the Chat Completions API.
+      class ChatCompletions < Protocols::ChatCompletions
+        include XAI::Chat
+        include XAI::ReportedCost
+        include XAI::Images
+        include XAI::Models
+        include XAI::Speech
+        include XAI::Transcription
+        include Protocols::XAI::Tokenization
+        include Protocols::XAI::StreamingTranscription
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/xai/images.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      # Image generation and editing for the xAI API. Generation rejects the
+      # size parameter. Editing takes JSON image references (URL, data URI,
+      # or file id) instead of multipart uploads, up to three per request,
+      # and has no mask support.
+      module Images
+        module_function
+
+        def images_url(with: nil, mask: nil)
+          editing?(with, mask) ? 'images/edits' : 'images/generations'
+        end
+
+        def render_image_payload(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {})
+          return render_edit_payload(prompt, model:, count:, with:, provider_options:) if editing?(with, mask)
+
+          RubyLLM.logger.debug { "Ignoring size #{size}. xAI image generation does not support a size parameter." }
+          payload = { model: model, prompt: prompt }
+          payload[:n] = count if count
+
+          payload.merge(provider_options)
+        end
+
+        def render_edit_payload(prompt, model:, with:, provider_options:, count: nil)
+          payload = {
+            model: model,
+            prompt: prompt,
+            images: image_references(with)
+          }
+          payload[:n] = count if count
+
+          payload.merge(provider_options)
+        end
+
+        def image_references(sources)
+          Array(sources).filter_map do |source|
+            next if blank_attachment?(source)
+
+            { type: 'image_url', url: image_reference_url(source) }
+          end
+        end
+
+        def image_reference_url(source)
+          attachment = Attachment.new(source, config: @config)
+          return attachment.provider_file_id if attachment.provider_file?
+          return attachment.source.to_s if attachment.url?
+
+          raise UnsupportedAttachmentError, attachment.mime_type unless attachment.image?
+
+          attachment.for_llm
+        end
+
+        def parse_image_response(response, model:)
+          parse_image_responses(response, model:).first
+        end
+
+        def parse_image_responses(response, model:)
+          data = response.body
+          entries = Array(data['data'])
+
+          raise Error, 'Unexpected response format from xAI image API' if entries.empty?
+
+          entries.map.with_index do |image_data, index|
+            Image.new(
+              url: image_data['url'],
+              data: image_data['b64_json'],
+              mime_type: image_data['mime_type'] || 'image/png',
+              model: model,
+              usage: index.zero? ? (data['usage'] || {}) : {}
+            )
+          end
+        end
+
+        def validate_paint_inputs!(with:, mask:) # rubocop:disable Lint/UnusedMethodArgument
+          raise Error, 'xAI image editing does not support a mask parameter' if mask
+        end
+
+        def editing?(with, mask)
+          Protocols::ChatCompletions::Images.editing?(with, mask)
+        end
+
+        def blank_attachment?(value)
+          Protocols::ChatCompletions::Images.blank_attachment?(value)
+        end
+      end
+    end
+  end
+end