ruby_llm 1.16.0 → 2.0.0.rc4

data/lib/ruby_llm/protocols/gemini/embedding_batches.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      # Inline embedding requests and results for Gemini's asynchronous Batch API.
+      module EmbeddingBatches
+        private
+
+        def embedding_batch?(requests)
+          kinds = requests.map { |request| embedding_batch_payload?(request.fetch(:payload)) }.uniq
+          raise Error, 'Gemini batches take chat or embedding requests, not both' if kinds.size > 1
+
+          kinds.first
+        end
+
+        def embedding_batch_payload?(payload)
+          payload.key?(:content) || payload.key?(:requests)
+        end
+
+        def embedding_batch_response?(data)
+          batch = data['metadata'] || data
+          batch['@type']&.end_with?('.EmbedContentBatch') ||
+            data.dig('response', '@type')&.end_with?('.EmbedContentBatchOutput') ||
+            data.dig('response', 'inlinedEmbedContentResponses') ||
+            data.dig('output', 'inlinedEmbedContentResponses')
+        end
+
+        def embedding_batch_requests(request, model)
+          payload = request.fetch(:payload)
+          array_input = payload.key?(:requests)
+          inputs = array_input ? payload.fetch(:requests) : [payload]
+          raise ArgumentError, 'Gemini embedding batches require at least one text per request' if inputs.empty?
+
+          inputs.each_with_index.map do |input, index|
+            {
+              request: input.merge(model: "models/#{model}"),
+              metadata: {
+                custom_id: request.fetch(:custom_id), model:, array_input:,
+                embedding_index: index, embedding_count: inputs.size
+              }
+            }
+          end
+        end
+
+        def parse_embedding_batch_results(responses)
+          groups = responses.each_with_index.group_by do |inline, index|
+            inline.dig('metadata', 'custom_id') || index.to_s
+          end
+          groups.filter_map do |key, indexed|
+            parse_embedding_batch_group(key, indexed.map(&:first))
+          end
+        end
+
+        def parse_embedding_batch_group(key, responses)
+          index = batch_result_index(key)
+          error = responses.find { |inline| inline['error'] }
+          return [index, nil, batch_failure(key, error.dig('error', 'message'))] if error
+
+          metadata = responses.first.fetch('metadata')
+          return if responses.size < metadata.fetch('embedding_count')
+
+          positions = responses.map { |inline| inline.dig('metadata', 'embedding_index') }
+          unless positions.sort == (0...responses.size).to_a
+            return [index, nil, batch_failure(key, 'Invalid or duplicate embedding record positions')]
+          end
+
+          vectors = embedding_batch_vectors(responses)
+          return [index, nil, batch_failure(key, 'Gemini returned no embedding')] unless vectors
+
+          embedding = Embedding.new(
+            vectors: metadata['array_input'] ? vectors : vectors.first,
+            model: metadata.fetch('model'), input_tokens: embedding_batch_tokens(responses)
+          )
+          [index, embedding]
+        end
+
+        def embedding_batch_vectors(responses)
+          ordered = responses.sort_by { |inline| inline.dig('metadata', 'embedding_index') }
+          vectors = ordered.map { |inline| inline.dig('response', 'embedding', 'values') }
+          vectors if vectors.all? { |vector| vector && !vector.empty? }
+        end
+
+        def embedding_batch_tokens(responses)
+          counts = responses.filter_map { |inline| inline.dig('response', 'usageMetadata', 'promptTokenCount') }
+          counts.sum unless counts.empty?
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/gemini/embeddings.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      # Embeddings methods for the Gemini API integration
+      module Embeddings
+        module_function
+
+        def embedding_url(model:)
+          "models/#{model}:batchEmbedContents"
+        end
+
+        def render_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil, with: [],
+                                     provider_options: {})
+          requests = if with.any?
+                       raise ArgumentError, 'embed one text at a time when embedding attachments' if text.is_a?(Array)
+
+                       [media_embedding_payload(text, with, model:, dimensions:, task_type:, title:)]
+                     else
+                       [text].flatten.map do |t|
+                         single_embedding_payload(t, model:, dimensions:, task_type:, title:)
+                       end
+                     end
+
+          Support::Utils.deep_merge({ requests: requests }, provider_options)
+        end
+
+        def supports_embedding_media?
+          true
+        end
+
+        def render_embedding(text, dimensions: nil, **options)
+          payload = render_embedding_payload(text, dimensions:, **options)
+          return payload if text.is_a?(Array) || !payload.key?(:requests)
+
+          payload.fetch(:requests).first
+        end
+        public :render_embedding
+
+        def parse_embedding_response(response, model:, text:)
+          vectors = response.body['embeddings']&.map { |e| e['values'] }
+          vectors = vectors.first if vectors&.length == 1 && !text.is_a?(Array)
+
+          Embedding.new(vectors:, model:)
+        end
+
+        private
+
+        def single_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil)
+          embedding_payload([{ text: text.to_s }], model:, dimensions:, task_type:, title:)
+        end
+
+        def media_embedding_payload(text, attachments, model:, dimensions:, task_type: nil, title: nil)
+          embedding_payload(Media.format_content(text, attachments), model:, dimensions:, task_type:, title:)
+        end
+
+        def embedding_payload(parts, model:, dimensions:, task_type:, title:)
+          {
+            model: "models/#{model}",
+            content: { parts: parts },
+            outputDimensionality: dimensions,
+            taskType: task_type,
+            title: title
+          }.compact
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/gemini/file_transcription.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      module FileTranscription # :nodoc: all
+        def transcribe(audio_file, model:, language:, format: nil, speaker_names: nil,
+                       speaker_references: nil, provider_options: {}, prompt: nil, temperature: nil)
+          raise_transcription_streaming_unsupported if block_given?
+          validate_transcription_request(format:, speaker_references:, temperature:)
+          attachments = Attachment.wrap(audio_file, config: @config)
+          unless attachments.one? && attachments.first.audio?
+            raise ArgumentError, 'Dedicated transcription requires exactly one audio file'
+          end
+
+          track_usage(:transcription) do
+            payload = render_transcription_payload(attachments.first, model:, language:, speaker_names:,
+                                                                      provider_options:, prompt:)
+            response = @connection.post(transcription_url(model), payload, usage: @usage_tracker)
+            parse_transcription_response(response, model:)
+          end
+        end
+
+        def validate_transcription_request(format:, speaker_references:, temperature:)
+          return unless format || speaker_references || temperature
+
+          raise ArgumentError, 'Dedicated transcription does not accept format, speaker references, or temperature'
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      # Gemini Files API.
+      class Files < Protocols::Files
+        PROCESSING_POLL_INTERVAL = 2
+        PROCESSING_TIMEOUT = 600
+
+        # 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:)
+          display_name = provider_options[:display_name] || attachment.filename
+          upload_url = start_resumable_upload(attachment, display_name:)
+          response = upload_file_bytes(upload_url, attachment)
+          await_active(parse_file_response(response.body.fetch('file')))
+        end
+
+        def download(file_id)
+          file = find(file_id)
+          download_uri = file.metadata['downloadUri']
+          raise Error, 'gemini file has no download URI' unless download_uri
+
+          gemini_connection(request_json: false, response_json: false).get(download_uri) do |request|
+            request.headers.merge!(@provider.headers)
+          end.body
+        end
+
+        private
+
+        # Gemini rejects a file reference until processing finishes, which
+        # for video takes far longer than the upload itself.
+        def await_active(file)
+          deadline = Time.now + PROCESSING_TIMEOUT
+
+          while file.status == 'PROCESSING'
+            raise Error, "gemini is still processing #{file.id}" if Time.now >= deadline
+
+            sleep PROCESSING_POLL_INTERVAL
+            file = find(file.id)
+          end
+
+          raise Error, "gemini failed to process #{file.id}" if file.status == 'FAILED'
+
+          file
+        end
+
+        def file_info_url(file_id)
+          gemini_file_name(file_id)
+        end
+
+        def start_resumable_upload(attachment, display_name:)
+          response = gemini_connection(request_json: true, response_json: true).post(gemini_upload_url) do |request|
+            request.headers.merge!(@provider.headers)
+            request.headers['X-Goog-Upload-Protocol'] = 'resumable'
+            request.headers['X-Goog-Upload-Command'] = 'start'
+            request.headers['X-Goog-Upload-Header-Content-Length'] = attachment.content.bytesize.to_s
+            request.headers['X-Goog-Upload-Header-Content-Type'] = file_content_type(attachment)
+            request.headers['Content-Type'] = 'application/json'
+            request.body = { file: { display_name: display_name } }
+          end
+
+          response.headers['x-goog-upload-url'] || response.headers['X-Goog-Upload-URL'] ||
+            raise(Error, 'gemini did not return an upload URL')
+        end
+
+        def upload_file_bytes(upload_url, attachment)
+          gemini_connection(request_json: false, response_json: true).post(upload_url) do |request|
+            request.headers.merge!(@provider.headers)
+            request.headers['Content-Length'] = attachment.content.bytesize.to_s
+            request.headers['X-Goog-Upload-Offset'] = '0'
+            request.headers['X-Goog-Upload-Command'] = 'upload, finalize'
+            request.body = attachment.content
+          end
+        end
+
+        def parse_file_response(data)
+          uploaded_file(
+            data,
+            id: data['name'],
+            filename: data['displayName'],
+            byte_size: data['sizeBytes']&.to_i,
+            created_at: timestamp(data['createTime']),
+            expires_at: timestamp(data['expirationTime']),
+            mime_type: data['mimeType'],
+            status: data['state'],
+            uri: data['uri']
+          )
+        end
+
+        def gemini_file_name(file_id)
+          file_id.to_s.start_with?('files/') ? file_id : "files/#{file_id}"
+        end
+
+        def gemini_upload_url
+          base = @provider.api_base.to_s.sub(%r{/+\z}, '')
+          "#{base.sub(%r{/v1beta\z}, '/upload/v1beta').sub(%r{/v1\z}, '/upload/v1')}/files"
+        end
+
+        def gemini_connection(request_json:, response_json:)
+          connection = Transport::Connection.basic(@config) do |connection|
+            connection.request :json if request_json
+            connection.response :json if response_json
+            connection.adapter @config.faraday_adapter
+            connection.use :llm_errors, provider: @provider
+          end
+          connection.url_prefix = @provider.api_base
+          connection
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/gemini/images.rb

