ruby_llm 2.1.1 → 2.2.0

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

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      # Gemini can pause a long response with a continuation token. Each
+      # hop resends the same request with the newest token, and the hops
+      # join into one assistant message.
+      module LongDecode # :nodoc:
+        def complete(messages, **, &)
+          @unfinished = messages.last if messages.last&.continuation
+          super
+        end
+
+        def complete_request(payload, headers, usage_recorder: nil, progress_recorder: nil, &)
+          segments = [@unfinished].compact
+          parts = Array(@unfinished&.raw_content)
+
+          loop do
+            hop = super(payload, headers, usage_recorder:, &)
+            segments << hop
+            parts.concat(hop_parts(hop))
+            return segments.one? ? hop : join_hops(segments, parts) unless hop.continuation
+
+            if segments[...-1].any? { |segment| segment.continuation == hop.continuation }
+              raise Error.new('Gemini repeated a continuation token', response: hop.raw)
+            end
+
+            progress_recorder&.call(join_hops(segments, parts, continuation: hop.continuation))
+            payload = payload.merge(continuationToken: hop.continuation)
+          end
+        end
+
+        private
+
+        def stream_response(...)
+          @stream_parts = @stream_metadata_calls = @stream_citation_parts = nil
+          super
+        end
+
+        def hop_parts(hop)
+          hop.raw_content || @stream_parts || hop.raw&.body&.dig('candidates', 0, 'content', 'parts') || []
+        end
+
+        def join_hops(segments, parts, continuation: nil)
+          final = segments.last
+          Message.new(
+            role: :assistant,
+            content: join_text(segments.map(&:content)),
+            thinking: join_thinking(segments),
+            attachments: segments.flat_map(&:attachments),
+            citations: join_citations(segments),
+            tool_calls: join_tool_calls(segments),
+            server_tool_calls: segments.flat_map(&:server_tool_calls),
+            raw_content: join_raw_content(parts, continuation),
+            continuation: continuation,
+            tokens: Tokens.aggregate(segments.map(&:tokens)),
+            usage_entries: segments.flat_map(&:ruby_llm_usage_entries),
+            finish_reason: final.finish_reason,
+            model: final.model,
+            raw: final.raw
+          )
+        end
+
+        def join_text(values)
+          text = values.compact.join
+          text unless text.empty?
+        end
+
+        def join_thinking(segments)
+          thoughts = segments.filter_map(&:thinking)
+          Thinking.build(text: join_text(thoughts.map(&:text)), signature: thoughts.filter_map(&:signature).last)
+        end
+
+        def join_tool_calls(segments)
+          calls = segments.filter_map(&:tool_calls).reduce({}, :merge)
+          calls unless calls.empty?
+        end
+
+        def join_raw_content(parts, continuation)
+          return without_search_suggestions(parts) if continuation
+          return unless parts.any? { |part| server_tool_part?(part) } ||
+                        parts.count { |part| part['thoughtSignature'] } > 1
+
+          without_search_suggestions(parts)
+        end
+
+        def join_citations(segments)
+          offset = 0
+          segments.flat_map do |segment|
+            citations = segment.citations.map do |citation|
+              next citation unless citation.start_index
+
+              Citation.new(citation.to_h.merge(start_index: citation.start_index + offset,
+                                               end_index: citation.end_index + offset))
+            end
+            offset += segment.content.to_s.length
+            citations
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -14,31 +14,27 @@ module RubyLLM
           %i[image audio video pdf text].include?(attachment.type)
         end
 
-        def format_content(content, attachments = [])
+        def render_content(content, attachments = [])
           parts = []
-          parts << format_text(content) if content
+          parts << render_text(content) if content
 
           attachments.each do |attachment|
-            parts << with_media_resolution(format_content_attachment(attachment), attachment)
+            part = with_media_resolution(render_content_attachment(attachment), attachment)
+            parts << with_media_processing(part, attachment)
           end
 
           parts
         end
 
