ruby_llm 1.16.0 → 2.0.0.rc2

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

@@ -3,60 +3,56 @@
 module RubyLLM
   module Providers
     class XAI
-      # Models metadata for xAI list models.
+      # Models metadata for xAI list models. The list endpoint omits the
+      # model-less TTS and STT services, so those are appended under the
+      # grok-tts and grok-stt names from xAI's model catalog.
       module Models
         module_function
 
-        def parse_list_models_response(response, slug, _capabilities)
+        AUDIO_MODELS = {
+          'grok-tts' => {
+            name: 'Grok TTS',
+            modalities: { input: ['text'], output: ['audio'] },
+            capabilities: []
+          },
+          'grok-stt' => {
+            name: 'Grok STT',
+            modalities: { input: ['audio'], output: ['text'] },
+            capabilities: ['transcription']
+          }
+        }.freeze
+
+        def parse_list_models_response(response, slug)
           Array(response.body['data']).map do |model_data|
             model_id = model_data['id']
 
-            Model::Info.new(
+            Model.new(
               id: model_id,
-              name: format_display_name(model_id),
+              name: model_id,
               provider: slug,
-              family: 'grok',
               created_at: model_data['created'] ? Time.at(model_data['created']) : nil,
-              context_window: nil,
-              max_output_tokens: nil,
-              modalities: modalities_for(model_id),
-              capabilities: capabilities_for(model_id),
-              pricing: {},
               metadata: {
                 object: model_data['object'],
                 owned_by: model_data['owned_by']
               }.compact
             )
-          end
+          end + audio_models(slug)
         end
 
-        def modalities_for(model_id)
-          if image_model?(model_id)
-            { input: %w[text image], output: ['image'] }
-          elsif video_model?(model_id)
-            { input: %w[text image video], output: ['video'] }
-          else
-            { input: ['text'], output: ['text'] }
+        def audio_models(slug)
+          AUDIO_MODELS.map do |model_id, info|
+            Model.new(
+              id: model_id,
+              name: info[:name],
+              provider: slug,
+              family: 'grok',
+              modalities: info[:modalities],
+              capabilities: info[:capabilities],
+              pricing: {},
+              metadata: {}
+            )
           end
         end
-
-        def capabilities_for(model_id)
-          return ['vision'] if image_model?(model_id) || video_model?(model_id)
-
-          ['streaming']
-        end
-
-        def format_display_name(model_id)
-          model_id.tr('-', ' ').split.map(&:capitalize).join(' ')
-        end
-
-        def image_model?(model_id)
-          model_id.match?(/\Agrok-(?:2-)?imagine-image/) || model_id == 'grok-2-image-1212'
-        end
-
-        def video_model?(model_id)
-          model_id.match?(/\Agrok-(?:2-)?imagine-video/)
-        end
       end
     end
   end

data/lib/ruby_llm/providers/xai/reported_cost.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      module ReportedCost # :nodoc:
+        USD_PER_TICK = 1e-10
+
+        private
+
+        def reported_cost(usage)
+          ticks = usage['cost_in_usd_ticks']
+          ticks ? ticks * USD_PER_TICK : nil
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      # xAI's dialect of the OpenAI Responses API, the primary protocol for
+      # Grok models. Usage counts agentic tool activity.
+      class Responses < Protocols::Responses
+        include XAI::ReportedCost
+        include XAI::Images
+        include XAI::Models
+        include XAI::Speech
+        include XAI::Transcription
+        include XAI::Videos
+        include Protocols::XAI::Tokenization
+        include Protocols::Responses::Compaction
+        include Protocols::XAI::StreamingTranscription
+
+        SERVER_TOOL_ALIASES = {
+          web_search: { tool: { type: 'web_search' } },
+          x_search: { tool: { type: 'x_search' } },
+          code_execution: { tool: { type: 'code_execution' } },
+          code_interpreter: { tool: { type: 'code_interpreter' } },
+          file_search: { tool: { type: 'file_search' } },
+          collections_search: { tool: { type: 'file_search' } },
+          image_generation: { tool: { type: 'image_generation' } },
+          mcp: Protocols::Responses::SERVER_TOOL_ALIASES.fetch(:mcp)
+        }.freeze
+
+        SERVER_TOOL_USAGE_COUNTERS = %w[num_sources_used num_server_side_tools_used].freeze
+
+        def server_tool_aliases
+          SERVER_TOOL_ALIASES
+        end
+
+        def parse_usage(usage)
+          super.merge(
+            server_tool_use: server_side_tool_usage(usage),
+            reported_cost: reported_cost(usage)
+          )
+        end
+
+        private
+
+        def server_side_tool_usage(usage)
+          counters = usage.slice(*SERVER_TOOL_USAGE_COUNTERS).reject { |_, count| count.to_i.zero? }
+          counters.empty? ? nil : counters
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      # Speech dialect for the xAI Voice API. The /v1/tts endpoint is
+      # model-less: it takes text, voice_id, and a required language
+      # (defaulted to auto-detection) and returns raw audio bytes.
+      module Speech
+        module_function
+
+        def speech_url(model:) # rubocop:disable Lint/UnusedMethodArgument
+          'tts'
+        end
+
+        def render_speech_payload(input, model:, voice:, format:, provider_options: {}) # rubocop:disable Lint/UnusedMethodArgument
+          payload = {
+            text: input,
+            voice_id: voice,
+            language: 'auto'
+          }.compact
+          payload[:output_format] = { codec: format.to_s } if format
+          Support::Utils.deep_merge(payload, provider_options)
+        end
+
+        def parse_speech_response(response, model:, voice:, format:)
+          RubyLLM::Speech.new(
+            data: response.body,
+            model: model,
+            voice: voice || 'eve',
+            format: (format || 'mp3').to_s
+          )
+        end
+
+        def stream_speech(payload, model:, voice:, format:, &)
+          if payload[:with_timestamps] || payload['with_timestamps']
+            raise ArgumentError, 'xAI streaming speech does not accept with_timestamps'
+          end
+
+          stream_speech_response(speech_url(model:), payload, model:, voice:, format:, &)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      # Transcription dialect for the xAI Voice API. The /v1/stt endpoint is
+      # model-less multipart with the audio file last in the form; giving
+      # speaker names turns on diarization, which labels each word with a
+      # numeric speaker index.
+      module Transcription
+        def render_transcription_options(timestamps:, **)
+          return {} if timestamps.nil? || timestamps == :word
+
+          raise ArgumentError, 'xAI transcription timestamps must be word'
+        end
+
+        module_function
+
+        def transcription_url
+          'stt'
+        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 = {}
+          payload[:language] = language if language
+          payload[:diarize] = true if speaker_names
+          payload.merge(provider_options).merge(file: file_part)
+        end
+
+        def parse_transcription_response(response, model:)
+          data = response.body
+
+          RubyLLM::Transcription.new(
+            text: data['text'],
+            model: model,
+            language: data['language'],
+            duration: data['duration'],
+            segments: data['channels'],
+            words: data['words']
+          )
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/xai/videos.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class XAI
+      # Grok Imagine generation, editing, and extension jobs.
+      module Videos
+        def video_url
+          'videos/generations'
+        end
+
+        def video_request_url(payload)
+          payload.key?(:video) ? 'videos/edits' : video_url
+        end
+
+        def video_extension_url
+          'videos/extensions'
+        end
+
+        def render_video_payload(prompt, model:, with: [], provider_options: {})
+          payload = { model: model, prompt: prompt }
+          if (attachment = with.first)
+            key = attachment.video? ? :video : :image
+            payload[key] = video_reference(attachment)
+          end
+
+          payload.merge(provider_options)
+        end
+
+        def render_video_extension_payload(prompt, model:, extend:, provider_options: {})
+          video = video_extension_attachment(extend)
+          { model: model, prompt: prompt, video: video_reference(video) }.merge(provider_options)
+        end
+
+        def parse_video_job(response, model:)
+          id = response.body['request_id']
+          raise Error.new('xAI did not return a video request id', response:) unless id
+
+          VideoJob.new(id: id, protocol: self, model: model, raw: response.body)
+        end
+
+        def video_job_url(job)
+          "videos/#{job.id}"
+        end
+
+        def parse_video_job_status(response, job:) # rubocop:disable Lint/UnusedMethodArgument
+          body = response.body
+          case body['status']
+          when 'done' then { status: :completed, raw: body }
+          when 'failed', 'expired' then { status: :failed, raw: body, error: body['error'] || body['status'] }
+          else { status: :pending, raw: body }
+          end
+        end
+
+        def download_video(job)
+          video = job.raw['video'] || {}
+
+          Video.new(
+            url: video['url'],
+            mime_type: 'video/mp4',
+            duration: video['duration'],
+            model: job.raw['model'] || job.model,
+            raw: job.raw
+          )
+        end
+
+        private
+
+        def validate_animate_inputs!(with:)
+          raise Error, 'xAI video generation takes a single reference image or video' if with.size > 1
+
+          with.each do |attachment|
+            next if attachment.provider_file? || attachment.url?
+            raise UnsupportedAttachmentError, attachment.mime_type unless attachment.image? || attachment.video?
+          end
+        end
+
+        def video_reference(attachment)
+          return { file_id: attachment.provider_file_id } if attachment.provider_file?
+          return { url: attachment.source.to_s } if attachment.url?
+
+          { url: attachment.for_llm }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/xai.rb

@@ -3,9 +3,10 @@
 module RubyLLM
   module Providers
     # xAI API integration
-    class XAI < OpenAI
-      include XAI::Chat
-      include XAI::Models
+    class XAI < Provider
+      protocol :responses, XAI::Responses
+      protocol :chat_completions, ChatCompletions, batches: XAI::ChatCompletions::Batches
+      protocol :files, Protocols::XAI::Files
 
       def api_base
         @config.xai_api_base || 'https://api.x.ai/v1'
@@ -13,12 +14,15 @@ module RubyLLM
 
       def headers
         {
-          'Authorization' => "Bearer #{@config.xai_api_key}",
-          'Content-Type' => 'application/json'
+          'Authorization' => "Bearer #{@config.xai_api_key}"
         }
       end
 
       class << self
+        def capabilities
+          XAI::Capabilities
+        end
+
         def configuration_options
           %i[xai_api_key xai_api_base]
         end
@@ -27,6 +31,12 @@ module RubyLLM
           %i[xai_api_key]
         end
       end
+
+      private
+
+      def batch_protocol
+        batch_protocol_for_name(:chat_completions)
+      end
     end
   end
 end

data/lib/ruby_llm/railtie.rb

@@ -2,8 +2,7 @@
 
 if defined?(Rails::Railtie)
   module RubyLLM
-    # Rails integration for RubyLLM
-    class Railtie < Rails::Railtie
+    class Railtie < Rails::Railtie # :nodoc:
       initializer 'ruby_llm.inflections' do
         ActiveSupport::Inflector.inflections(:en) do |inflect|
           inflect.acronym 'RubyLLM'
@@ -17,23 +16,15 @@ if defined?(Rails::Railtie)
       initializer 'ruby_llm.active_record' do
         ActiveSupport.on_load :active_record do
           require 'ruby_llm/active_record/payload_helpers'
+          require 'ruby_llm/active_record/model'
+          require 'ruby_llm/active_record/tool_call'
+          require 'ruby_llm/active_record/usage'
+          require 'ruby_llm/active_record/batch'
           require 'ruby_llm/active_record/chat_methods'
           require 'ruby_llm/active_record/message_methods'
-          require 'ruby_llm/active_record/model_methods'
-          require 'ruby_llm/active_record/tool_call_methods'
 
-          if RubyLLM.config.use_new_acts_as
-            require 'ruby_llm/active_record/acts_as'
-            ::ActiveRecord::Base.include RubyLLM::ActiveRecord::ActsAs
-          else
-            require 'ruby_llm/active_record/acts_as_legacy'
-            ::ActiveRecord::Base.include RubyLLM::ActiveRecord::ActsAsLegacy
-
-            RubyLLM.deprecator.warn(
-              "\n!!! RubyLLM's legacy acts_as API is deprecated and will be removed in RubyLLM 2.0.0. " \
-              "Please consult the migration guide at https://rubyllm.com/upgrading-to-1-7/\n"
-            )
-          end
+          require 'ruby_llm/active_record/acts_as'
+          ::ActiveRecord::Base.include RubyLLM::ActiveRecord::ActsAs
         end
       end
 

data/lib/ruby_llm/rerank.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # A Rerank is the result of ordering documents by relevance to a query.
+  # RubyLLM.rerank returns one:
+  #
+  #   rerank = RubyLLM.rerank("what is ruby",
+  #                           ["Ruby is a programming language", "Paris is in France"],
+  #                           model: "voyageai/rerank-2.5-lite", provider: :openrouter)
+  #   rerank.results.first.document # => "Ruby is a programming language"
+  #   rerank.results.first.score    # => 0.80078125
+  #
+  class Rerank
+    include Support::Inspectable
+    include Accounting::Usage::Result
+
+    # One reranked document: its +index+ in the input array, the
+    # +document+ text, and the provider's relevance +score+.
+    Result = Struct.new(:index, :document, :score, keyword_init: true)
+
+    # The reranked results, most relevant first, as an array of Result.
+    attr_reader :results
+
+    # The id of the model that ranked the documents, as a String.
+    attr_reader :model
+
+    # The raw provider response body.
+    attr_reader :raw
+
+    def initialize(results:, model:, raw: nil, input_tokens: nil, reported_cost: nil) # :nodoc:
+      @results = results
+      @model = model
+      @raw = raw
+      @input_tokens = input_tokens
+      @reported_cost = reported_cost
+    end
+
+    # Returns usage aggregated across every provider attempt.
+    def tokens
+      return ruby_llm_usage_tokens unless ruby_llm_usage_entries.empty?
+
+      Tokens.new(input: @input_tokens, reported_cost: @reported_cost)
+    end
+
+    # Returns the rerank cost across every provider attempt.
+    def cost
+      return ruby_llm_usage_cost unless ruby_llm_usage_entries.empty?
+
+      Cost.new(tokens:, model: model_info, category: :embeddings)
+    end
+
+    def model_info # :nodoc:
+      @model_info ||= RubyLLM.models.find(model)
+    rescue ModelNotFoundError
+      nil
+    end
+
+    def inspect_attributes # :nodoc:
+      { model: model, results: results.length }
+    end
+
+    # Ranks +documents+ (an array of strings) by relevance to +query+ and
+    # returns a Rerank. +model:+ is required; rerank catalogs are
+    # provider-specific and have no cross-provider default. +top_n:+
+    # limits how many results come back. +provider_options:+ merges
+    # options into the request in the provider's vocabulary.
+    #
+    #   RubyLLM.rerank "what is ruby", docs, model: "voyageai/rerank-2.5-lite", provider: :openrouter
+    #
+    def self.rerank(query, documents,
+                    model:,
+                    provider: nil,
+                    assume_model_exists: false,
+                    context: nil,
+                    top_n: nil,
+                    provider_options: {},
+                    metadata: nil)
+      config = context&.config || RubyLLM.config
+      model, provider_instance = Models.resolve(model, provider: provider, assume_model_exists: assume_model_exists,
+                                                       config: config)
+      empty_tokens = Tokens.new
+      payload = {
+        provider: provider_instance.slug,
+        provider_class: provider_instance.class.display_name,
+        model: model.id,
+        model_info: model,
+        query: query,
+        document_count: documents.length,
+        top_n: top_n,
+        provider_options: provider_options,
+        metadata: metadata,
+        tokens: empty_tokens,
+        cost: Cost.new(tokens: empty_tokens, model:, category: :embeddings)
+      }
+
+      RubyLLM.instrument('rerank.ruby_llm', payload, config: config) do |event|
+        result = provider_instance.rerank(query, documents, model:, top_n:, provider_options:)
+        event[:result] = result
+        event[:tokens] = result.tokens
+        event[:cost] = result.cost
+        result
+      end
+    end
+  end
+end