@@ -0,0 +1,183 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      # Image generation methods for the Gemini API implementation
+      module Images
+        IMAGE_SIZES = %w[512 512P 512PX 1K 2K 4K].freeze
+        ASPECT_RATIO = /\A\d+:\d+\z/
+        PIXEL_DIMENSIONS = /\A(\d+)\s*[x×]\s*(\d+)\z/i
+
+        def images_url(with: nil, mask: nil) # rubocop:disable Lint/UnusedMethodArgument
+          id = model_id(@model)
+
+          "models/#{id}:#{image_endpoint_action(id)}"
+        end
+
+        def render_image_payload(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {}) # rubocop:disable Lint/UnusedMethodArgument
+          @model = model
+          payload = if gemini_image_model?(model)
+                      render_gemini_image_payload(prompt, with:, count:, size:)
+                    else
+                      render_imagen_payload(prompt, count:, size:)
+                    end
+
+          Support::Utils.deep_merge(payload, provider_options)
+        end
+
+        def parse_image_response(response, model:)
+          parse_image_responses(response, model:).first
+        end
+
+        def parse_image_responses(response, model:)
+          data = response.body
+          return parse_gemini_image_responses(data, model:) if gemini_image_model?(model)
+
+          parse_imagen_responses(data, model:)
+        end
+
+        private
+
+        def validate_paint_inputs!(with:, mask:)
+          if gemini_image_model?(@model)
+            raise UnsupportedAttachmentError, 'image mask' if mask
+
+            return
+          end
+
+          return if with.nil? && mask.nil?
+
+          raise UnsupportedAttachmentError, 'image reference'
+        end
+
+        def render_imagen_payload(prompt, count: nil, size: nil)
+          RubyLLM.logger.debug { "Ignoring size #{size}. Imagen sizing is not supported." } if size
+
+          {
+            instances: [
+              {
+                prompt: prompt
+              }
+            ],
+            parameters: {
+              sampleCount: count || 1
+            }
+          }
+        end
+
+        def render_gemini_image_payload(prompt, with:, count: nil, size: nil)
+          generation_config = { responseModalities: %w[TEXT IMAGE] }
+          generation_config[:candidateCount] = count if count && count > 1
+          image_config = build_image_config(size)
+          generation_config[:imageConfig] = image_config if image_config
+
+          {
+            contents: [
+              {
+                role: 'user',
+                parts: Media.format_content(prompt, image_attachments(with))
+              }
+            ],
+            generationConfig: generation_config
+          }
+        end
+
+        # Gemini sizes an image by aspect ratio and resolution tier rather
+        # than by pixel dimensions, so WxH becomes the ratio it reduces to.
+        def build_image_config(size)
+          value = size.to_s.strip
+          return nil if value.empty?
+          return { imageSize: value.upcase } if IMAGE_SIZES.include?(value.upcase)
+          return { aspectRatio: value } if value.match?(ASPECT_RATIO)
+
+          match = value.match(PIXEL_DIMENSIONS)
+          raise ArgumentError, unsupported_size_message(size) unless match
+
+          { aspectRatio: aspect_ratio(match[1].to_i, match[2].to_i) }
+        end
+
+        def aspect_ratio(width, height)
+          raise ArgumentError, unsupported_size_message("#{width}x#{height}") unless width.positive? && height.positive?
+
+          divisor = width.gcd(height)
+          "#{width / divisor}:#{height / divisor}"
+        end
+
+        def unsupported_size_message(size)
+          "Gemini cannot generate an image of size #{size.inspect}. Give pixel dimensions such as " \
+            '"1024x1024", an aspect ratio such as "16:9", or a resolution ' \
+            "such as #{IMAGE_SIZES.join(', ')}."
+        end
+
+        def parse_imagen_responses(data, model:)
+          predictions = Array(data['predictions']).select { |prediction| prediction['bytesBase64Encoded'] }
+          raise Error, 'Unexpected response format from Gemini image generation API' if predictions.empty?
+
+          predictions.map do |image_data|
+            Image.new(
+              data: image_data['bytesBase64Encoded'],
+              mime_type: image_data['mimeType'] || 'image/png',
+              model: model
+            )
+          end
+        end
+
+        def parse_gemini_image_responses(data, model:)
+          parts = gemini_image_parts(data)
+          raise Error, 'Unexpected response format from Gemini image generation API' if parts.empty?
+
+          parts.map.with_index do |image_data, index|
+            Image.new(
+              data: image_data['data'],
+              mime_type: image_data['mimeType'] || 'image/png',
+              model: data['modelVersion'] || model,
+              usage: index.zero? ? gemini_image_usage(data) : {}
+            )
+          end
+        end
+
+        def gemini_image_parts(data)
+          Array(data['candidates']).filter_map do |candidate|
+            parts = candidate.dig('content', 'parts') || []
+            parts.filter_map { |part| part['inlineData'] }.find do |inline_data|
+              image_mime_type?(inline_data['mimeType']) && inline_data['data']
+            end
+          end
+        end
+
+        def image_mime_type?(mime_type)
+          mime_type.nil? || mime_type.start_with?('image/')
+        end
+
+        def gemini_image_usage(data)
+          {
+            'input_tokens' => input_tokens(data),
+            'output_tokens' => calculate_output_tokens(data)
+          }.compact
+        end
+
+        def image_attachments(sources)
+          Attachment.wrap(sources, config: @config).each do |attachment|
+            raise UnsupportedAttachmentError, attachment.mime_type unless attachment.image?
+          end
+        end
+
+        def gemini_image_model?(model)
+          id = model_id(model).downcase
+          return true if id.start_with?('nano-banana', 'nanobanana')
+
+          id.start_with?('gemini-') && id.include?('-image')
+        end
+
+        def image_endpoint_action(model)
+          gemini_image_model?(model) ? 'generateContent' : 'predict'
+        end
+
+        def model_id(model)
+          model.respond_to?(:id) ? model.id : model.to_s
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/gemini/live_transcription.rb

