Image to image with gemini-2.0-flash-preview-image-generation

lib/ruby_llm/message.rb

@@ -7,7 +7,7 @@ module RubyLLM
   class Message
     ROLES = %i[system user assistant tool].freeze
 
-    attr_reader :role, :tool_calls, :tool_call_id, :input_tokens, :output_tokens, :model_id
+    attr_reader :role, :tool_calls, :tool_call_id, :input_tokens, :output_tokens, :model_id, :images
 
     def initialize(options = {})
       @role = options.fetch(:role).to_sym
@@ -17,6 +17,7 @@ def initialize(options = {})
       @output_tokens = options[:output_tokens]
       @model_id = options[:model_id]
       @tool_call_id = options[:tool_call_id]
+      @images = options[:images]
 
       ensure_valid_role
     end
@@ -49,7 +50,8 @@ def to_h
         tool_call_id: tool_call_id,
         input_tokens: input_tokens,
         output_tokens: output_tokens,
-        model_id: model_id
+        model_id: model_id,
+        images: images
       }.compact
     end
 

lib/ruby_llm/providers/gemini/capabilities.rb

@@ -280,6 +280,9 @@ def modalities_for(model_id)
           # Embedding output
           modalities[:output] << 'embeddings' if model_id.match?(/embedding|gemini-embedding/)
 
+          # Image output
+          modalities[:output] << 'image' if model_id.match?(/image-generation/)
+
           modalities
         end
 

lib/ruby_llm/providers/gemini/chat.rb

@@ -16,7 +16,8 @@ def render_payload(messages, tools:, temperature:, model:, stream: false) # rubo
           payload = {
             contents: format_messages(messages),
             generationConfig: {
-              temperature: temperature
+              temperature: temperature,
+              responseModalities: capabilities.modalities_for(model)[:output]
             }
           }
           payload[:tools] = format_tools(tools) if tools.any?
@@ -72,6 +73,7 @@ def parse_completion_response(response)
           Message.new(
             role: :assistant,
             content: extract_content(data),
+            images: extract_images(data),
             tool_calls: tool_calls,
             input_tokens: data.dig('usageMetadata', 'promptTokenCount'),
             output_tokens: data.dig('usageMetadata', 'candidatesTokenCount'),
@@ -94,6 +96,27 @@ def extract_content(data)
           text_parts.map { |p| p['text'] }.join
         end
 
+        def extract_images(data)
+          candidate = data.dig('candidates', 0)
+          return '' unless candidate
+
+          # Content will be empty for function calls
+          return '' if function_call?(candidate)
+
+          # Extract image content
+          parts = candidate.dig('content', 'parts')
+          image_parts = parts&.select { |p| p['inlineData'] }
+          return '' unless image_parts&.any?
+
+          image_parts.map do |p|
+            Image.new(
+              data: p['inlineData']['data'],
+              mime_type: p['inlineData']['mimeType'],
+              model_id: data['modelVersion']
+            )
+          end
+        end
+
         def function_call?(candidate)
           parts = candidate.dig('content', 'parts')
           parts&.any? { |p| p['functionCall'] }