ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/providers/mistral/chat_completions/batches.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Mistral
+      class ChatCompletions
+        # Mistral batch jobs for chat completions and embeddings.
+        module Batches
+          include RubyLLM::Batch::Helpers
+
+          TERMINAL_STATUSES = %w[SUCCESS FAILED TIMEOUT_EXCEEDED CANCELLED].freeze
+          Response = Struct.new(:body)
+          private_constant :TERMINAL_STATUSES, :Response
+
+          def create_batch(requests)
+            model = single_batch_model!(requests, 'mistral')
+            response = @connection.post('batch/jobs', {
+                                          endpoint: mistral_batch_endpoint(requests),
+                                          model: model,
+                                          requests: requests.map { |request| mistral_batch_request(request) }
+                                        }, idempotent: false)
+
+            parse_batch_response(response.body)
+          end
+
+          def find_batch(id)
+            attempts = 0
+            begin
+              parse_batch_response @connection.get(batch_url(id)).body
+            rescue Error => e
+              attempts += 1
+              raise unless e.response&.status == 404 && attempts < 3
+
+              sleep(0.5 * attempts)
+              retry
+            end
+          end
+
+          def cancel_batch(id)
+            parse_batch_response @connection.post("#{batch_url(id)}/cancel", {}).body
+          end
+
+          def batch_results(id)
+            response = @connection.get(batch_url(id)) { |request| request.params[:inline] = true }
+            Array(response.body['outputs']).filter_map { |line| parse_batch_result(line) }
+          end
+
+          private
+
+          def batch_url(id)
+            "batch/jobs/#{id}"
+          end
+
+          def mistral_batch_request(request)
+            body = batch_payload(request, except: :model)
+            custom_id = request[:custom_id]
+            custom_id = "#{custom_id}:array" if (body[:input] || body['input']).is_a?(Array)
+
+            {
+              custom_id: custom_id,
+              body: body
+            }
+          end
+
+          def mistral_batch_endpoint(requests)
+            endpoints = requests.map do |request|
+              payload = request.fetch(:payload)
+              payload.key?(:input) || payload.key?('input') ? '/v1/embeddings' : '/v1/chat/completions'
+            end.uniq
+            return endpoints.first if endpoints.one?
+
+            raise Error, 'Mistral batches cannot mix chat and embedding requests'
+          end
+
+          def parse_batch_response(data)
+            {
+              id: data['id'],
+              raw_status: data['status'],
+              completed: TERMINAL_STATUSES.include?(data['status']),
+              request_count: data['total_requests'],
+              request_counts: {
+                'total' => data['total_requests'],
+                'completed' => data['completed_requests'],
+                'succeeded' => data['succeeded_requests'],
+                'failed' => data['failed_requests']
+              }.compact
+            }
+          end
+
+          def parse_batch_status(raw_status, completed:)
+            return :pending unless completed
+            return :succeeded if raw_status == 'SUCCESS'
+            return :cancelled if raw_status == 'CANCELLED'
+
+            :failed
+          end
+
+          def parse_batch_result(line)
+            custom_id, shape = line['custom_id'].split(':', 2)
+            index = batch_result_index(custom_id)
+            response = line['response']
+
+            if response && response['body']
+              body = response['body']
+              [index, parse_mistral_batch_body(body, shape:)]
+            else
+              [index, nil, batch_failure(line['custom_id'], batch_error_message(line))]
+            end
+          end
+
+          def parse_mistral_batch_body(body, shape:)
+            return parse_completion_body(body, raw: body) unless body['data'].is_a?(Array)
+
+            parse_embedding_response(Response.new(body), model: body['model'], text: shape == 'array' ? [] : nil)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/mistral/chat_completions.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Mistral
+      # Mistral's dialect of the Chat Completions API.
+      class ChatCompletions < Protocols::ChatCompletions
+        include Mistral::Chat
+        include Mistral::Embeddings
+        include Mistral::Media
+        include Mistral::Models
+        include Mistral::OCR
+        include Mistral::Speech
+        include Mistral::Transcription
+        include Protocols::Mistral::MultiCompletion
+
+        public :server_tool_aliases
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/mistral/conversations.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Mistral
+      # Mistral Conversations with its document and image attachment formats.
+      class Conversations < Protocols::Mistral::Conversations
+        include Mistral::Media
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/mistral/embeddings.rb

