ruby_llm 2.0.0 → 2.1.0

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

@@ -10,6 +10,10 @@ module RubyLLM
       module Tools
         MULTIMODAL_FUNCTION_RESPONSE_GENERATION = Gem::Version.new('3')
 
+        # The signature Google documents for a function call Gemini did not
+        # make, such as one from another provider. Vertex AI accepts no other.
+        PLACEHOLDER_SIGNATURE = 'skip_thought_signature_validator'
+
         def format_tools(tools)
           return [] if tools.empty?
 
@@ -18,14 +22,13 @@ module RubyLLM
           }]
         end
 
-        def format_tool_call(msg) # rubocop:disable Metrics/PerceivedComplexity
+        # 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)
           parts = []
 
           parts.concat(Media.format_content(msg.content, msg.attachments)) if msg.content && !msg.content.empty?
 
-          fallback_signature = msg.thinking&.signature
-          used_fallback = false
-
           msg.tool_calls.each_value do |tool_call|
             part = {
               functionCall: {
@@ -33,13 +36,7 @@ module RubyLLM
                 args: tool_call.arguments
               }
             }
-
-            signature = tool_call.thought_signature
-            if signature.nil? && fallback_signature && !used_fallback
-              signature = fallback_signature
-              used_fallback = true
-            end
-            part[:thoughtSignature] = signature if signature
+            part[:thoughtSignature] = tool_call.thought_signature if tool_call.thought_signature
             parts << part
           end
 
@@ -95,6 +92,15 @@ module RubyLLM
 
         private
 
+        # Gemini 3 refuses a step of the current turn whose first function
+        # call carries no signature.
+        def sign_step(parts)
+          first_call = parts.find { |part| part.key?(:functionCall) }
+          return parts if first_call.nil? || first_call.key?(:thoughtSignature)
+
+          parts.map { |part| part.equal?(first_call) ? part.merge(thoughtSignature: PLACEHOLDER_SIGNATURE) : part }
+        end
+
         # functionResponse.parts only accepts inline bytes, and pre-Gemini 3 models reject it
         def partition_tool_result_attachments(attachments)
           return [[], []] if attachments.empty?

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

@@ -75,6 +75,7 @@ module RubyLLM
         def parse_transcription_response(response, model:)
           data = response.body
           text = extract_text(data)
+          raise_if_transcription_blocked(data, text, response)
 
           usage = extract_usage(data)
 
@@ -86,6 +87,18 @@ module RubyLLM
           )
         end
 
+        # Gemini returns no candidates only for a prompt it blocked.
+        def raise_if_transcription_blocked(data, text, response)
+          return unless text.to_s.empty? && data.is_a?(Hash)
+
+          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_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)
           candidate = data.is_a?(Hash) ? data.dig('candidates', 0) : nil
           return unless candidate

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

@@ -18,11 +18,16 @@ module RubyLLM
           Support::Utils.deep_merge({ instances: [instance] }, provider_options)
         end
 
-        def parse_video_job(response, model:)
+        def parse_video_job(response, model:, **request)
           name = response.body['name']
           raise Error.new('Gemini did not return a video generation operation', response:) unless name
 
-          VideoJob.new(id: name, protocol: self, model: model, raw: response.body)
+          VideoJob.new(id: name, protocol: self, model: model, raw: response.body, **request)
+        end
+
+        def video_request_settings(request)
+          parameters = request[:parameters] || {}
+          { duration: parameters[:durationSeconds], resolution: parameters[:resolution] }
         end
 
         def render_video_extension_payload(prompt, extend:, provider_options: {}, **)

data/lib/ruby_llm/protocols/gemini.rb

@@ -11,6 +11,7 @@ module RubyLLM
       include Gemini::Videos
       include Gemini::Media
       include Gemini::Models
+      include Gemini::RequestShapes
       include Gemini::Streaming
       include Gemini::Tools
       include Gemini::Speech

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

@@ -10,8 +10,6 @@ module RubyLLM
         end
 
         def render_video_payload(prompt, model:, with: [], provider_options: {})
