ruby_llm 2.1.1 → 2.2.0

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

@@ -7,11 +7,14 @@ module RubyLLM
       # operations. The finished video is a Files API URI that requires the
       # API key to download.
       module Videos
+        OPERATION_NAME = %r{\A(?:[\w-][\w.-]*/)*operations/[\w-][\w.-]*\z}
+        private_constant :OPERATION_NAME
+
         def video_url
           "models/#{model_id(@model)}:predictLongRunning"
         end
 
-        def render_video_payload(prompt, model:, with: [], provider_options: {}) # rubocop:disable Lint/UnusedMethodArgument
+        def render_video_payload(prompt, with: [], provider_options: {}, **)
           instance = { prompt: prompt }
           instance[:image] = render_video_image(with.first) if with.first
 
@@ -36,10 +39,10 @@ module RubyLLM
         end
 
         def video_job_url(job)
-          job.id
+          video_operation_name(job)
         end
 
-        def parse_video_job_status(response, job:) # rubocop:disable Lint/UnusedMethodArgument
+        def parse_video_job_status(response, **)
           body = response.body
           if body['error']
             { status: :failed, raw: body, error: body.dig('error', 'message') }
@@ -68,6 +71,14 @@ module RubyLLM
 
         private
 
+        # The name is the URL, so a job id from outside must not be able to
+        # send the API key to another host.
+        def video_operation_name(job)
+          return job.id if job.id.match?(OPERATION_NAME)
+
+          raise ArgumentError, "#{job.id.inspect} is not a video operation name"
+        end
+
         def render_video_extension(source)
           if source.is_a?(Video) && source.raw && (video = generated_video(source.raw)) && video['uri']
             return { uri: video['uri'] }

data/lib/ruby_llm/protocols/gemini.rb

@@ -16,6 +16,7 @@ module RubyLLM
       include Gemini::Tools
       include Gemini::Speech
       include Gemini::Transcription
+      include Gemini::LongDecode
 
       # Gemini nests each tool's options inside its key, so aliases are
       # lambdas placing the options there.
@@ -28,6 +29,21 @@ module RubyLLM
         web_fetch: ->(options) { { tool: { url_context: Support::Utils.deep_symbolize_keys(options) } } }
       ).freeze
 
+      # Search results arrive as Google's redirect links. A page the model
+      # chose to read through url_context arrives in the same list under its
+      # own address, and need not be a result. That rules out the links to
+      # Google's own results page too, whose query the model could choose.
+      SEARCH_RESULT_LINK = 'https://vertexaisearch.cloud.google.com/grounding-api-redirect/'
+
+      # Returns the pages the message's searches returned, as WebSource objects.
+      def self.parse_web_sources(message)
+        searches = message.server_tool_calls.select { |call| call.type == 'google_search' }
+        pages = searches.flat_map { |call| Array(call.result) }.map { |chunk| chunk['web'] }
+        pages.select { |page| page['uri'].to_s.start_with?(SEARCH_RESULT_LINK) }.map do |page|
+          WebSource.new(url: page['uri'], title: page['title'])
+        end
+      end
+
       def server_tool_aliases
         SERVER_TOOL_ALIASES
       end

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

@@ -40,12 +40,12 @@ module RubyLLM
           super(payload, merge_mcp_filters(tools))
         end
 
-        def format_assistant_items(message)
+        def render_assistant_items(message)
           super.flat_map do |item|
             data = Support::Utils.deep_symbolize_keys(item)
             next item unless data[:type] && !Protocols::Responses::Chat::CLIENT_OUTPUT_ITEM_TYPES.include?(data[:type])
 
-            format_server_tool_history(data)
+            render_server_tool_history(data)
           end
         end
 
@@ -81,7 +81,7 @@ module RubyLLM
           existing[:allowed_tools] |= tool[:allowed_tools]
         end
 