@@ -11,16 +11,18 @@ module RubyLLM
           'embeddings'
         end
 
-        def render_embedding_payload(text, model:, dimensions:) # rubocop:disable Lint/UnusedMethodArgument
+        # rubocop:disable-next Lint/UnusedMethodArgument
+        def render_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil, provider_options: {})
           {
             model: model,
-            input: text
-          }
+            input: text,
+            output_dimension: dimensions
+          }.compact.merge(provider_options)
         end
 
         def parse_embedding_response(response, model:, text:)
           data = response.body
-          input_tokens = data.dig('usage', 'prompt_tokens') || 0
+          input_tokens = data.dig('usage', 'prompt_tokens')
           vectors = data['data'].map { |d| d['embedding'] }
 
           vectors = vectors.first if vectors.length == 1 && !text.is_a?(Array)

data/lib/ruby_llm/providers/mistral/media.rb

@@ -7,33 +7,21 @@ module RubyLLM
       module Media
         module_function
 
-        def format_content(content) # rubocop:disable Metrics/PerceivedComplexity
-          if content.is_a?(RubyLLM::Content::Raw)
-            value = content.value
-            return value.is_a?(Hash) ? value.to_json : value
-          end
-          return content.to_json if content.is_a?(Hash) || content.is_a?(Array)
-          return content unless content.is_a?(Content)
-
-          parts = []
-          parts << OpenAI::Media.format_text(content.text) if content.text
-
-          content.attachments.each do |attachment|
+        def format_content(content, attachments = [])
+          Protocols::ChatCompletions::Media.format_parts(content, attachments) do |attachment|
             case attachment.type
             when :image
-              parts << format_image(attachment)
+              format_image(attachment)
             when :audio
-              parts << OpenAI::Media.format_audio(attachment)
+              Protocols::ChatCompletions::Media.format_audio(attachment)
             when :pdf, :document
-              parts << format_document(attachment)
+              format_document(attachment)
             when :text
-              parts << OpenAI::Media.format_text_file(attachment)
+              Protocols::ChatCompletions::Media.format_text_file(attachment)
             else
               raise UnsupportedAttachmentError, attachment.mime_type
             end
           end
-
-          parts
         end
 
         def format_image(image)

data/lib/ruby_llm/providers/mistral/models.rb

@@ -1,49 +1,83 @@
 # frozen_string_literal: true
 
-require 'time'
-
 module RubyLLM
   module Providers
     class Mistral
       # Model information for Mistral
       module Models
+        # Mistral's capability flags, named as Models::Schema::CAPABILITIES names.
+        CAPABILITY_FLAGS = {
+          'function_calling' => 'function_calling',
+          'reasoning' => 'reasoning',
+          'fine_tuning' => 'fine_tuning',
+          'vision' => 'vision',
+          'ocr' => 'vision',
+          'moderation' => 'moderation',
+          'audio_transcription' => 'transcription',
+          'audio_transcription_realtime' => 'realtime',
+          'audio_speech' => 'speech_generation'
+        }.freeze
+
         module_function
 
         def models_url
           'models'
         end
 
-        def headers(config)
-          {
-            'Authorization' => "Bearer #{config.mistral_api_key}"
-          }
+        def models_dev_alias(model_id, models_dev_by_key, provider_model = nil)
+          source = Array(provider_model&.metadata&.[](:aliases)).filter_map do |alias_id|
+            models_dev_by_key["mistral:#{alias_id}"]
+          end.first
+          Model.new(source.to_h.merge(id: model_id)) if source
         end
 
-        def parse_list_models_response(response, slug, capabilities)
+        def parse_list_models_response(response, slug)
           Array(response.body['data']).map do |model_data|
             model_id = model_data['id']
+            flags = model_data['capabilities'] || {}
 