-          raise ArgumentError, 'vLLM-Omni video generation requires a prompt' if prompt.nil?
-
           options = video_options(provider_options)
           payload = { model: model, prompt: prompt }.merge(video_references(with)).merge(options).compact
           payload.transform_values do |value|
@@ -25,12 +23,16 @@ module RubyLLM
           end
         end
 
-        def parse_video_job(response, model:)
+        def parse_video_job(response, model:, **request)
           body = response.body
           id = body['id']
           raise Error.new('GPUStack did not return a video job id', response:) unless id
 
-          VideoJob.new(id: id, protocol: self, model: body['model'] || model, **video_job_state(body))
+          VideoJob.new(id: id, protocol: self, model: body['model'] || model, **video_job_state(body), **request)
+        end
+
+        def video_request_settings(request)
+          { duration: request[:seconds], resolution: request[:size] }
         end
 
         def video_job_url(job)
@@ -60,7 +62,7 @@ module RubyLLM
 
         def video_references(attachments)
           attachments.group_by(&:type).to_h do |type, group|
-            values = group.map { |attachment| { "#{type}_url" => video_reference(attachment) } }
+            values = group.map { |attachment| { "#{type}_url" => attachment.for_llm } }
             [:"#{type}_reference", values.one? ? values.first : values]
           end
         end
@@ -86,10 +88,6 @@ module RubyLLM
             raise UnsupportedAttachmentError, attachment.mime_type
           end
         end
-
-        def video_reference(attachment)
-          attachment.url? ? attachment.source.to_s : attachment.for_llm
-        end
       end
     end
   end

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

@@ -5,6 +5,7 @@ module RubyLLM
     class Interactions
       module Chat # :nodoc:
         FINISH_REASONS = { 'completed' => :stop, 'requires_action' => :tool_calls, 'incomplete' => :max_tokens }.freeze
+        SERVER_TOOL_USAGE_NAMES = { 'google_search' => 'web_search' }.freeze
 
         module_function
 
@@ -34,9 +35,6 @@ module RubyLLM
             raise ArgumentError, 'Gemini Interactions does not expose a thinking-off control'
           end
           raise ArgumentError, 'Gemini Interactions accepts thinking effort, not a token budget' if thinking.budget
-          if thinking.effort && !%i[minimal low medium high].include?(thinking.effort)
-            raise ArgumentError, 'Gemini Interactions thinking effort must be minimal, low, medium, or high'
-          end
 
           { thinking_level: thinking.effort&.to_s,
             thinking_summaries: render_interaction_summaries(thinking.display) }.compact
@@ -50,20 +48,25 @@ module RubyLLM
           display == :omitted ? 'none' : 'auto'
         end
 
+        # 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] }
-          messages.reject { |message| message.role == :system }.flat_map do |message|
-            render_interaction_turn(message, calls)
+          turn_start = messages.rindex { |message| message.role == :user }
+          messages.each_with_index.flat_map do |message, index|
+            next [] if message.role == :system
+
+            render_interaction_turn(message, calls, current: turn_start.nil? || index > turn_start)
           end
         end
 
-        def render_interaction_turn(message, calls)
+        def render_interaction_turn(message, calls, current:)
           if interaction_state?(message.raw_content)
             render_interaction_history(message.raw_content.dig('response', 'steps') || [])
           elsif message.tool_result?
             [render_interaction_result(message, calls)]
           else
-            render_interaction_message(message)
+            render_interaction_message(message, current:)
           end
         end
 
@@ -73,19 +76,24 @@ module RubyLLM
         end
 
         def render_interaction_history(steps)
+          answered = interaction_answered_calls(steps)
           steps.map do |step|
-            step['type'].to_s.start_with?('mcp_server_') ? step.except('signature') : step
+            if step['type'].to_s.start_with?('mcp_server_') || answered.include?(step['id'] || step['call_id'])
+              step.except('signature')
+            else
+              step
+            end
           end
         end
 
-        def render_interaction_message(message)
+        def render_interaction_message(message, current:)
           steps = []
