ruby_llm 2.1.1 → 2.2.0

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

@@ -6,15 +6,43 @@ module RubyLLM
       # Chat methods for Mistral API
       module Chat
         PROMPT_CACHE_OPTIONS = %i[key].freeze
+        FINISH_REASONS = Protocols::ChatCompletions::Chat::FINISH_REASONS.merge('model_length' => :max_tokens).freeze
 
         module_function
 
-        def format_role(role)
+        def finish_reasons = FINISH_REASONS
+
+        # A tool that supplies numbered sources gets reference chunks back,
+        # whose reference_ids are the keys of that tool's own result.
+        def extract_citations(message_data, data, content)
+          references = reference_citations(message_data['content'])
+          references.any? ? references : super
+        end
+
+        def parse_chunk_citations(delta, data)
+          references = reference_citations(delta['content'], offsets: false)
+          references.any? ? references : super
+        end
+
+        def reference_citations(content, offsets: true)
+          return [] unless content.is_a?(Array)
+
+          position = 0
+          content.flat_map do |part|
+            position += part['text'].to_s.length if part['type'] == 'text'
+            span = offsets ? { start_index: position, end_index: position } : {}
+            Array(part['reference_ids']).map { |id| Citation.new(source_index: id, **span) }
+          end
+        end
+
+        def render_role(role)
           role.to_s
         end
 
         def render_payload(messages, tools:, temperature:, model:, stream: false, max_output_tokens: nil,
                            schema: nil, thinking: nil, citations: false, caching: nil, tool_prefs: nil)
+          raise ArgumentError, 'Mistral has no thinking budget; use effort:' if thinking&.budget
+
           payload = super
           payload.delete(:stream_options)
           normalize_required_tool_choice(payload)
@@ -42,6 +70,8 @@ module RubyLLM
           raise ArgumentError, "Mistral prompt caching accepts :key, got #{format_cache_option_keys(unsupported)}"
         end
 
+        def cache_option_names = PROMPT_CACHE_OPTIONS
+
         def format_cache_option_keys(keys)
           keys.map { |key| ":#{key}" }.join(', ')
         end
@@ -49,7 +79,7 @@ module RubyLLM
         # extra_content carries Gemini's thought signatures on Google's
         # OpenAI-compatible endpoints, and Mistral refuses any field it does
         # not define.
-        def format_tool_calls(tool_calls)
+        def render_tool_calls(tool_calls)
           super&.map { |call| call.except(:extra_content) }
         end
 
@@ -73,7 +103,7 @@ module RubyLLM
 
         # Mistral carries reasoning in content blocks, not in the top-level
         # reasoning fields the rest of the Chat Completions family uses.
-        def format_message_content(msg, **)
+        def render_message_content(msg, **)
           formatted_content = super
           return formatted_content unless msg.role == :assistant && msg.thinking
 
@@ -83,7 +113,7 @@ module RubyLLM
           content_blocks
         end
 
-        def format_thinking(_msg)
+        def render_thinking(_msg)
           {}
         end
 

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

@@ -99,14 +99,15 @@ module RubyLLM
             custom_id, shape = line['custom_id'].split(':', 2)
             index = batch_result_index(custom_id)
             response = line['response']
-            body = response && response['body']
+            body = response['body'] if batch_response_succeeded?(response)
 
-            if body && body['data'].is_a?(Array) && !embedding_positions?(body['data'])
+            if body && body['data'].is_a?(Array) && !valid_embedding_positions?(body['data'])
               return [index, nil, batch_failure(line['custom_id'], 'Invalid or duplicate embedding record positions')]
             end
 
             if body
-              [index, parse_mistral_batch_body(body, shape:)]
+              billed = body unless body['data'].is_a?(Array)
+              parsed_batch_result(index, custom_id, billed) { parse_mistral_batch_body(body, shape:) }
             else
               [index, nil, batch_failure(line['custom_id'], batch_error_message(line))]
             end
@@ -115,17 +116,7 @@ module RubyLLM
           def parse_mistral_batch_body(body, shape:)
             return parse_completion_body(body, raw: body) unless body['data'].is_a?(Array)
 