-        def format_server_tool_history(item)
+        def render_server_tool_history(item)
           if item[:type] == 'mcp_call' && item[:name] && item[:id] && !item[:output].nil?
             output = item[:output].is_a?(String) ? item[:output] : JSON.generate(item[:output])
             [{ type: 'function_call', call_id: item[:id], name: item[:name], arguments: item[:arguments] },

data/lib/ruby_llm/protocols/gpustack/tokenization.rb

@@ -5,6 +5,8 @@ module RubyLLM
     module GPUStack
       # vLLM's text tokenizer through GPUStack's model proxy.
       module Tokenization
+        COUNT_TOKENS_KEYS = %i[model messages tools].freeze
+
         def tokenization_url
           "#{@provider.backend_api_base}/tokenize"
         end
@@ -16,6 +18,18 @@ module RubyLLM
         def parse_tokenization_response(response, model:)
           RubyLLM::Tokenization.new(ids: response.body.fetch('tokens'), model: model, raw: response.body)
         end
+
+        def count_tokens_url
+          tokenization_url
+        end
+
+        def render_count_tokens_payload(messages, **)
+          render_payload(messages, temperature: nil, **).slice(*COUNT_TOKENS_KEYS)
+        end
+
+        def parse_count_tokens_response(response)
+          response.body.fetch('count')
+        end
       end
     end
   end

data/lib/ruby_llm/protocols/interactions/chat.rb

@@ -16,13 +16,14 @@ module RubyLLM
         end
 
         def render_payload(messages, tools:, temperature:, model:, stream: false, max_output_tokens: nil,
-                           schema: nil, thinking: nil, tool_prefs: nil, **)
+                           schema: nil, thinking: nil, caching: nil, tool_prefs: nil, **)
+          ensure_interaction_carries(messages, caching)
           config = { temperature: temperature, max_output_tokens: max_output_tokens }.compact
           config.merge!(render_interaction_thinking(thinking)) if thinking
           choice = tool_prefs&.dig(:choice)
           config[:tool_choice] = render_interaction_choice(choice) if choice
           payload = {
-            model: model.id, input: format_interaction_input(messages), stream: stream, store: false,
+            model: model.id, input: render_interaction_input(messages), stream: stream, store: false,
             system_instruction: messages.select { |message| message.role == :system }.map(&:content).join("\n\n"),
             generation_config: config, tools: render_interaction_tools(tools)
           }
@@ -30,6 +31,15 @@ module RubyLLM
           payload
         end
 
+        def ensure_interaction_carries(messages, caching)
+          if caching.is_a?(Hash) && caching[:id]
+            raise ArgumentError, 'Gemini Interactions does not accept a cached content id'
+          end
+
+          attachment = messages.select { |message| message.role == :system }.flat_map(&:attachments).first
+          raise UnsupportedAttachmentError, attachment.mime_type if attachment
+        end
+
         def render_interaction_thinking(thinking)
           if thinking.disabled? || thinking.budget&.zero?
             raise ArgumentError, 'Gemini Interactions does not expose a thinking-off control'
@@ -50,8 +60,8 @@ module RubyLLM
 
         # Gemini checks signatures only in the current turn, which starts at
         # the last user message.
-        def format_interaction_input(messages)
-          calls = messages.flat_map { |message| message.tool_calls.to_h.values }.to_h { |call| [call.id, call] }
+        def render_interaction_input(messages)
+          calls = messages.flat_map { |message| message.tool_calls.to_a }.to_h { |call| [call.id, call] }
           turn_start = messages.rindex { |message| message.role == :user }
           messages.each_with_index.flat_map do |message, index|
             next [] if message.role == :system
@@ -62,7 +72,9 @@ module RubyLLM
 
         def render_interaction_turn(message, calls, current:)
           if interaction_state?(message.raw_content)
-            render_interaction_history(message.raw_content.dig('response', 'steps') || [])
+            steps = message.raw_content.dig('response', 'steps') || []
+            steps = Protocol::ReplayedArguments.blank(steps, message, id: 'id', arguments: 'arguments', empty: {})
+            render_interaction_history(steps)
           elsif message.tool_result?
             [render_interaction_result(message, calls)]
           else