-          if message.content || message.attachments.any?
-            steps << { type: message.role == :assistant ? 'model_output' : 'user_input',
-                       content: render_interaction_content(message.content, message.attachments) }
-          end
-          message.tool_calls&.each_value do |call|
-            steps << { type: 'function_call', id: call.id, name: call.name, arguments: call.arguments }
+          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|
+            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,
+                       signature: }.compact
           end
           steps
         end
@@ -104,19 +112,39 @@ module RubyLLM
           steps = data.fetch('steps', [])
           content = parse_interaction_content(steps)
           calls = parse_interaction_calls(steps)
-          if data['status'] == 'requires_action' && calls.empty?
+          status = interaction_status(data, calls, steps)
+          if status == 'requires_action' && calls.empty?
             raise Error.new('Gemini interaction requires an unsupported action', response: raw)
           end
 
           Message.new(role: :assistant, content: content[:text], attachments: content[:attachments],
                       citations: content[:citations], thinking: parse_interaction_thinking(steps),
                       tool_calls: calls, server_tool_calls: parse_interaction_server_calls(steps),
-                      raw_content: { 'response' => data },
+                      raw_content: kept_interaction(data, steps),
                       model: model, raw: raw, cost: cost,
-                      finish_reason: interaction_finish_reason(data['status'], calls),
+                      finish_reason: interaction_finish_reason(status, calls),
+                      server_tool_use: parse_interaction_server_tool_use(data['usage'] || {}),
                       **parse_interaction_usage(data['usage'] || {}))
         end
 
+        def interaction_status(data, calls, steps)
+          return data['status'] unless data['status'] == 'requires_action' && calls.empty?
+          return data['status'] if interaction_answered_calls(steps).empty?
+
+          steps.any? { |step| step['type'] == 'model_output' } ? 'completed' : data['status']
+        end
+
+        def parse_interaction_server_tool_use(usage)
+          Array(usage['grounding_tool_count']).to_h do |grounding|
+            tool = grounding['type']
+            ["#{SERVER_TOOL_USAGE_NAMES.fetch(tool, tool)}_requests", grounding['count']]
+          end
+        end
+
+        def kept_interaction(data, steps)
+          { 'response' => data.merge('steps' => steps.map { |step| without_search_suggestions(step) }) }
+        end
+
         def interaction_finish_reason(status, calls)
           return :tool_calls unless calls.empty?
 

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

@@ -11,15 +11,22 @@ module RubyLLM
         module_function
 
         def render_interaction_content(text, attachments)
-          parts = text.nil? ? [] : [{ type: 'text', text: text.to_s }]
+          parts = text.to_s.empty? ? [] : [{ type: 'text', text: text.to_s }]
           parts + attachments.map { |attachment| render_interaction_attachment(attachment) }
         end
 
+        def supported_attachment?(attachment)
+          attachment.text? || !interaction_attachment_type(attachment).nil?
+        end
+
         def render_interaction_attachment(attachment)
+          unless supported_attachment?(attachment)
+            raise ArgumentError, "Gemini Interactions does not support #{attachment.mime_type} input"
+          end
+
           return { type: 'text', text: attachment.content } if attachment.text?
 
           type = interaction_attachment_type(attachment)