-            release_date = capabilities.release_date_for(model_id)
-            created_at = release_date ? Time.parse(release_date) : nil
-
-            Model::Info.new(
+            Model.new(
               id: model_id,
-              name: capabilities.format_display_name(model_id),
+              name: model_id,
               provider: slug,
-              family: capabilities.model_family(model_id),
-              created_at: created_at,
-              context_window: capabilities.context_window_for(model_id),
-              max_output_tokens: capabilities.max_tokens_for(model_id),
-              modalities: capabilities.modalities_for(model_id),
-              capabilities: capabilities.capabilities_for(model_id),
-              pricing: capabilities.pricing_for(model_id),
-              metadata: {
-                object: model_data['object'],
-                owned_by: model_data['owned_by']
-              }
+              context_window: model_data['max_context_length'],
+              modalities: modalities_from(flags, model_data),
+              capabilities: capabilities_from(flags),
+              metadata: metadata_from(model_data)
             )
           end
         end
+
+        def capabilities_from(flags)
+          CAPABILITY_FLAGS.filter_map { |flag, capability| capability if flags[flag] }
+        end
+
+        def modalities_from(flags, model_data)
+          return { input: ['text'], output: ['embeddings'] } if embedding_model?(model_data)
+          return { input: ['audio'], output: ['text'] } if flags['audio_transcription']
+
+          input = ['text']
+          input << 'image' if flags['vision'] || flags['ocr']
+          output = ['text']
+          output << 'audio' if flags['audio_speech']
+          { input: input, output: output }
+        end
+
+        def embedding_model?(model_data)
+          values = [model_data['id'], model_data['description'], *Array(model_data['aliases'])]
+          values.any? { |value| value.to_s.match?(/(?:\A|[-_ ])embed(?:ding)?(?:\z|[-_ ])/i) }
+        end
+
+        def metadata_from(model_data)
+          {
+            object: model_data['object'],
+            owned_by: model_data['owned_by'],
+            description: model_data['description'],
+            aliases: model_data['aliases'],
+            deprecation: model_data['deprecation'],
+            deprecation_replacement_model: model_data['deprecation_replacement_model']
+          }.reject { |_, value| value.nil? || (value.respond_to?(:empty?) && value.empty?) }
+        end
       end
     end
   end

data/lib/ruby_llm/providers/mistral/ocr.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Mistral
+      # Document OCR methods for the Mistral API. Remote files are passed as
+      # URLs; local files are inlined as base64 data URIs. Images go through
+      # the image_url document variant, everything else through document_url.
+      module OCR
+        module_function
+
+        def ocr_url
+          'ocr'
+        end
+
+        def render_ocr_payload(file, model:, pages: nil, provider_options: {})
+          attachment = file.is_a?(Attachment) ? file : Attachment.new(file, config: @config)
+
+          payload = { model: model, document: ocr_document_part(attachment) }
+          payload[:pages] = pages if pages
+          payload.merge(provider_options)
+        end
+
+        def ocr_document_part(attachment)
+          reference = attachment.url? ? attachment.source.to_s : attachment.for_llm
+
+          if attachment.image?
+            { type: 'image_url', image_url: reference }
+          else
+            { type: 'document_url', document_url: reference }
+          end
+        end
+
+        def parse_ocr_response(response, model:)
+          data = response.body
+
+          RubyLLM::OCR.new(
+            pages: data['pages'],
+            model: data['model'] || model,
+            usage: data['usage_info'],
+            raw: data
+          )
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/mistral/speech.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Mistral
+      # Speech dialect for the Mistral API, which requires a voice from its own
+      # catalog and returns JSON with base64 audio instead of raw bytes.
+      module Speech
+        module_function
+
+        def stream_speech(payload, model:, voice:, format:)
+          audio = +''.b
+          final = nil
+          stream_events(speech_url(model:), payload.merge(stream: true)) do |event|
+            case event['type']
+            when 'speech.audio.delta'
+              data = Base64.strict_decode64(event.fetch('audio_data'))
+              audio << data
+              yield SpeechChunk.new(data:, format: format || 'mp3') unless data.empty?
+            when 'speech.audio.done'
+              final = event
+            end
+          end
+          raise Error, 'Mistral speech stream ended before its completion event' unless final
+
+          usage = final['usage'] || {}
+          RubyLLM::Speech.new(data: audio, model:, voice:, format: format || 'mp3',
+                              input_tokens: usage['prompt_tokens'], output_tokens: usage['completion_tokens'])
+        end
+
+        def render_speech_payload(input, model:, voice:, format:, provider_options: {})
+          {
+            model: model,
+            input: input,
+            voice_id: voice,
+            response_format: format
+          }.compact.merge(provider_options)
+        end
+
+        def parse_speech_response(response, model:, voice:, format:)
+          RubyLLM::Speech.new(
+            data: Base64.decode64(response.body['audio_data'].to_s),
+            model: model,
+            voice: voice,
+            format: (format || 'mp3').to_s
+          )
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/mistral/transcription.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Mistral
+      # Transcription dialect for the Mistral audio API. Giving speaker
+      # names turns on diarization, which labels each segment with a
+      # speaker id; a prompt rides along as context biasing terms.
+      module Transcription
+        STREAM_TYPES = {
+          'transcription.text.delta' => RubyLLM::TranscriptionChunk::DELTA,
+          'transcription.segment' => RubyLLM::TranscriptionChunk::SEGMENT,
+          'transcription.done' => RubyLLM::TranscriptionChunk::DONE
+        }.freeze
+
+        def render_transcription_options(timestamps:, **)
+          return {} if timestamps.nil?
+          raise ArgumentError, 'Mistral transcription timestamps must be segment' unless timestamps == :segment
+
+          { timestamp_granularities: ['segment'] }
+        end
+
+        module_function
+
+        def build_transcription_chunk(data)
+          type = STREAM_TYPES[data['type']]
+          return super unless type
+
+          RubyLLM::TranscriptionChunk.new(
+            type: type,
+            delta: (data['text'] if type == RubyLLM::TranscriptionChunk::DELTA),
+            text: (data['text'] if type == RubyLLM::TranscriptionChunk::DONE),
+            segment: (data.except('type') if type == RubyLLM::TranscriptionChunk::SEGMENT),
+            raw: data
+          )
+        end
+
+        def transcription_duration(usage)
+          usage['prompt_audio_seconds'] || super
+        end
+
+        # rubocop:disable-next Lint/UnusedMethodArgument
+        def render_transcription_payload(file_part, model:, language:, format: nil, speaker_names: nil,
+                                         speaker_references: nil, provider_options: {}, prompt: nil,
+                                         temperature: nil)
+          payload = {
+            model: model,
+            file: file_part,
+            language: language,
+            temperature: temperature,
+            context_bias: prompt ? Array(prompt) : nil
+          }.compact
+          if speaker_names
+            payload[:diarize] = true
+            payload[:timestamp_granularities] = ['segment']
+          end
+          payload.merge(provider_options)
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/mistral.rb

@@ -3,10 +3,16 @@
 module RubyLLM
   module Providers
     # Mistral API integration.
-    class Mistral < OpenAI
-      include Mistral::Chat
-      include Mistral::Models
-      include Mistral::Embeddings
+    class Mistral < Provider
+      protocol :chat_completions, ChatCompletions, batches: Mistral::ChatCompletions::Batches
+      protocol :conversations, Conversations
+      protocol :files, Protocols::Mistral::Files
+
+      def protocol_for(model, operation: nil, **)
+        return Conversations if operation == :paint
+
+        super
+      end
 
       def api_base
         @config.mistral_api_base || 'https://api.mistral.ai/v1'
@@ -18,11 +24,17 @@ module RubyLLM
         }
       end
 