@@ -0,0 +1,140 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      class LiveTranscription < Protocol # :nodoc: all
+        def transcribe(audio_file, model:, language:, format: nil, speaker_names: nil,
+                       speaker_references: nil, provider_options: {}, prompt: nil, temperature: nil, &block)
+          validate_transcription_request(format:, speaker_names:, speaker_references:, temperature:)
+          audio = transcription_audio(audio_file)
+          setup = render_transcription_setup(model:, language:, prompt:, provider_options:)
+
+          track_usage(:transcription) do
+            @usage_tracker.start
+            events = collect_transcription(audio, setup, &block)
+            result = parse_transcription_events(events, audio:, model:)
+            block&.call(TranscriptionChunk.new(type: TranscriptionChunk::DONE, text: result.text, raw: events.last))
+            result
+          end
+        end
+
+        def validate_transcription_request(format:, speaker_names:, speaker_references:, temperature:)
+          return unless format || speaker_names || speaker_references || temperature
+
+          raise ArgumentError, 'Google Live transcription does not accept format, diarization, or temperature'
+        end
+
+        def transcription_audio(file)
+          attachments = Attachment.wrap(file, config: @config)
+          raise ArgumentError, 'Transcription requires exactly one audio file' unless attachments.one?
+
+          audio = RubyLLM::Transcription::WavAudio.new(attachments.first.content)
+          unless [audio.encoding, audio.channels, audio.bits_per_sample] == [1, 1, 16] &&
+                 audio.data.bytesize.positive? && audio.data.bytesize.even?
+            raise ArgumentError, 'Google Live transcription requires non-empty mono 16-bit PCM WAV audio'
+          end
+
+          audio
+        end
+
+        def render_transcription_setup(model:, language:, prompt:, provider_options:)
+          payload = { model: transcription_model_name(model), generationConfig: { responseModalities: ['TEXT'] },
+                      inputAudioTranscription: { languageCodes: language && Array(language),
+                                                 customVocabulary: prompt && Array(prompt) }.compact,
+                      realtimeInputConfig: { automaticActivityDetection: { disabled: true } } }
+          payload = Support::Utils.deep_merge(payload, provider_options)
+          validate_transcription_setup(payload)
+          { setup: payload }
+        end
+
+        def validate_transcription_setup(payload)
+          config = payload.fetch(:inputAudioTranscription)
+          if config[:diarization] || config[:wordTimestamp]
+            raise ArgumentError, 'Google Live transcription does not support diarization or word timestamps'
+          end
+          return if payload.dig(:realtimeInputConfig, :automaticActivityDetection, :disabled) == true
+
+          raise ArgumentError, 'Google file transcription requires manual activity boundaries'
+        end
+
+        def transcription_model_name(model)
+          "models/#{model}"
+        end
+
+        def websocket_service
+          'google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent'
+        end
+
+        def transcription_websocket_url
+          uri = URI.parse(@provider.api_base)
+          uri.scheme = uri.scheme == 'https' ? 'wss' : 'ws'
+          uri.path = "/ws/#{websocket_service}"
+          uri.to_s
+        end
+
+        def collect_transcription(audio, setup, &block)
+          events = []
+          ready = Queue.new
+          Transport::WebsocketConnection.open(
+            transcription_websocket_url, headers: @provider.headers.transform_keys(&:to_s), config: @config
+          ) do |socket|
+            socket.send_text(JSON.generate(setup))
+            write = lambda { |connection|
+              ready.pop
+              send_transcription_audio(connection, audio)
+            }
+            socket.each_message(write:) do |message|
+              event = JSON.parse(message)
+              events << event
+              ready << true if event.key?('setupComplete')
+              process_transcription_event(event, &block)
+              socket.close if event.dig('serverContent', 'generationComplete')
+            end
+          end
+          unless events.last&.dig('serverContent', 'generationComplete')
+            raise Error, 'Google Live transcription ended before generation completed'
+          end
+
+          events
+        end
+
+        def send_transcription_audio(socket, audio)
+          socket.send_text(JSON.generate(realtimeInput: { activityStart: {} }))
+          bytes = [audio.sample_rate / 10, 1].max * 2
+          offset = 0
+          while offset < audio.data.bytesize
+            chunk = audio.data.byteslice(offset, bytes)
+            input = { audio: { mimeType: "audio/pcm;rate=#{audio.sample_rate}", data: Base64.strict_encode64(chunk) } }
+            socket.send_text(JSON.generate(realtimeInput: input))
+            offset += chunk.bytesize
+          end
+          socket.send_text(JSON.generate(realtimeInput: { activityEnd: {} }))
+        end
+
+        def process_transcription_event(event)
+          raise Error, event.dig('error', 'message') || 'Google Live transcription failed' if event['error']
+          return unless block_given?
+
+          content = event['serverContent'] || {}
+          if content['interimInputTranscription']
+            yield TranscriptionChunk.new(type: TranscriptionChunk::PARTIAL,
+                                         text: content.dig('interimInputTranscription', 'text'), raw: event)
+          elsif content['inputTranscription']
+            yield TranscriptionChunk.new(type: TranscriptionChunk::DELTA,
+                                         delta: content.dig('inputTranscription', 'text'), raw: event)
+          end
+        end
+
+        def parse_transcription_events(events, audio:, model:)
+          text = events.filter_map { |event| event.dig('serverContent', 'inputTranscription', 'text') }.join
+          usage = events.reverse.find { |event| event['usageMetadata'] }&.fetch('usageMetadata') || {}
+          RubyLLM::Transcription.new(
+            text:, model:, duration: audio.duration,
+            input_tokens: usage['promptTokenCount'], output_tokens: usage['candidatesTokenCount']
+          )
+        end
+      end
+    end
+  end
+end