-          raise ArgumentError, "Gemini Interactions does not support #{attachment.mime_type} input" unless type
 
           part = { type: type, mime_type: attachment.mime_type }
           if attachment.provider_file?

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

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Interactions
+      # Describes Interactions requests as RequestShape objects. Each input
+      # step is a turn. A function result answers its call by id, and like
+      # Gemini, the first call of each step in the current turn needs a
+      # signature.
+      module RequestShapes
+        SPEAKERS = { 'user_input' => :user, 'function_result' => :tool }.freeze
+        MEDIA_KINDS = { 'image' => :image, 'audio' => :audio, 'video' => :video, 'document' => :document }.freeze
+
+        def parse_request_shape(payload)
+          input = payload['input']
+          return unless input.is_a?(Array)
+
+          turns = input.each_with_index.map do |entry, index|
+            entry = shape_hash(entry)
+            shape_turn(index, entry['type'], SPEAKERS.fetch(entry['type'], :model), shape_interaction_entry(entry))
+          end
+          shape_request(turns, payload:, rules: %i[signed_steps],
+                               instructions: shape_instructions(payload['system_instruction']),
+                               tool_names: shape_list(payload['tools']).map { |tool| shape_tool_name(tool) },
+                               thinking_settings: shape_settings(payload['generation_config'], 'thinking_level',
+                                                                 'thinking_summaries'))
+        end
+
+        private
+
+        def shape_interaction_entry(entry)
+          signed = shape_signature?(entry['signature'])
+          case entry['type']
+          when 'user_input', 'model_output' then shape_interaction_content(entry['content'])
+          when 'function_call' then [shape_tool_call(entry['name'], entry['arguments'], id: entry['id'], signed:)]
+          when 'function_result'
+            shape_tool_result(entry['result'], name: entry['name'], id: entry['call_id']) do |part|
+              shape_interaction_part(part)
+            end
+          when 'thought' then [shape_summary(entry['summary'], signed:)]
+          else [shape_other(entry['type'], signed:)]
+          end
+        end
+
+        def shape_interaction_content(content)
+          return [shape_text(content)] if content.is_a?(String)
+
+          shape_list(content).map { |part| shape_interaction_part(shape_hash(part)) }
+        end
+
+        def shape_interaction_part(part)
+          kind = MEDIA_KINDS[part['type']]
+          return shape_text(part['text']) if part['type'] == 'text'
+          return shape_other(part['type']) unless kind
+          return shape_inline(part['data'], kind:, mime_type: part['mime_type']) if part.key?('data')
+
+          shape_url(part['uri'], kind:, mime_type: part['mime_type'])
+        end
+      end
+    end
+  end
+end

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

@@ -4,6 +4,13 @@ module RubyLLM
   module Protocols
     class Interactions
       module Streaming # :nodoc:
+        ERROR_STATUSES = {
+          'invalid_request' => 400, 'failed_precondition' => 400, 'authentication' => 401,
+          'payment_required' => 402, 'permission_denied' => 403, 'not_found' => 404,
+          'rate_limit_exceeded' => 429, 'quota_exceeded' => 429, 'too_many_requests' => 429,
+          'api_error' => 500, 'service_unavailable' => 503, 'deadline_exceeded' => 504
+        }.freeze
+
         module_function
 
         def stream_response(payload, additional_headers = {})
@@ -78,6 +85,14 @@ module RubyLLM
           @interaction_response.merge('steps' => steps)
         end
 