+      def batch_cost_multiplier(**) = 0.5
+
       class << self
         def capabilities
           Mistral::Capabilities
         end
 
+        def models_dev_alias(...)
+          Mistral::Models.models_dev_alias(...)
+        end
+
         def configuration_options
           %i[mistral_api_key mistral_api_base]
         end

data/lib/ruby_llm/providers/ollama/chat.rb

@@ -7,22 +7,16 @@ module RubyLLM
       module Chat
         module_function
 
-        def format_messages(messages)
-          messages.map do |msg|
-            {
-              role: format_role(msg.role),
-              content: format_message_content(msg),
-              tool_calls: format_tool_calls(msg.tool_calls),
-              tool_call_id: msg.tool_call_id
-            }.compact.merge(OpenAI::Chat.format_thinking(msg))
+        def render_payload(messages, tools:, temperature:, model:, stream: false, max_output_tokens: nil,
+                           schema: nil, thinking: nil, citations: false, caching: nil, tool_prefs: nil)
+          if thinking&.budget
+            RubyLLM.logger.debug { "Ollama has no thinking budgets; ignoring budget #{thinking.budget}" }
           end
+          super
         end
 
-        def format_message_content(msg)
-          content = Ollama::Media.format_content(msg.content)
-          return '' if content.nil? && OpenAI::Chat.thinking_only_assistant_message?(msg)
-
-          content
+        def format_content(content, attachments = [])
+          Ollama::Media.format_content(content, attachments)
         end
 
         def format_role(role)