-            ordered = body.merge('data' => body['data'].sort_by { |row| row['index'] })
-            parse_embedding_response(Response.new(ordered), model: body['model'], text: shape == 'array' ? [] : nil)
-          end
-
-          # A row's index names the input its vector belongs to, so rows are put
-          # back in that order and positions that are not exactly 0...N fail the
-          # request rather than pairing vectors with the wrong inputs.
-          def embedding_positions?(rows)
-            positions = rows.map { |row| row.is_a?(Hash) ? row['index'] : nil }
-
-            positions.all?(Integer) && positions.sort == (0...positions.size).to_a
+            parse_embedding_response(Response.new(body), model: body['model'], text: shape == 'array' ? [] : nil)
           end
         end
       end

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

@@ -11,24 +11,13 @@ module RubyLLM
           'embeddings'
         end
 
-        # rubocop:disable-next Lint/UnusedMethodArgument
-        def render_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil, provider_options: {})
+        def render_embedding_payload(text, model:, dimensions:, provider_options: {}, **)
           {
             model: model,
             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')
-          vectors = data['data'].map { |d| d['embedding'] }
-
-          vectors = vectors.first if vectors.length == 1 && !text.is_a?(Array)
-
-          Embedding.new(vectors:, model:, input_tokens:)
-        end
       end
     end
   end

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

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Mistral
+      # Counts a request through Mistral's Responses input-token endpoint,
+      # the only token counting route Mistral offers, whichever protocol the
+      # chat itself runs on.
+      class InputTokens < Protocols::Responses
+        include Protocols::Responses::TokenCounting
+
+        def render_count_tokens_payload(messages, **)
+          payload = super(messages.flat_map { |message| countable_messages(message) }, **)
+          payload.merge(input: payload[:input].map { |item| typed_item(item) })
+        end
+
+        private
+
+        # Mistral rejects an input item that doesn't name its type.
+        def typed_item(item)
+          item.key?(:role) ? { type: 'message' }.merge(item) : item
+        end
+
+        # Native content from other protocols doesn't fit Responses input,
+        # and the endpoint can't take unsigned thinking, so thinking goes in
+        # as text to keep it in the count.
+        def countable_messages(message)
+          plain = message.without_native_content
+          text = message.thinking&.text
+          text ? [Message.new(role: :assistant, content: text), plain] : [plain]
+        end
+      end
+    end
+  end
+end

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

@@ -11,31 +11,31 @@ module RubyLLM
           %i[image audio pdf document text].include?(attachment.type)
         end
 
-        def format_content(content, attachments = [])
-          Protocols::ChatCompletions::Media.format_parts(content, attachments) do |attachment|
+        def render_content(content, attachments = [])
+          Protocols::ChatCompletions::Media.render_parts(content, attachments) do |attachment|
             raise UnsupportedAttachmentError, attachment.mime_type unless supported_attachment?(attachment)
 
             case attachment.type
             when :image
-              format_image(attachment)
+              render_image(attachment)
             when :audio
-              Protocols::ChatCompletions::Media.format_audio(attachment)
+              Protocols::ChatCompletions::Media.render_audio(attachment)
             when :pdf, :document
-              format_document(attachment)
+              render_document(attachment)
             when :text
-              Protocols::ChatCompletions::Media.format_text_file(attachment)
+              Protocols::ChatCompletions::Media.render_text_file(attachment)
             end
           end
         end
 
-        def format_image(image)
+        def render_image(image)
           {
             type: 'image_url',
             image_url: image.url_or_data_uri
           }
         end
 
-        def format_document(document)
+        def render_document(document)
           {
             type: 'document_url',
             document_url: document.url_or_data_uri

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

@@ -3,8 +3,8 @@
 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
+      # Transcription dialect for the Mistral audio API. Requesting speaker
+      # labels turns on diarization, which labels each segment with a
       # speaker id; a prompt rides along as context biasing terms.
       module Transcription
         STREAM_TYPES = {
@@ -20,6 +20,10 @@ module RubyLLM
           { timestamp_granularities: ['segment'] }
         end
 
+        def rejected_options(operation)
+          operation == :transcribe ? %i[format speaker_references named_speakers] : super
+        end
+
         module_function
 
         def build_transcription_chunk(data)
@@ -39,10 +43,8 @@ module RubyLLM
           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)
+        def render_transcription_payload(file_part, model:, language:, speaker_names: nil, provider_options: {},
+                                         prompt: nil, temperature: nil, **)
           payload = {
             model: model,
             file: file_part,

data/lib/ruby_llm/providers/mistral.rb

@@ -14,8 +14,12 @@ module RubyLLM
         super
       end
 
+      def count_tokens(messages, model:, tools: {}, protocol: nil, **) # rubocop:disable Lint/UnusedMethodArgument
+        InputTokens.new(self, model).count_tokens(messages, tools:, **)
+      end
+
       def api_base
-        @config.mistral_api_base || 'https://api.mistral.ai/v1'
+        configured_api_base(:mistral_api_base, 'https://api.mistral.ai/v1', endpoints: Protocols::ChatCompletions::ENDPOINT_PATHS)
       end
 
       def headers
@@ -27,6 +31,14 @@ module RubyLLM
       def batch_cost_multiplier(**) = 0.5
 
       class << self
+        def open_telemetry_name # :nodoc:
+          'mistral_ai'
+        end
+
+        def models_dev_provider_ids # :nodoc:
+          %w[mistral]
+        end
+
         def capabilities
           Mistral::Capabilities
         end

data/lib/ruby_llm/providers/moonshot/capabilities.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Moonshot
+      # Feature capability gaps not represented in upstream model catalogs.
+      module Capabilities
+        def self.augment(capabilities, **)
+          return capabilities unless capabilities.include?('function_calling')
+
+          capabilities | ['tool_choice']
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Moonshot
+      # Chat methods of the Moonshot AI API integration
+      # https://platform.kimi.ai/docs/api/chat
+      module Chat
+        module_function
+
+        def render_role(role)
+          role.to_s
+        end
+
+        def max_output_tokens_field(_model)
+          :max_completion_tokens
+        end
+
+        def render_payload(messages, tools:, temperature:, model:, stream: false, max_output_tokens: nil,
+                           schema: nil, thinking: nil, citations: false, caching: nil, tool_prefs: nil)
+          payload = super
+          raise ArgumentError, 'Moonshot does not accept calls: :one' if payload.delete(:parallel_tool_calls) == false
+
+          payload[:thinking] = { type: thinking.enabled ? 'enabled' : 'disabled' } unless thinking&.enabled.nil?
+          payload.merge!(prompt_cache_params(caching)) if caching
+          payload
+        end
+
+        def count_tokens_url
+          'tokenizers/estimate-token-count'
+        end
+
+        def render_count_tokens_payload(messages, model:, **)
+          render_payload(messages, model:, temperature: nil, **).slice(:model, :messages, :tools, :response_format)
+        end
+
+        def parse_count_tokens_response(response)
+          response.body.dig('data', 'total_tokens')
+        end
+
+        def apply_end_user(payload, identifier)
+          payload.merge(safety_identifier: identifier)
+        end
+
+        # Moonshot takes prompt_cache_key and prompt_cache_options, and
+        # refuses the explicit cache breakpoints OpenAI accepts.
+        def openai_prompt_caching?
+          false
+        end
+
+        def render_thinking(msg)
+          text = msg.thinking&.text if msg.role == :assistant
+          text ? { reasoning_content: text } : {}
+        end
+
+        def supported_attachment?(attachment)
+          return attachment.image? || own_video?(attachment) if attachment.provider_file?
+
+          attachment.video? || Protocols::ChatCompletions::Media.supported_attachment?(
+            attachment, document_attachments: :none, audio_attachments: false
+          )
+        end
+
+        def render_content(content, attachments = [])
+          Protocols::ChatCompletions::Media.render_parts(content, attachments) do |attachment|
+            if attachment.image?
+              render_inline_image(attachment)
+            elsif attachment.video?
+              render_inline_video(attachment)
+            else
+              Protocols::ChatCompletions::Media.render_attachment(
+                attachment, document_attachments: :none, image_attachments: true, audio_attachments: false
+              )
+            end
+          end
+        end
+
+        def render_inline_image(image)
+          { type: 'image_url', image_url: { url: inline_url(image) } }
+        end
+
+        def own_video?(attachment)
+          attachment.video? && attachment.source.provider == 'moonshot'
+        end
+
+        def render_inline_video(video)
+          raise UnsupportedAttachmentError, video.mime_type unless supported_attachment?(video)
+
+          { type: 'video_url', video_url: { url: inline_url(video) } }
+        end
+
+        # Moonshot reads media only from data URIs and its own file references,
+        # so remote files travel inline.
+        def inline_url(attachment)
+          attachment.provider_file? ? "ms://#{attachment.provider_file_id}" : attachment.for_llm
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/moonshot/documents.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+require 'stringio'
+
+module RubyLLM
+  module Providers
+    class Moonshot
+      # Kimi models cannot read a document attached to a message. Kimi
+      # documents one way to ask about a file: extract its text through the
+      # Files API and send that text in the conversation.
+      # https://platform.kimi.ai/docs/guide/ask-questions-about-pdf-content
+      module Documents
+        EXTRACT_PURPOSE = 'file-extract'
+
+        def preprocess_message(message, unsupported_attachment: nil)
+          return super unless message.attachments.any? { |attachment| extracted_document?(attachment) }
+
+          super(message.with_attachments(message.attachments.map do |attachment|
+            extracted_document?(attachment) ? extracted_attachment(attachment) : attachment
+          end), unsupported_attachment:)
+        end
+
+        # Kimi extracts a whole document without a model, so +model:+ only
+        # labels the result. The upload exists only to read it back.
+        def ocr(file, model:, pages: nil, provider_options: {})
+          raise ArgumentError, "#{@provider.name} extracts whole documents; pages: is not available" if pages
+
+          track_usage(:ocr) do
+            upload = @provider.upload_file(file, purpose: EXTRACT_PURPOSE, provider_options:)
+            begin
+              text = @provider.download_file(upload.id)
+              RubyLLM::OCR.new(pages: [{ 'index' => 0, 'markdown' => text }], model:, raw: upload.metadata)
+            ensure
+              delete_upload(upload)
+            end
+          end
+        end
+
+        private
+
+        def provider_file_upload_options(_attachment)
+          { purpose: EXTRACT_PURPOSE }
+        end
+
+        def extracted_document?(attachment)
+          return attachment.source.purpose == EXTRACT_PURPOSE if attachment.provider_file?
+
+          %i[pdf document].include?(attachment.type)
+        end
+
+        def extracted_attachment(attachment)
+          filename = "#{File.basename(attachment.filename.to_s, '.*')}.md"
+          Attachment.new(StringIO.new(extracted_text(attachment)), filename:, config: @config)
+        end
+
+        def extracted_text(attachment)
+          return @provider.download_file(attachment.provider_file_id) if attachment.provider_file?
+
+          download_upload(attachment)
+        rescue Error => e
+          message = Message.new(role: :user, content: nil, attachments: [attachment])
+          raise if attachment.provider_file? || discard_missing_uploads([message], e).empty?
+
+          download_upload(attachment)
+        end
+
+        def download_upload(attachment)
+          @provider.download_file(provider_upload(attachment).id)
+        end
+
+        def delete_upload(upload)
+          @connection.delete("files/#{upload.id}")
+        rescue Error
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/moonshot/responses.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Moonshot
+      # Moonshot's dialect of the OpenAI Responses API, the route to Kimi's
+      # server-side web search. It holds no state: store is always false.
+      # https://platform.kimi.ai/docs/api/responses
+      class Responses < Protocols::Responses
+        include Moonshot::Documents
+
+        SERVER_TOOL_ALIASES = {
+          web_search: { tool: { type: 'web_search' } }
+        }.freeze
+
+        def server_tool_aliases
+          SERVER_TOOL_ALIASES
+        end
+
+        def web_search_sources?
+          true
+        end
+
+        def apply_end_user(payload, identifier)
+          payload.merge(safety_identifier: identifier)
+        end
+
+        # Kimi reads images only from data URIs, so remote images travel inline.
+        def render_image(image, image_url: image.for_llm, **)
+          super
+        end
+
+        private
+
+        def provider_file_references?
+          false
+        end
+
+        # Kimi silently drops the web_search tool for a message whose content
+        # is a bare string, so text always travels as input_text parts.
+        def render_item(msg, caching: nil)
+          item = super
+          return item unless item.is_a?(Hash) && item[:role] && item[:content].is_a?(String)
+
+          item.merge(content: [{ type: 'input_text', text: item[:content] }])
+        end
+
+        def parse_server_tool_use(response)
+          searches = Array(response['output']).count { |item| item['type'] == 'web_search_call' }
+          searches.positive? ? { 'web_search_requests' => searches } : {}
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/moonshot.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    # Moonshot AI API integration for the Kimi models.
+    class Moonshot < Provider
+      # Moonshot's dialect of the Chat Completions API.
+      class ChatCompletions < Protocols::ChatCompletions
+        include Moonshot::Chat
+        include Moonshot::Documents
+      end
+
+      protocol :chat_completions, ChatCompletions, batches: Protocols::ChatCompletions::Batches
+      protocol :responses, Moonshot::Responses
+      protocol :files, Protocols::Moonshot::Files
+
+      def api_base
+        configured_api_base(:moonshot_api_base, 'https://api.moonshot.ai/v1', endpoints: Protocols::ChatCompletions::ENDPOINT_PATHS)
+      end
+
+      def headers
+        { 'Authorization' => "Bearer #{@config.moonshot_api_key}" }
+      end
+
+      def account_identity
+        account_digest(api_base, @config.moonshot_api_key)
+      end
+
+      def batch_cost_multiplier(**) = 0.6
+
+      class << self
+        def open_telemetry_name # :nodoc:
+          'moonshot_ai'
+        end
+
+        def models_dev_provider_ids # :nodoc:
+          %w[moonshotai]
+        end
+
+        def capabilities
+          Moonshot::Capabilities
+        end
+
+        def configuration_options
+          %i[moonshot_api_key moonshot_api_base]
+        end
+
+        def configuration_requirements
+          %i[moonshot_api_key]
+        end
+      end
+    end
+  end
+end

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

@@ -9,17 +9,16 @@ module RubyLLM
 
         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
+          raise ArgumentError, 'Ollama has no thinking budget; use effort:' if thinking&.budget
+
           super
         end
 
-        def format_content(content, attachments = [])
-          Ollama::Media.format_content(content, attachments)
+        def render_content(content, attachments = [])
+          Ollama::Media.render_content(content, attachments)
         end
 
-        def format_role(role)
+        def render_role(role)
           role.to_s
         end
       end

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

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class Ollama
+      # Ollama embeds text, an image, and audio together when the input is
+      # one object holding each as a field.
+      module Embeddings
+        EMBEDDING_MEDIA_FIELDS = %i[image audio].freeze
+
+        module_function
+
+        def embedding_media?
+          true
+        end
+
+        def render_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil, with: [],
+                                     provider_options: {})
+          input = with.any? ? media_embedding_input(text, with) : text
+          super(input, model:, dimensions:, task_type:, title:, provider_options:)
+        end
+
+        def media_embedding_input(text, attachments)
+          raise ArgumentError, 'embed one text at a time when embedding attachments' if text.is_a?(Array)
+
+          input = text.to_s.empty? ? {} : { text: text }
+          attachments.each do |attachment|
+            unless EMBEDDING_MEDIA_FIELDS.include?(attachment.type)
+              raise UnsupportedAttachmentError,
+                    attachment.mime_type
+            end
+            raise ArgumentError, "Ollama embeds one #{attachment.type} per input" if input.key?(attachment.type)
+
+            input[attachment.type] = attachment.encoded
+          end
+          input
+        end
+      end
+    end
+  end
+end