+        def parse_streaming_error(data)
+          error = JSON.parse(data)
+          error = error['error'] if error.is_a?(Hash)
+          return super unless error.is_a?(Hash)
+
+          [ERROR_STATUSES.fetch(error['code'], 500), error['message']]
+        end
+
         def final_interaction_chunk
           message = parse_completion_body(streamed_interaction, raw: nil)
           Chunk.new(role: :assistant, content: nil, model: message.model, tokens: message.tokens,

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

@@ -23,13 +23,20 @@ module RubyLLM
         end
 
         def parse_interaction_calls(steps)
-          steps.select { |step| step['type'] == 'function_call' }.to_h do |step|
+          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'])]
           end
         end
 
+        def interaction_answered_calls(steps)
+          calls = steps.filter_map { |step| step['id'] if step['type'] == 'function_call' }
+          results = steps.filter_map { |step| step['call_id'] if step['type'] == 'function_result' }
+          calls & results
+        end
+
         def parse_interaction_arguments(arguments)
           return {} if arguments.nil? || (arguments.is_a?(String) && arguments.empty?)
 
@@ -39,14 +46,31 @@ module RubyLLM
         end
 
         def parse_interaction_server_calls(steps)
+          answered = interaction_answered_calls(steps)
           steps.filter_map do |step|
             type = step['type'].to_s
-            next unless type.end_with?('_call', '_result') && !type.start_with?('function_')
+            next unless type.end_with?('_call', '_result')
+            next if type.start_with?('function_') && !answered.include?(step['id'] || step['call_id'])
 
-            ServerToolCall.new(type: type, id: step['id'] || step['call_id'], name: step['name'],
-                               input: step['arguments'], result: step['result'], raw: step)
+            kept = without_search_suggestions(step)
+            ServerToolCall.new(type: type, id: kept['id'] || kept['call_id'], name: kept['name'],
+                               input: kept['arguments'], result: kept['result'], raw: kept,
+                               search_suggestions: search_suggestions(step))
           end
         end
+
+        # Google's terms forbid storing search suggestions, and a replayed
+        # google_search_result keeps the shape it needs without them.
+        def without_search_suggestions(step)
+          return step unless step['type'] == 'google_search_result'
+
+          step.merge('result' => Array(step['result']).map { |item| item.except('search_suggestions') })
+        end
+
+        def search_suggestions(step)
+          suggestions = Array(step['result']).filter_map { |item| item['search_suggestions'] if item.is_a?(Hash) }
+          suggestions.join unless suggestions.empty?
+        end
       end
     end
   end

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

@@ -24,17 +24,7 @@ module RubyLLM
           payload = { model:, store: false,
                       input: [{ type: 'audio', mime_type: attachment.mime_type, data: attachment.encoded }],
                       generation_config: { transcription_config: config } }
-          payload = Support::Utils.deep_merge(payload, provider_options)
-          validate_transcription_config(payload.dig(:generation_config, :transcription_config))
-          payload
-        end
-
-        def validate_transcription_config(config)
-          mode = config[:mode]
-          return unless config[:custom_vocabulary] && mode.is_a?(Hash)
-          return unless mode[:diarization_mode] || mode[:timestamp_granularities]
-
-          raise ArgumentError, 'Gemini custom vocabulary cannot be combined with diarization or word timestamps'
+          Support::Utils.deep_merge(payload, provider_options)
         end
 
         def parse_transcription_response(response, model:)

data/lib/ruby_llm/protocols/interactions.rb

@@ -6,6 +6,7 @@ module RubyLLM
     class Interactions < Protocol
       include Interactions::Chat
       include Interactions::Content
+      include Interactions::RequestShapes
       include Interactions::Tools
       include Interactions::Streaming
       include Interactions::Transcription

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

@@ -18,18 +18,13 @@ module RubyLLM
 
         private
 
-        def render_embedding_payload(text, model:, dimensions:, provider_options:, task_type: nil)
+        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?
-
-          if dimensions
-            raise Error, "#{model} does not support custom dimensions" unless cohere_v4?(model)
-
-            payload[:output_dimension] = dimensions
-          end
+          payload[:output_dimension] = dimensions if dimensions
 
           deep_merge_provider_options(payload, provider_options)
         end
@@ -41,10 +36,6 @@ module RubyLLM
 
           Embedding.new(vectors:, model:)
         end
-
-        def cohere_v4?(model)
-          model.to_s.include?('cohere.embed-v4')
-        end
       end
     end
   end

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

@@ -17,8 +17,8 @@ module RubyLLM
           requests.each do |request|
             values = request.fetch(:text)
             values = [values] unless values.is_a?(Array)
-            unless values.any? && values.all? { |value| value.is_a?(String) && !value.empty? }
-              raise ArgumentError, 'Bedrock embedding batches require nonempty text strings'
+            unless values.any? && values.all?(String)
+              raise ArgumentError, 'Bedrock embedding batches require a nonempty list of text strings'
             end
             unless request.fetch(:payload).key?(:inputText)
               raise ArgumentError, 'Bedrock embedding batches require Titan embedding requests'

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

@@ -11,7 +11,6 @@ module RubyLLM
         INPAINT_MODEL = 'stability.stable-image-inpaint-v1:0'
         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
-        IMAGE_TYPES = %w[image/jpeg image/png image/webp].freeze
 
         def paint(prompt, model:, size:, count: nil, with: nil, mask: nil, provider_options: {})
           track_usage(:image) do
@@ -39,10 +38,7 @@ module RubyLLM
           end
 
           images.map do |data|
-            mime_type = RubyLLM::Files::MimeType.for(StringIO.new(Base64.decode64(data)))
-            raise Error.new('Bedrock returned invalid image data', response:) unless IMAGE_TYPES.include?(mime_type)
-
-            Image.new(data:, model:, mime_type:)
+            Image.new(data:, model:, mime_type: RubyLLM::Files::MimeType.for(StringIO.new(Base64.decode64(data))))
           end
         end
 
@@ -71,7 +67,7 @@ module RubyLLM
         end
 
         def encoded_image(attachment)
-          raise UnsupportedAttachmentError, attachment.mime_type unless IMAGE_TYPES.include?(attachment.mime_type)
+          raise UnsupportedAttachmentError, attachment.mime_type unless attachment.image?
 
           attachment.encoded
         end

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

@@ -23,14 +23,8 @@ module RubyLLM
         # The G1 and V1 models take inputText alone; Bedrock rejects the V2
         # tuning keys as extraneous.
         def render_embedding_payload(text, model:, dimensions:, provider_options:, **)
-          payload = { inputText: text.to_s }
-
-          if titan_v2?(model)
-            payload[:dimensions] = dimensions if dimensions
-            payload[:normalize] = true
-          elsif dimensions
-            raise Error, "#{model} does not support custom dimensions"
-          end
+          payload = { inputText: text.to_s, dimensions: }.compact
+          payload[:normalize] = true if titan_v2?(model)
 
           deep_merge_provider_options(payload, provider_options)
         end

data/lib/ruby_llm/protocols/invoke_model.rb

@@ -10,10 +10,7 @@ module RubyLLM
       private
 
       def signed_post(url, payload, additional_headers = {})
-        body = JSON.generate(payload)
-
-        @connection.post(url, payload, usage: @usage_tracker) do |req|
-          req.headers.merge!(@provider.sign_headers('POST', url, body))
+        @connection.post(url, JSON.generate(payload), usage: @usage_tracker) do |req|
           req.headers.merge!(additional_headers) unless additional_headers.empty?
         end
       end

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

@@ -11,10 +11,23 @@ module RubyLLM
           output.select { |entry| entry['type'] == 'message.output' }.each do |entry|
             parse_conversation_parts(entry['content'], result)
           end
+          result[:attachments].concat(output.filter_map do |entry|
+            next unless entry['type'] == 'tool.execution' && entry['name'] == 'image_generation'
+
+            parse_generated_image(entry.dig('info', 'result'))
+          end)
+          result[:attachments].uniq!(&:source)
           result[:thinking] = nil if result[:thinking].empty?
           result
         end
 
+        def parse_generated_image(result)
+          return unless result
+
+          url = JSON.parse(result)['url']
+          Attachment.new(url, config: @config) if url
+        end
+
         def parse_conversation_parts(content, result)
           return result[:text] << content if content.is_a?(String)
 

data/lib/ruby_llm/protocols/mistral/conversations/chat.rb

@@ -5,6 +5,10 @@ module RubyLLM
     module Mistral
       class Conversations
         module Chat # :nodoc:
+          SERVER_TOOL_USAGE_NAMES = {
+            'code_interpreter' => 'code_execution', 'document_library' => 'file_search'
+          }.freeze
+
           module_function
 
           def completion_url
@@ -150,9 +154,15 @@ module RubyLLM
             input = usage['prompt_tokens'] && (usage['prompt_tokens'] + usage.fetch('connector_tokens', 0).to_i)
             {
               input_tokens: input,
-              output_tokens: usage['completion_tokens'], server_tool_use: usage['connectors']
+              output_tokens: usage['completion_tokens'], server_tool_use: parse_connector_usage(usage['connectors'])
             }
           end
+
+          def parse_connector_usage(connectors)
+            connectors.to_h.transform_keys do |connector|
+              "#{SERVER_TOOL_USAGE_NAMES.fetch(connector, connector)}_requests"
+            end
+          end
         end
       end
     end