-        def format_content_attachment(attachment)
+        def render_content_attachment(attachment)
           raise UnsupportedAttachmentError, attachment.mime_type unless supported_attachment?(attachment)
 
-          case attachment.type
-          when :text
-            format_text_file(attachment)
-          else
-            format_attachment(attachment)
-          end
-        end
+          return render_file_data(attachment) if attachment.provider_file?
 
-        def format_attachment(attachment)
-          return format_file_data(attachment) if attachment.provider_file?
+          attachment.type == :text ? render_text_file(attachment) : render_attachment(attachment)
+        end
 
+        def render_attachment(attachment)
           {
             inline_data: {
               mime_type: attachment.mime_type,
@@ -55,7 +51,13 @@ module RubyLLM
           part.merge(media_resolution: { level: "MEDIA_RESOLUTION_#{level.upcase}" })
         end
 
-        def format_file_data(attachment)
+        def with_media_processing(part, attachment)
+          return part unless attachment.processing && attachment.video?
+
+          part.merge(media_processing: attachment.processing.to_s.upcase)
+        end
+
+        def render_file_data(attachment)
           uri = attachment.provider_file_uri || attachment.provider_file_id
           {
             file_data: {
@@ -65,17 +67,28 @@ module RubyLLM
           }
         end
 
-        def format_text_file(text_file)
+        def render_text_file(text_file)
           {
             text: text_file.for_llm
           }
         end
 
-        def format_text(text)
+        def render_text(text)
           {
             text: text
           }
         end
+
+        def attachment_filename(mime_type, index)
+          return "gemini_attachment_#{index + 1}" unless mime_type
+
+          extension = mime_type.split('/').last.to_s
+          extension = 'jpg' if extension == 'jpeg'
+          extension = 'txt' if extension == 'plain'
+          extension = 'mp3' if mime_type == 'audio/mpeg'
+          extension = extension.tr('+', '.')
+          "gemini_attachment_#{index + 1}.#{extension}"
+        end
       end
 
       def build_response_content(parts) # rubocop:disable Metrics/PerceivedComplexity
@@ -108,7 +121,7 @@ module RubyLLM
         io = StringIO.new(decoded)
         io.set_encoding(Encoding::BINARY) if io.respond_to?(:set_encoding)
 
-        filename = attachment_filename(mime_type, index)
+        filename = Media.attachment_filename(mime_type, index)
         RubyLLM::Attachment.new(io, filename:)
       rescue ArgumentError => e
         RubyLLM.logger.warn "Failed to decode Gemini inline data attachment: #{e.message}"
@@ -119,19 +132,9 @@ module RubyLLM
         uri = file_data['fileUri']
         return unless uri
 
-        filename = file_data['filename'] || attachment_filename(file_data['mimeType'], index)
+        filename = file_data['filename'] || Media.attachment_filename(file_data['mimeType'], index)
         RubyLLM::Attachment.new(uri, filename:)
       end
-
-      def attachment_filename(mime_type, index)
-        return "gemini_attachment_#{index + 1}" unless mime_type
-
-        extension = mime_type.split('/').last.to_s
-        extension = 'jpg' if extension == 'jpeg'
-        extension = 'txt' if extension == 'plain'
-        extension = extension.tr('+', '.')
-        "gemini_attachment_#{index + 1}.#{extension}"
-      end
     end
   end
 end

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

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      # Google says how long to wait in the error body, as a
+      # google.rpc.RetryInfo detail, rather than in a header.
+      module RetryInfo # :nodoc: all
+        DURATION = /\A\d{1,12}(\.\d{1,9})?s\z/
+
+        def retry_delay(response)
+          retry_info_details(parse_error_body(response)).filter_map { |detail| retry_info_seconds(detail) }.max
+        end
+
+        private
+
+        def retry_info_details(body)
+          errors = [body].flatten.grep(Hash).map { |part| part['error'] }.grep(Hash)
+          errors.flat_map { |error| error['details'] }.grep(Hash)
+        end
+
+        def retry_info_seconds(detail)
+          delay = detail['retryDelay']
+          delay.to_f if delay.is_a?(String) && delay.match?(DURATION)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Gemini
+      # JSON Schemas as Gemini accepts them. Gemini refuses a schema nested
+      # more than 32 levels deep with a bare "Request contains an invalid
+      # argument", so a deeper one goes with its deep subschemas moved into
+      # $defs and referenced from where they were, which describes the same
+      # values in fewer levels.
+      module Schemas
+        MAX_DEPTH = 32
+        HOIST_DEPTH = 16
+        DEFINITION_KEYS = %w[$defs definitions].freeze
+        SCHEMA_MAPS = %w[properties patternProperties dependentSchemas $defs definitions].freeze
+        SCHEMA_LISTS = %w[anyOf oneOf allOf prefixItems].freeze
+        SCHEMA_VALUES = %w[items additionalProperties additionalItems unevaluatedProperties unevaluatedItems
+                           not if then else contains propertyNames].freeze
+
+        module_function
+
+        def render_schema(schema)
+          return schema unless schema.is_a?(Hash) && nesting(schema) > MAX_DEPTH
+
+          schema = Support::Utils.deep_stringify_keys(schema)
+          movable?(schema) ? Hoisting.new(schema).schema : schema
+        end
+
+        def movable?(schema)
+          references_only_definitions?(schema) && !identified?(schema.except('$id'))
+        end
+
+        def nesting(value)
+          case value
+          when Hash then 1 + (value.values.map { |child| nesting(child) }.max || 0)
+          when Array then 1 + (value.map { |child| nesting(child) }.max || 0)
+          else 0
+          end
+        end
+
+        # A reference into the schema's own structure, such as
+        # "#/properties/address", would point at a subschema that moved.
+        def references_only_definitions?(value)
+          case value
+          when Hash
+            value.all? do |key, child|
+              key == '$ref' ? definition_reference?(child) : references_only_definitions?(child)
+            end
+          when Array then value.all? { |child| references_only_definitions?(child) }
+          else true
+          end
+        end
+
+        # Only a whole definition, since a pointer into one could pass
+        # through a subschema that moved.
+        def definition_reference?(reference)
+          !reference.is_a?(String) || !reference.start_with?('#') ||
+            DEFINITION_KEYS.any? { |key| reference.match?(%r{\A#/#{Regexp.escape(key)}/[^/]+\z}) }
+        end
+
+        # A nested $id starts a resource of its own, whose references
+        # resolve against it rather than the root.
+        def identified?(value)
+          case value
+          when Hash then value.key?('$id') || value.values.any? { |child| identified?(child) }
+          when Array then value.any? { |child| identified?(child) }
+          else false
+          end
+        end
+
+        # Rewrites one schema, moving each subschema at HOIST_DEPTH or
+        # deeper into the root's definitions.
+        class Hoisting
+          attr_reader :schema
+
+          def initialize(schema)
+            @key = DEFINITION_KEYS.find { |key| schema[key].is_a?(Hash) } || '$defs'
+            @definitions = schema.fetch(@key, {}).dup
+            @count = 0
+            @definitions.keys.each { |name| @definitions[name] = rewrite(@definitions[name], 3) }
+            @schema = rewrite(schema.except(@key), 1).merge(@key => @definitions)
+          end
+
+          private
+
+          def rewrite(schema, depth)
+            return schema unless schema.is_a?(Hash)
+            return hoist(schema) if depth >= HOIST_DEPTH
+
+            schema.to_h { |key, value| [key, rewrite_value(key, value, depth)] }
+          end
+
+          def rewrite_value(key, value, depth)
+            case [position(key), value]
+            in [:map, Hash] then value.transform_values { |child| rewrite(child, depth + 2) }
+            in [:list, Array] then value.map { |child| rewrite(child, depth + 2) }
+            in [:value, _] then rewrite(value, depth + 1)
+            else value
+            end
+          end
+
+          def position(key)
+            if SCHEMA_MAPS.include?(key) then :map
+            elsif SCHEMA_LISTS.include?(key) then :list
+            elsif SCHEMA_VALUES.include?(key) then :value
+            end
+          end
+
+          def hoist(schema)
+            name = next_name
+            @definitions[name] = rewrite(schema, 3)
+            { '$ref' => "#/#{@key}/#{name}" }
+          end
+
+          def next_name
+            loop do
+              @count += 1
+              name = "nested#{@count}"
+              return name unless @definitions.key?(name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -12,7 +12,9 @@ module RubyLLM
         end
 
         def render_speech_payload(input, model:, voice:, format:, provider_options: {})
-          RubyLLM.logger.debug { "Ignoring speech format #{format}. Gemini returns PCM audio." } if format
+          unless format.nil? || format.to_s == 'pcm'
+            raise ArgumentError, "Gemini speech returns PCM audio and does not accept format #{format.inspect}"
+          end
 
           payload = {
             contents: [
@@ -39,15 +41,19 @@ module RubyLLM
           Support::Utils.deep_merge(payload, provider_options)
         end
 
-        def parse_speech_response(response, model:, voice:, format:) # rubocop:disable Lint/UnusedMethodArgument
+        def parse_speech_response(response, model:, voice:, **)
           audio = response.body.dig('candidates', 0, 'content', 'parts', 0, 'inlineData', 'data')
           raise Error, 'Unexpected response format from Gemini speech generation API' unless audio
 
+          usage = response.body['usageMetadata'] || {}
+
           RubyLLM::Speech.new(
             data: Base64.decode64(audio),
             model: model,
             voice: voice || 'Kore',
-            format: 'pcm'
+            format: 'pcm',
+            input_tokens: usage['promptTokenCount'],
+            output_tokens: usage['candidatesTokenCount']
           )
         end
       end

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

@@ -18,21 +18,23 @@ module RubyLLM
           Chunk.new(
             role: :assistant,
             model: data['modelVersion'],
-            content: extract_text_content(parts),
-            citations: extract_citations(data, nil),
+            content: parse_text_content(parts),
+            attachments: build_response_content(parts.reject { |part| part['thought'] }).last,
+            citations: parse_stream_citations(data),
             thinking: Thinking.build(
-              text: extract_thought_parts(parts),
-              signature: extract_thought_signature(parts)
+              text: parse_thought_parts(parts),
+              signature: parse_thought_signature(parts)
             ),
             input_tokens: input_tokens(data),
-            output_tokens: extract_output_tokens(data),
+            output_tokens: parse_output_tokens(data),
             cache_read_tokens: data.dig('usageMetadata', 'cachedContentTokenCount'),
             thinking_tokens: data.dig('usageMetadata', 'thoughtsTokenCount'),
             server_tool_use: parse_server_tool_use(data),
             finish_reason: normalize_finish_reason(
               data.dig('candidates', 0, 'finishReason') || data.dig('promptFeedback', 'blockReason')
             ),
-            tool_calls: extract_tool_calls(data),
+            tool_calls: parse_tool_calls(data),
+            continuation: parse_continuation(data),
             **stream_end_fields(data)
           )
         end
@@ -42,12 +44,31 @@ module RubyLLM
         # Accumulates streamed parts so code-execution turns can be replayed
         # verbatim, mirroring the non-streaming path.
         def track_stream_parts(data, parts)
+          track_citation_parts(parts)
           @stream_parts = (@stream_parts || []).concat(parts)
           @saw_server_part = true if parts.any? { |part| server_tool_part?(part) }
           metadata_calls = metadata_server_tool_calls(data)
           @stream_metadata_calls = metadata_calls if metadata_calls.any?
         end
 
+        def track_citation_parts(parts)
+          @stream_citation_parts ||= []
+          previous = @stream_citation_parts.last
+          first = parts.first
+          if continued_text?(previous, first)
+            previous['text'] << first['text']
+            parts = parts.drop(1)
+          end
+          @stream_citation_parts.concat(parts.map do |part|
+            part['text'] ? part.merge('text' => part['text'].dup) : part
+          end)
+        end
+
+        def continued_text?(previous, first)
+          previous && first && previous.key?('text') && first.key?('text') &&
+            previous.fetch('thought', false) == first.fetch('thought', false)
+        end
+
         def stream_end_fields(data)
           return {} unless data.dig('candidates', 0, 'finishReason')
 
@@ -60,21 +81,24 @@ module RubyLLM
             )
           end
           calls = part_calls + @stream_metadata_calls.to_a
-          return {} if calls.empty?
+          replay = @saw_server_part || parse_continuation(data)
+          fields = { server_tool_calls: calls, raw_content: (without_search_suggestions(@stream_parts) if replay) }
+          fields.compact.reject { |_, value| value.empty? }
+        end
+
+        def parse_stream_citations(data)
+          return [] unless data.dig('candidates', 0, 'groundingMetadata')
 
-          {
-            server_tool_calls: calls,
-            raw_content: @saw_server_part ? without_search_suggestions(@stream_parts) : nil
-          }
+          parse_citations(data, @stream_citation_parts)
         end
 
-        def extract_text_content(parts)
+        def parse_text_content(parts)
           text_parts = parts.reject { |p| p['thought'] }
           text = text_parts.filter_map { |p| p['text'] }.join
           text.empty? ? nil : text
         end
 
-        def extract_output_tokens(data)
+        def parse_output_tokens(data)
           candidates = data.dig('usageMetadata', 'candidatesTokenCount') || 0
           thoughts = data.dig('usageMetadata', 'thoughtsTokenCount') || 0
           total = candidates + thoughts

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

@@ -14,7 +14,7 @@ module RubyLLM
         # make, such as one from another provider. Vertex AI accepts no other.
         PLACEHOLDER_SIGNATURE = 'skip_thought_signature_validator'
 
-        def format_tools(tools)
+        def render_tools(tools)
           return [] if tools.empty?
 
           [{
@@ -24,12 +24,13 @@ module RubyLLM
 
         # Gemini signs only the first of parallel calls, so each call goes back
         # with its own signature or none, never one another part carried.
-        def format_tool_call(msg)
+        def render_tool_call(msg)
           parts = []
 
-          parts.concat(Media.format_content(msg.content, msg.attachments)) if msg.content && !msg.content.empty?
+          text = msg.content unless msg.content.to_s.empty?
+          parts.concat(Media.render_content(text, msg.attachments))
 
-          msg.tool_calls.each_value do |tool_call|
+          msg.tool_calls.each do |tool_call|
             part = {
               functionCall: {
                 name: tool_call.name,
@@ -43,27 +44,28 @@ module RubyLLM
           parts
         end
 
-        def format_tool_result(msg, function_name = nil)
+        def render_tool_result(msg, function_name = nil)
           function_name ||= msg.tool_call_id
           content = msg.content
           content = nil if content && content.empty?
           content = '(no output)' if content.nil? && msg.attachments.empty?
 
+          texts, files = split_text_attachments(msg.attachments)
           function_response = {
             name: function_name,
             response: {
               name: function_name,
-              content: Media.format_content(content)
+              content: Media.render_content(content) + texts.map { |text| Media.render_content_attachment(text) }
             }
           }
 
-          media_parts, sibling_parts = partition_tool_result_attachments(msg.attachments)
+          media_parts, sibling_parts = partition_tool_result_attachments(files)
           function_response[:parts] = media_parts if media_parts.any?
 
           [{ functionResponse: function_response }, *sibling_parts]
         end
 
-        def extract_tool_calls(data) # rubocop:disable Metrics/PerceivedComplexity
+        def parse_tool_calls(data) # rubocop:disable Metrics/PerceivedComplexity
           return nil unless data
 
           candidate = data.is_a?(Hash) ? data.dig('candidates', 0) : nil
@@ -101,11 +103,17 @@ module RubyLLM
           parts.map { |part| part.equal?(first_call) ? part.merge(thoughtSignature: PLACEHOLDER_SIGNATURE) : part }
         end
 
+        # Vertex refuses text parts beside a functionResponse, so text files
+        # travel inside its response.
+        def split_text_attachments(attachments)
+          attachments.partition { |attachment| attachment.type == :text && !attachment.provider_file? }
+        end
+
         # functionResponse.parts only accepts inline bytes, and pre-Gemini 3 models reject it
         def partition_tool_result_attachments(attachments)
           return [[], []] if attachments.empty?
 
-          parts = attachments.map { |attachment| Media.format_content_attachment(attachment) }
+          parts = attachments.map { |attachment| Media.render_content_attachment(attachment) }
           return [[], parts] unless multimodal_function_responses_supported?(@model)
 
           parts.partition { |part| part.key?(:inline_data) }
@@ -132,7 +140,7 @@ module RubyLLM
             description: tool.description
           }
 
-          declaration[:parametersJsonSchema] = parameters_schema if parameters_schema
+          declaration[:parametersJsonSchema] = Schemas.render_schema(parameters_schema) if parameters_schema
 
           return declaration if tool.provider_options.empty?
 

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

@@ -7,13 +7,17 @@ module RubyLLM
       module Transcription
         DEFAULT_PROMPT = 'Transcribe the provided audio and respond with only the transcript text.'
 
+        def rejected_options(operation)
+          operation == :transcribe ? %i[speaker_names speaker_references] : super
+        end
+
         # rubocop:disable-next  Lint/UnusedMethodArgument
         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?
 
           track_usage(:transcription) do
-            attachment = Attachment.new(audio_file, config: @config)
+            attachment = audio_file.is_a?(Attachment) ? audio_file : Attachment.new(audio_file, config: @config)
             payload = render_transcription_payload(attachment, language:, format:, provider_options:, prompt:,
                                                                temperature:)
             response = @connection.post(transcription_url(model), payload, usage: @usage_tracker)
@@ -30,7 +34,7 @@ module RubyLLM
         def render_transcription_payload(attachment, language:, format: nil, provider_options: {}, prompt: nil,
                                          temperature: nil)
           prompt = build_prompt(prompt, language)
-          audio_part = format_audio_part(attachment)
+          audio_part = render_audio_part(attachment)
 
           raise UnsupportedAttachmentError, attachment.mime_type unless attachment.audio?
 
@@ -63,7 +67,7 @@ module RubyLLM
           prompt
         end
 
-        def format_audio_part(attachment)
+        def render_audio_part(attachment)
           {
             inline_data: {
               mime_type: attachment.mime_type,
@@ -74,10 +78,10 @@ module RubyLLM
 
         def parse_transcription_response(response, model:)
           data = response.body
-          text = extract_text(data)
+          text = parse_text(data)
           raise_if_transcription_blocked(data, text, response)
 
-          usage = extract_usage(data)
+          usage = parse_usage(data)
 
           RubyLLM::Transcription.new(
             text: text,
@@ -94,21 +98,21 @@ module RubyLLM
           reason = data.dig('candidates', 0, 'finishReason') || data.dig('promptFeedback', 'blockReason')
           return unless data.dig('candidates', 0).nil? || normalize_finish_reason(reason) == :content_filter
 
-          usage = extract_usage(data)
+          usage = parse_usage(data)
           @usage_tracker&.observe_tokens(Tokens.new(input: usage[:input_tokens], output: usage[:output_tokens]))
           raise ContentFilterError.new(['Gemini blocked the transcription', reason].compact.join(': '), response:)
         end
 
-        def extract_text(data)
+        def parse_text(data)
           candidate = data.is_a?(Hash) ? data.dig('candidates', 0) : nil
           return unless candidate
 
           parts = candidate.dig('content', 'parts') || []
-          texts = parts.filter_map { |part| part['text'] }
+          texts = parts.reject { |part| part['thought'] }.filter_map { |part| part['text'] }
           texts.join if texts.any?
         end
 
-        def extract_usage(data)
+        def parse_usage(data)
           metadata = data.is_a?(Hash) ? data['usageMetadata'] : nil
           return { input_tokens: nil, output_tokens: nil } unless metadata