@@ -71,7 +83,7 @@ module RubyLLM
         end
 
         def render_interaction_result(message, calls)
-          { type: 'function_result', call_id: message.tool_call_id, name: calls[message.tool_call_id]&.name,
+          { type: 'function_result', call_id: message.provider_call_id, name: calls[message.tool_call_id]&.name,
             result: render_interaction_content(message.content, message.attachments) }.compact
         end
 
@@ -90,9 +102,9 @@ module RubyLLM
           steps = []
           content = render_interaction_content(message.content, message.attachments)
           steps << { type: message.role == :assistant ? 'model_output' : 'user_input', content: } if content.any?
-          message.tool_calls.to_h.values.each_with_index do |call, position|
+          message.tool_calls.each_with_index do |call, position|
             signature = call.thought_signature || (Gemini::Tools::PLACEHOLDER_SIGNATURE if current && position.zero?)
-            steps << { type: 'function_call', id: call.id, name: call.name, arguments: call.arguments,
+            steps << { type: 'function_call', id: call.provider_call_id, name: call.name, arguments: call.arguments,
                        signature: }.compact
           end
           steps
@@ -102,16 +114,19 @@ module RubyLLM
           content.is_a?(Hash) && content.dig('response', 'object') == 'interaction'
         end
 
+        def raise_failed_interaction(data, raw)
+          message = Array(data['errors']).filter_map { |error| error['message'] }.join('; ')
+          raise Error.new("Gemini interaction ended with status #{data['status']}", response: raw) if message.empty?
+
+          raise Error.new(provider_message: message, response: raw, redact: @config&.redact_error_messages)
+        end
+
         def parse_completion_body(data, raw:, model: data['model'] || @model&.id, cost: nil)
-          unless %w[completed requires_action incomplete].include?(data['status'])
-            message = Array(data['errors']).filter_map { |error| error['message'] }.join('; ')
-            raise Error.new(message.empty? ? "Gemini interaction ended with status #{data['status']}" : message,
-                            response: raw)
-          end
+          raise_failed_interaction(data, raw) unless %w[completed requires_action incomplete].include?(data['status'])
 
           steps = data.fetch('steps', [])
           content = parse_interaction_content(steps)
-          calls = parse_interaction_calls(steps)
+          calls = parse_interaction_calls(steps, response: raw, finish_reason: finish_reasons[data['status']])
           status = interaction_status(data, calls, steps)
           if status == 'requires_action' && calls.empty?
             raise Error.new('Gemini interaction requires an unsupported action', response: raw)

data/lib/ruby_llm/protocols/interactions/content.rb

@@ -16,7 +16,7 @@ module RubyLLM
         end
 
         def supported_attachment?(attachment)
-          attachment.text? || !interaction_attachment_type(attachment).nil?
+          (attachment.text? && !attachment.provider_file?) || !interaction_attachment_type(attachment).nil?
         end
 
         def render_interaction_attachment(attachment)
@@ -24,18 +24,18 @@ module RubyLLM
             raise ArgumentError, "Gemini Interactions does not support #{attachment.mime_type} input"
           end
 
-          return { type: 'text', text: attachment.content } if attachment.text?
+          return { type: 'text', text: attachment.content } if attachment.text? && !attachment.provider_file?
 
-          type = interaction_attachment_type(attachment)
+          part = { type: interaction_attachment_type(attachment), mime_type: attachment.mime_type }
+          part[:processing] = attachment.processing.to_s if attachment.processing && attachment.video?
+          part.merge(interaction_attachment_source(attachment))
+        end
 
-          part = { type: type, mime_type: attachment.mime_type }
-          if attachment.provider_file?
-            part.merge(uri: attachment.provider_file_uri)
-          elsif attachment.url?
-            part.merge(uri: attachment.source.to_s)
-          else
-            part.merge(data: attachment.encoded)
-          end
+        def interaction_attachment_source(attachment)
+          return { uri: attachment.provider_file_uri } if attachment.provider_file?
+          return { uri: attachment.source.to_s } if attachment.url?
+
+          { data: attachment.encoded }
         end
 
         def interaction_attachment_type(attachment)
@@ -43,7 +43,7 @@ module RubyLLM
           return 'audio' if attachment.audio?
           return 'video' if attachment.video?
 
-          'document' if attachment.pdf?
+          'document' if attachment.pdf? || (attachment.provider_file? && attachment.mime_type == 'text/csv')
         end
 
         def parse_interaction_content(steps)
@@ -60,11 +60,16 @@ module RubyLLM
             result[:citations].concat(parse_interaction_citations(part, result[:text].length))
             result[:text] << text
           elsif part['data'] || part['uri']
-            source = part['uri'] || StringIO.new(Base64.decode64(part['data']))
-            result[:attachments] << Attachment.new(source, config: @config)
+            result[:attachments] << parse_interaction_attachment(part, result[:attachments].size)
           end
         end
 
+        def parse_interaction_attachment(part, index)
+          source = part['uri'] || StringIO.new(Base64.decode64(part['data']))
+          filename = Gemini::Media.attachment_filename(part['mime_type'], index) if part['mime_type']
+          Attachment.new(source, filename:, config: @config)
+        end
+
         def parse_interaction_citations(part, offset)
           Array(part['annotations']).filter_map do |annotation|
             next unless CITATION_TYPES.include?(annotation['type'])

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

@@ -80,7 +80,7 @@ module RubyLLM
           steps = @interaction_steps.sort.map do |_index, step|
             next step unless step['type'] == 'function_call'
 
-            step.merge('arguments' => parse_interaction_arguments(step['arguments']))
+            step.merge('arguments' => ToolCall.parse_arguments(step['arguments']).fetch(:arguments, step['arguments']))
           end
           @interaction_response.merge('steps' => steps)
         end

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

@@ -22,13 +22,18 @@ module RubyLLM
           { allowed_tools: { mode: 'any', tools: [choice.to_s] } }
         end
 
-        def parse_interaction_calls(steps)
+        def parse_interaction_calls(steps, response: nil, finish_reason: nil)
           answered = interaction_answered_calls(steps)
-          steps.select { |step| step['type'] == 'function_call' && !answered.include?(step['id']) }.to_h do |step|
-            [step.fetch('id'), ToolCall.new(id: step.fetch('id'), name: step.fetch('name'),
-                                            arguments: parse_interaction_arguments(step['arguments']),
-                                            thought_signature: step['signature'])]
+          pending = steps.select { |step| step['type'] == 'function_call' && !answered.include?(step['id']) }
+          cut_off = finish_reason == :max_tokens && pending.any? && [nil, '', {}].include?(pending.last['arguments'])
+          raise ToolCallTruncatedError.new(response:) if cut_off
+
+          calls = pending.map do |step|
+            ToolCall.new(provider_call_id: step.fetch('id'), name: step.fetch('name'),
+                         **ToolCall.parse_arguments(step['arguments'], response:, finish_reason:),
+                         thought_signature: step['signature'])
           end
+          ToolCall.index(calls, by: :provider_call_id)
         end
 
         def interaction_answered_calls(steps)
@@ -37,14 +42,6 @@ module RubyLLM
           calls & results
         end
 
-        def parse_interaction_arguments(arguments)
-          return {} if arguments.nil? || (arguments.is_a?(String) && arguments.empty?)
-
-          arguments.is_a?(String) ? JSON.parse(arguments) : arguments
-        rescue JSON::ParserError => e
-          raise ToolCallParseError.new(finish_reason: :tool_calls), cause: e
-        end
-
         def parse_interaction_server_calls(steps)
           answered = interaction_answered_calls(steps)
           steps.filter_map do |step|

data/lib/ruby_llm/protocols/interactions/videos.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Interactions
+      # Gemini Omni video generation. A background interaction is the only way
+      # Gemini lets a job be polled, and it requires store: true, so the job
+      # stays on Google's servers. Each job is a single turn, not a conversation.
+      module Videos # :nodoc:
+        JOB_ID = /\A[\w-]+\z/
+        FINISHED_WITHOUT_VIDEO = %w[failed cancelled incomplete budget_exceeded].freeze
+        private_constant :JOB_ID, :FINISHED_WITHOUT_VIDEO
+
+        def video_url
+          'interactions'
+        end
+
+        def render_video_payload(prompt, model:, with: [], provider_options: {}, **)
+          text = prompt.to_s.empty? ? [] : [{ type: 'text', text: prompt.to_s }]
+          payload = { model:, input: with.map { |attachment| render_interaction_attachment(attachment) } + text,
+                      background: true, store: true, response_format: { type: 'video' } }
+
+          Support::Utils.deep_merge(payload, provider_options)
+        end
+
+        # Gemini refuses uploaded sources in some regions but continues a
+        # video it generated everywhere, so a generated Video names its job.
+        def render_video_extension_payload(prompt, extend:, model:, provider_options: {}, **)
+          generated = extend.is_a?(Video) && extend.raw && generated_video(extend.raw)
+          unless generated
+            return render_video_payload(prompt, model:, with: [video_extension_attachment(extend)],
+                                                provider_options:)
+          end
+
+          render_video_payload(prompt, model:, provider_options:).merge(previous_interaction_id: extend.raw['id'])
+        end
+
+        def video_request_settings(request)
+          request[:response_format].slice(:duration, :resolution)
+        end
+
+        def parse_video_job(response, model:, **request)
+          id = response.body['id']
+          raise Error.new('Gemini did not return a video generation interaction', response:) unless id
+
+          VideoJob.new(id:, protocol: self, model:, **request, **parse_video_job_status(response))
+        end
+
+        def video_job_url(job)
+          return "interactions/#{job.id}" if job.id.match?(JOB_ID)
+
+          raise ArgumentError, "#{job.id.inspect} is not a video interaction id"
+        end
+
+        def parse_video_job_status(response, **)
+          body = response.body
+          return { status: :pending, raw: body } unless finished_video_job?(body)
+
+          tokens = Tokens.new(**parse_video_tokens(body))
+          return { status: :completed, raw: body, tokens: } if body['status'] == 'completed' && generated_video(body)
+
+          { status: :failed, raw: body, tokens:, error: video_job_error(body) }
+        end
+
+        def download_video(job)
+          video = generated_video(job.raw)
+
+          Video.new(data: Base64.decode64(video['data']), mime_type: video['mime_type'] || 'video/mp4',
+                    model: job.raw['model'] || job.model, raw: job.raw)
+        end
+
+        private
+
+        def validate_animate_inputs!(with:); end
+
+        def finished_video_job?(body)
+          body['status'] == 'completed' || FINISHED_WITHOUT_VIDEO.include?(body['status'])
+        end
+
+        def parse_video_tokens(body)
+          usage = parse_interaction_usage(body['usage'] || {})
+          { input: usage[:input_tokens], output: usage[:output_tokens], cache_read: usage[:cache_read_tokens],
+            thinking: usage[:thinking_tokens] }
+        end
+
+        def video_job_error(body)
+          message = Array(body['errors']).filter_map { |error| error['message'] }.join('; ')
+          message.empty? ? 'Gemini returned no video' : message
+        end
+
+        def generated_video(body)
+          Array(body['steps']).select { |step| step['type'] == 'model_output' }
+                              .flat_map { |step| Array(step['content']) }
+                              .find { |part| part['type'] == 'video' && part['data'] }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/interactions.rb

@@ -10,6 +10,7 @@ module RubyLLM
       include Interactions::Tools
       include Interactions::Streaming
       include Interactions::Transcription
+      include Interactions::Videos
 
       public :render
 

data/lib/ruby_llm/protocols/invoke_model/cohere_embeddings.rb

@@ -5,11 +5,15 @@ module RubyLLM
     class InvokeModel
       # Cohere embedding models over Bedrock InvokeModel.
       class CohereEmbeddings < InvokeModel
+        def rejected_options(operation)
+          operation == :embed ? %i[title] : super
+        end
+
         # rubocop:disable-next Lint/UnusedMethodArgument
         def embed(text, model:, dimensions:, task_type: nil, title: nil, with: nil, provider_options: {})
-          ensure_no_embedding_media!(with)
+          attachments = embedding_attachments(text, with)
           track_usage(:embedding) do
-            payload = render_embedding_payload(text, model:, dimensions:, task_type:, provider_options:)
+            payload = render_embedding_payload(text, model:, dimensions:, task_type:, attachments:, provider_options:)
             response = signed_post(embedding_url(model:), payload)
 
             parse_embedding_response(response, model:, text:)
@@ -18,23 +22,37 @@ module RubyLLM
 
         private
 
-        def render_embedding_payload(text, dimensions:, provider_options:, task_type: nil, **)
-          payload = {
-            input_type: task_type || 'search_document'
-          }
-          texts = [text].flatten.compact.map(&:to_s).reject(&:empty?)
-          payload[:texts] = texts unless texts.empty?
+        def render_embedding_payload(text, dimensions:, provider_options:, task_type: nil, attachments: [], **)
+          payload = { input_type: task_type || 'search_document' }
+          payload.merge!(attachments.any? ? media_embedding_inputs(text, attachments) : text_embedding_inputs(text))
           payload[:output_dimension] = dimensions if dimensions
 
           deep_merge_provider_options(payload, provider_options)
         end
 
+        def text_embedding_inputs(text)
+          texts = text.is_a?(Array) ? text.map(&:to_s) : [text.to_s].reject(&:empty?)
+          texts.empty? ? {} : { texts: }
+        end
+
+        def media_embedding_inputs(text, attachments)
+          content = Cohere::Media.render_content(text.to_s.empty? ? nil : text, attachments)
+          { inputs: [{ content: }], embedding_types: ['float'] }
+        end
+
         def parse_embedding_response(response, model:, text:)
           vectors = response.body['embeddings']
           vectors = vectors['float'] || vectors.values.first if vectors.is_a?(Hash)
-          vectors = vectors.first unless text.is_a?(Array)
+          vectors = text.is_a?(Array) ? aligned_vectors(vectors, text, response) : vectors.first
+
+          Embedding.new(vectors:, model:,
+                        input_tokens: response.headers['x-amzn-bedrock-input-token-count']&.to_i)
+        end
+
+        def aligned_vectors(vectors, text, response)
+          return vectors if text.empty? || vectors.size == text.size
 
-          Embedding.new(vectors:, model:)
+          raise Error.new("Expected #{text.size} embeddings, got #{vectors.size}", response:)
         end
       end
     end

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

@@ -102,7 +102,7 @@ module RubyLLM
         end
 
         def embedding_batch_vectors(bodies)
-          vectors = bodies.map { |body| extract_embedding(body) }
+          vectors = bodies.map { |body| parse_embedding(body) }
           vectors if vectors.all? { |vector| vector.is_a?(Array) && !vector.empty? && vector.all?(Numeric) }
         end
       end

data/lib/ruby_llm/protocols/invoke_model/nova_embeddings.rb

@@ -6,13 +6,19 @@ module RubyLLM
       # Amazon Nova multimodal embedding models over Bedrock InvokeModel.
       class NovaEmbeddings < InvokeModel
         DEFAULT_EMBEDDING_PURPOSE = 'GENERIC_INDEX'
+        INPUT_TOKENS_HEADER = 'x-amzn-bedrock-input-token-count'
+        MEDIA_TYPES = %i[image audio video].freeze
+
+        def rejected_options(operation)
+          operation == :embed ? %i[title] : super
+        end
 
         # rubocop:disable-next Lint/UnusedMethodArgument
         def embed(text, model:, dimensions:, task_type: nil, title: nil, with: nil, provider_options: {})
-          ensure_no_embedding_media!(with)
+          attachment = embedding_attachment(text, embedding_attachments(text, with))
           track_usage(:embedding) do
             responses = [text].flatten.map do |value|
-              payload = render_embedding_payload(value, dimensions:, task_type:, provider_options:)
+              payload = render_embedding_payload(value, dimensions:, task_type:, attachment:, provider_options:)
               signed_post(embedding_url(model:), payload).tap { |response| record_embedding_attempt(response) }
             end
 
@@ -22,14 +28,22 @@ module RubyLLM
 
         private
 
-        def render_embedding_payload(text, dimensions:, task_type:, provider_options:)
-          params = {
-            embeddingPurpose: task_type || DEFAULT_EMBEDDING_PURPOSE,
-            text: {
-              truncationMode: 'END',
-              value: text.to_s
-            }
-          }
+        def embedding_attachment(text, attachments)
+          attachment = single_embedding_attachment(attachments, 'Nova Multimodal Embeddings')
+          return unless attachment
+
+          unless text.to_s.empty?
+            raise ArgumentError,
+                  'Nova Multimodal Embeddings embeds text or an attachment, not both'
+          end
+          raise UnsupportedAttachmentError, attachment.mime_type unless MEDIA_TYPES.include?(attachment.type)
+
+          attachment
+        end
+
+        def render_embedding_payload(text, dimensions:, task_type:, provider_options:, attachment: nil)
+          params = { embeddingPurpose: task_type || DEFAULT_EMBEDDING_PURPOSE }
+          params.merge!(attachment ? render_embedding_media(attachment) : render_embedding_text(text))
           params[:embeddingDimension] = dimensions if dimensions
 
           deep_merge_provider_options(
@@ -41,9 +55,29 @@ module RubyLLM
           )
         end
 
-        def extract_embedding(body)
+        def render_embedding_text(text)
+          { text: { truncationMode: 'END', value: text.to_s } }
+        end
+
+        def render_embedding_media(attachment)
+          source = { source: { bytes: attachment.encoded } }
+          case attachment.type
+          when :video
+            format = attachment.extension || attachment.format
+            { video: { format:, embeddingMode: 'AUDIO_VIDEO_COMBINED' }.merge(source) }
+          else
+            { attachment.type => { format: attachment.format }.merge(source) }
+          end
+        end
+
+        def parse_embedding(body)
           body.dig('embeddings', 0, 'embedding')
         end
+
+        # Nova leaves the count out of its body and sends it as a header.
+        def embedding_attempt_tokens(responses)
+          responses.map { |response| Tokens.new(input: response.headers[INPUT_TOKENS_HEADER]&.to_i) }
+        end
       end
     end
   end

data/lib/ruby_llm/protocols/invoke_model/stability_images.rb

@@ -12,6 +12,10 @@ module RubyLLM
         MODELS = (GENERATION_MODELS + [INPAINT_MODEL]).freeze
         ASPECT_RATIOS = %w[16:9 1:1 21:9 2:3 3:2 4:5 5:4 9:16 9:21].freeze
 
+        def rejected_options(operation)
+          operation == :paint ? %i[count] : super
+        end
+
         def paint(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {})
           track_usage(:image) do
             payload = render_image_payload(prompt, model:, size:, count:, with:, mask:, provider_options:)
@@ -22,8 +26,7 @@ module RubyLLM
           end
         end
 
-        def render_image_payload(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {})
-          RubyLLM.logger.debug { 'Stability image models return one image per request' } if count && count > 1
+        def render_image_payload(prompt, model:, size:, with: nil, mask: nil, provider_options: {}, **)
           attachments = Attachment.wrap(with, config: @config)
           payload = { prompt: }.merge(render_image_inputs(attachments, model:, mask:, size:))
           payload.merge(provider_options.transform_keys(&:to_sym))