data/lib/ruby_llm/providers/ollama/media.rb

@@ -5,30 +5,21 @@ module RubyLLM
     class Ollama
       # Handles formatting of media content (images, audio) for Ollama APIs
       module Media
-        extend OpenAI::Media
-
         module_function
 
-        def format_content(content)
-          return content.value if content.is_a?(RubyLLM::Content::Raw)
-          return content.to_json if content.is_a?(Hash) || content.is_a?(Array)
-          return content unless content.is_a?(Content)
-
-          parts = []
-          parts << format_text(content.text) if content.text
-
-          content.attachments.each do |attachment|
+        def format_content(content, attachments = [])
+          Protocols::ChatCompletions::Media.format_parts(content, attachments) do |attachment|
             case attachment.type
             when :image
-              parts << Ollama::Media.format_image(attachment)
+              format_image(attachment)
+            when :audio
+              Protocols::ChatCompletions::Media.format_audio(attachment)
             when :text
-              parts << format_text_file(attachment)
+              Protocols::ChatCompletions::Media.format_text_file(attachment)
             else
               raise UnsupportedAttachmentError, attachment.mime_type
             end
           end
-
-          parts
         end
 
         def format_image(image)

data/lib/ruby_llm/providers/ollama/models.rb

@@ -5,24 +5,37 @@ module RubyLLM
     class Ollama
       # Models methods for the Ollama API integration
       module Models
+        SHOW_URL = '../api/show'
+        CAPABILITY_MAP = {
+          'tools' => 'function_calling',
+          'vision' => 'vision',
+          'thinking' => 'reasoning'
+        }.freeze
+
         def models_url
           'models'
         end
 
-        def parse_list_models_response(response, slug, _capabilities)
-          data = response.body['data'] || []
-          data.map do |model|
-            Model::Info.new(
+        def model_capabilities
+          %w[streaming structured_output]
+        end
+
+        def list_models
+          response = @connection.get models_url
+          parse_list_models_response(response, @provider.slug, details: show_models(response))
+        end
+
+        def parse_list_models_response(response, slug, details: {})
+          Array(response.body['data']).map do |model|
+            reported = Array(details[model['id']])
+            Model.new(
               id: model['id'],
               name: model['id'],
               provider: slug,
               family: 'ollama',
               created_at: model['created'] ? Time.at(model['created']) : nil,
-              modalities: {
-                input: %w[text image],
-                output: %w[text]
-              },
-              capabilities: %w[streaming function_calling structured_output vision],
+              modalities: build_modalities(reported),
+              capabilities: build_capabilities(reported),
               pricing: {},
               metadata: {
                 owned_by: model['owned_by']
@@ -30,6 +43,34 @@ module RubyLLM
             )
           end
         end
+
+        private
+
+        def show_models(response)
+          Array(response.body['data']).filter_map { |model| model['id'] }.to_h { |id| [id, show_model(id)] }
+        end
+
+        def show_model(id)
+          Array(@connection.post(SHOW_URL, { model: id }).body['capabilities'])
+        rescue Error, Faraday::Error => e
+          RubyLLM.logger.debug "Ollama did not report capabilities for #{id} (#{e.message})."
+          []
+        end
+
+        def build_capabilities(reported)
+          return [] if reported.include?('embedding')
+
+          derived = CAPABILITY_MAP.filter_map { |native, capability| capability if reported.include?(native) }
+          model_capabilities + derived
+        end
+
+        def build_modalities(reported)
+          return { input: %w[text], output: %w[embeddings] } if reported.include?('embedding')
+
+          input = %w[text]
+          input << 'image' if reported.include?('vision')
+          { input: input, output: %w[text] }
+        end
       end
     end
   end