ruby_llm 2.0.0 → 2.1.0

data/lib/generators/ruby_llm/upgrade/upgrade_generator.rb

@@ -6,182 +6,25 @@ require_relative '../generator_helpers'
 
 module RubyLLM
   module Generators
-    # Upgrades the schema generated by RubyLLM 1.16 to RubyLLM 2.0.
+    # Upgrades the schema of an application on RubyLLM 2.0 to 2.1.
     class UpgradeGenerator < Rails::Generators::Base
       include Rails::Generators::Migration
       include RubyLLM::Generators::GeneratorHelpers
 
-      MODEL_MAPPING_DEFAULTS = {
-        chat: 'Chat',
-        message: 'Message',
-        model: 'Model',
-        tool_call: 'ToolCall'
-      }.freeze
-      MODEL_MAPPING_TYPES = MODEL_MAPPING_DEFAULTS.keys.map(&:to_s).freeze
-
       namespace 'ruby_llm:upgrade'
       source_root File.expand_path('templates', __dir__)
 
-      class_option :phase, type: :string, enum: %w[prepare backfill finish cleanup],
-                           desc: 'Generate one upgrade phase; cleanup runs in a later deployment'
-      class_option :mode, type: :string, enum: %w[rename copy], default: 'rename',
-                          desc: 'Copy prepares and backfills online, with protected 1.16 rollback until cleanup'
-      class_option :discard_incomplete_tool_calls, type: :boolean, default: false,
-                                                   desc: 'Discard all incomplete legacy tool calls at copy finish'
-
-      argument :model_mappings,
-               type: :array,
-               default: [],
-               banner: 'chat:ChatName message:MessageName model:ModelName tool_call:ToolCallName'
+      argument :model_mappings, type: :array, default: [], banner: 'chat:ChatName message:MessageName'
 
-      desc 'Upgrades a RubyLLM 1.16 Rails schema to 2.0'
+      desc 'Upgrades a RubyLLM 2.0 Rails schema to 2.1\n' \
+           'Usage: bin/rails g ruby_llm:upgrade [chat:ChatName] [message:MessageName]'
 
       def self.next_migration_number(dirname)
         ::ActiveRecord::Generators::Base.next_migration_number(dirname)
       end
 
-      def create_migration_files
-        validate_cleanup_options
-        parse_model_mappings(allowed_types: MODEL_MAPPING_TYPES, defaults: MODEL_MAPPING_DEFAULTS)
-        say_status :models, resolved_model_mappings
-
-        if copy_mode? && !options[:phase].in?(%w[backfill finish cleanup])
-          template 'ruby_llm_upgrade.rb.tt', 'app/models/concerns/ruby_llm_upgrade.rb'
-          template 'upgrade_initializer.rb.tt', 'config/initializers/ruby_llm_upgrade.rb'
-        end
-
-        phases = options[:phase] ? [options[:phase]] : %w[prepare backfill finish]
-        phases.each do |phase|
-          template = phase == 'backfill' ? 'backfill_v2_data' : "#{phase}_v2_upgrade"
-          migration = phase == 'backfill' ? 'backfill_ruby_llm_v2_data' : "#{phase}_ruby_llm_v2_upgrade"
-          migration_template "#{template}.rb.tt", "db/migrate/#{migration}.rb"
-        end
-      end
-
-      def show_next_steps
-        say_status :success, 'Upgrade prepared!', :green
-        if copy_mode?
-          say <<~INSTRUCTIONS
-
-            Copy mode retains the legacy tables and a separate chat model reference.
-            Install the generated RubyLLMUpgrade concern and initializer in BOTH the
-            running 1.16 build and the 2.0 build BEFORE running prepare. Restart
-            all affected 1.16 processes so they load the guards and explicit column
-            selects. The guards protect whole conversations
-            changed by 2.0. Direct SQL, bulk updates/deletes and attachment purges
-            bypass these guards; review those application paths before upgrading.
-
-            Rehearse prepare, backfill and finish on a production database copy.
-            Run prepare and backfill from the 2.0 build while 1.16 serves traffic.
-            Stop affected traffic, workers, scheduled jobs and retries BEFORE finish.
-            Finish catches up intervening writes, validates and activates 2.0.
-            Then load models and restart the 2.0 application.
-            Keep the old model/tool-call classes in the 1.16 build.
-
-            PostgreSQL, MySQL and SQLite support this workflow. Schema changes can
-            block writes; SQLite also serializes backfill and app writes.
-            Use a direct connection or session-mode pool for copy migrations.
-            db:migrate runs ALL pending phases. Stop at the backfill timestamp to
-            defer finish, or pause AI before running all three together.
-
-            With all affected processes stopped, use the 2.0 build to run:
-              bin/rails ruby_llm:upgrade:rollback  # activates the protected 1.16 view
-              bin/rails ruby_llm:upgrade:resume    # reconciles 1.16 writes, activates 2.0
-
-            Finish pending tool calls, approvals and batches before switching.
-            Once rollback is no longer needed, run ruby_llm:upgrade:finalize, then
-            generate --mode copy --phase cleanup in a later deployment. Keep the
-            compatibility files until cleanup succeeds; remove them afterward.
-
-            Pass the same --mode and model mappings to every generator invocation.
-            Full guide: https://rubyllm.com/next/upgrading/
-
-          INSTRUCTIONS
-          return
-        end
-        say <<~INSTRUCTIONS
-
-          These migrations expect the schema generated by RubyLLM 1.16. They move
-          RubyLLM's model and tool-call records, preserve message content and usage,
-          and retain the 1.x message columns until a later cleanup deployment.
-
-          Next steps:
-          1. Review the generated migrations
-          2. Rehearse them on a recent production snapshot
-          3. Stop affected web processes, workers, scheduled jobs, and retries
-          4. Run preparation, backfill, and finish before reopening traffic
-          5. Run: bin/rails ruby_llm:load_models
-          6. Keep only the current declarations in your application models:
-             #{chat_model_path}: #{acts_as_chat_declaration}
-             #{message_model_path}: #{acts_as_message_declaration}
-          7. Remove the 1.x model and tool-call classes, config.model_registry_class,
-             and config.use_new_acts_as
-          8. Update 1.x API calls using the upgrade guide
-
-          Generate one phase with --phase prepare, backfill, or finish when the
-          backfill needs to run outside your deployment's release phase.
-          After verifying the upgrade, generate --phase cleanup in a later release.
-          Pass the same model mappings to every invocation.
-
-          The migrations are irreversible. Retained columns do not make a code
-          rollback safe; recovery requires the previous database and application.
-
-          Full upgrade guide:
-             https://rubyllm.com/next/upgrading/
-
-        INSTRUCTIONS
-      end
-
-      private
-
-      def copy_mode? = options[:mode] == 'copy'
-      def discard_incomplete_tool_calls? = options[:discard_incomplete_tool_calls]
-
-      def validate_cleanup_options
-        return unless discard_incomplete_tool_calls?
-        return if copy_mode? && options[:phase].in?([nil, 'finish'])
-
-        raise Thor::Error, 'Discard incomplete tool calls only with --mode copy and the finish phase'
-      end
-
-      def copy_upgrade_settings
-        {
-          chat_table: chat_table_name, message_table: message_table_name,
-          model_table: v1_model_table_name, tool_call_table: v1_tool_call_table_name,
-          chat_class: chat_model_name, message_class: message_model_name,
-          model_foreign_key: v1_model_foreign_key, tool_call_foreign_key: v1_tool_call_foreign_key,
-          chat_foreign_key: chat_foreign_key, message_foreign_key: message_foreign_key
-        }
-      end
-
-      def prepare_migration_class_name = 'prepare_ruby_llm_v2_upgrade'.camelize
-      def backfill_migration_class_name = 'backfill_ruby_llm_v2_data'.camelize
-      def finish_migration_class_name = 'finish_ruby_llm_v2_upgrade'.camelize
-      def cleanup_migration_class_name = 'cleanup_ruby_llm_v2_upgrade'.camelize
-
-      def resolved_model_mappings
-        @model_names.map do |type, model_name|
-          "#{type}:#{model_name} -> #{table_name_for(model_name)}"
-        end.join(', ')
-      end
-
-      def v1_model_name(type)
-        @model_names.fetch(type)
-      end
-
-      def v1_model_table_name = table_name_for(v1_model_name(:model))
-      def v1_tool_call_table_name = table_name_for(v1_model_name(:tool_call))
-      def v1_model_foreign_key = "#{v1_model_table_name.singularize}_id"
-      def v1_tool_call_foreign_key = "#{v1_tool_call_table_name.singularize}_id"
-      def message_foreign_key = "#{message_table_name.singularize}_id"
-      def chat_foreign_key = "#{chat_table_name.singularize}_id"
-
-      def chat_model_path
-        "app/models/#{chat_model_name.underscore}.rb"
-      end
-
-      def message_model_path
-        "app/models/#{message_model_name.underscore}.rb"
+      def create_migration_file
+        migration_template 'upgrade_ruby_llm_to_2_1.rb.tt', 'db/migrate/upgrade_ruby_llm_to_2_1.rb'
       end
     end
   end

data/lib/ruby_llm/accounting/usage.rb

@@ -6,19 +6,55 @@ module RubyLLM
   module Accounting # :nodoc:
     # Internal accounting for physical provider attempts.
     class Usage # :nodoc: all
-      def self.instrument(entry, config:) # :nodoc:
-        RubyLLM.instrument(
-          'usage.ruby_llm',
-          {
-            operation: entry.operation,
-            provider: entry.provider,
-            model: entry.model,
-            status: entry.status,
-            tokens: entry.tokens,
-            cost: entry.cost
-          },
-          config:
-        )
+      OWNER_KEY = :ruby_llm_usage_owner
+
+      class << self
+        # Persists finished entries that no chat records, when the
+        # application provides a store, such as the Rails usage ledger.
+        attr_accessor :ledger
+
+        # The owner lives in fiber storage, which isolates concurrent
+        # fibers and threads and is inherited by those started inside it.
+        def owner
+          Fiber[OWNER_KEY]
+        end
+
+        def with_owner(owner)
+          previous = Fiber[OWNER_KEY]
+          Fiber[OWNER_KEY] = owner
+          yield
+        ensure
+          Fiber[OWNER_KEY] = previous
+        end
+
+        def owned_by(owner, &)
+          owner.nil? ? yield : with_owner(owner, &)
+        end
+
+        def record(entry)
+          ledger&.record(entry)
+        end
+
+        def report(entry, config:)
+          record(entry)
+          instrument(entry, config:)
+        end
+
+        def instrument(entry, config:)
+          RubyLLM.instrument(
+            'usage.ruby_llm',
+            {
+              operation: entry.operation,
+              provider: entry.provider,
+              model: entry.model,
+              status: entry.status,
+              tokens: entry.tokens,
+              cost: entry.cost,
+              owner: entry.owner
+            },
+            config:
+          )
+        end
       end
 
       # Shared storage for result objects carrying attempt accounting.
@@ -48,10 +84,10 @@ module RubyLLM
 
       # One physical provider request attempt and its accounting facts.
       class Entry
-        OPERATIONS = %i[chat embedding moderation image speech transcription ocr rerank].freeze
+        OPERATIONS = %i[chat embedding moderation image speech transcription ocr rerank judgment video research].freeze
         STATUSES = %i[pending succeeded failed cancelled].freeze
 
-        attr_reader :operation, :provider, :model, :status, :tokens, :cost
+        attr_reader :operation, :provider, :model, :status, :tokens, :owner
         attr_accessor :message
 
         include Support::Inspectable
@@ -60,17 +96,23 @@ module RubyLLM
           { operation: operation, provider: provider, model: model, status: status, tokens: tokens }
         end
 
-        def initialize(operation:, provider:, model:, status: :pending, tokens: nil, cost: nil, message: nil)
+        def initialize(operation:, provider:, model:, status: :pending, tokens: nil, cost: nil, message: nil,
+                       owner: nil)
           @operation = operation.to_sym
           @provider = provider.to_s
           @model = model&.to_s
           @status = status.to_sym
           @tokens = tokens || Tokens.new
-          @cost = cost || Cost.new(tokens: @tokens)
+          @cost = cost
           @message = message
+          @owner = owner
           validate!
         end
 
+        def cost
+          @cost ||= Cost.new(tokens:)
+        end
+
         def pending? = status == :pending
         def succeeded? = status == :succeeded
         def failed? = status == :failed
@@ -92,7 +134,7 @@ module RubyLLM
         def finish(status:, tokens: nil, cost: nil) # :nodoc:
           @status = status.to_sym
           @tokens = tokens || Tokens.new
-          @cost = cost || Cost.new(tokens: @tokens)
+          @cost = cost
           validate!
           self
         end
@@ -115,15 +157,17 @@ module RubyLLM
           speech: :audio_tokens,
           transcription: :audio_tokens,
           ocr: :text_tokens,
-          rerank: :embeddings
+          rerank: :embeddings,
+          judgment: :text_tokens
         }.freeze
 
         attr_reader :entries
 
-        def initialize(operation:, provider:, model:, config:, on_finish: nil)
+        def initialize(operation:, provider:, model:, config:, on_finish: nil, model_id: nil)
           @operation = operation.to_sym
           @provider = provider
           @model_info = model
+          @model_id = model_id
           @config = config
           @on_finish = on_finish
           @entries = []
@@ -131,21 +175,25 @@ module RubyLLM
         end
 
         def start
-          entry = Entry.new(operation: @operation, provider: @provider.slug, model: @model_info&.id)
+          entry = Entry.new(operation: @operation, provider: @provider.slug, model: @model_id || @model_info&.id,
+                            owner: Usage.owner)
           @entries << entry
           @pending << entry
           entry
         end
 
         def observe(chunk)
+          observe_tokens(chunk.tokens) if chunk.respond_to?(:tokens)
+        end
+
+        # Keeps the usage a response reported for the current attempt, so an
+        # attempt that fails after the provider answered, such as a blocked
+        # one, still records what it billed.
+        def observe_tokens(tokens)
           entry = @pending.last
           return unless entry
-          return unless chunk.respond_to?(:tokens)
 
-          entry.finish(
-            status: :pending,
-            tokens: merge_stream_tokens(entry.tokens, chunk.tokens)
-          )
+          entry.finish(status: :pending, tokens: entry.tokens.merge(tokens))
         end
 
         def fail_attempt(entry, error)
@@ -162,8 +210,9 @@ module RubyLLM
         def succeed(result)
           # A request that produced several results, like a multi-image
           # generation, is billed once: the first result carries the call.
-          billed = result.is_a?(Array) ? result.first : result
-          billed.model_info = message_model(billed) if billed.is_a?(Message)
+          results = result.is_a?(Array) ? result : [result]
+          results.each { |item| item.model_info = result_model(item) if item.respond_to?(:model_info=) }
+          billed = results.first
           pending = @pending.dup
           if pending.empty?
             attach_to_result(billed)
@@ -188,10 +237,10 @@ module RubyLLM
 
         private
 
-        def message_model(message)
-          return @model_info if message.model.nil? || message.model == @model_info&.id
+        def result_model(result)
+          return @model_info if result.model.nil? || result.model == @model_info&.id
 
-          RubyLLM.models.find(message.model, provider: @provider.slug, config: @config)
+          RubyLLM.models.find(result.model, provider: @provider.slug, config: @config)
         rescue ModelNotFoundError
           @model_info
         end
@@ -225,29 +274,17 @@ module RubyLLM
           )
           entry.finish(status:, tokens:, cost:)
           @pending.delete(entry)
-          @on_finish&.call(entry)
-          instrument(entry)
-        end
-
-        def instrument(entry)
-          Accounting::Usage.instrument(entry, config: @config)
+          if @on_finish
+            @on_finish.call(entry)
+            Accounting::Usage.instrument(entry, config: @config)
+          else
+            Accounting::Usage.report(entry, config: @config)
+          end
         end
 
         def attach_to_result(result)
           result&.ruby_llm_usage_entries = entries
         end
-
-        def merge_stream_tokens(existing, incoming)
-          Tokens.new(
-            input: incoming.input.nil? ? existing.input : incoming.input,
-            output: incoming.output.nil? ? existing.output : incoming.output,
-            cache_read: incoming.cache_read.nil? ? existing.cache_read : incoming.cache_read,
-            cache_write: incoming.cache_write.nil? ? existing.cache_write : incoming.cache_write,
-            thinking: incoming.thinking.nil? ? existing.thinking : incoming.thinking,
-            server_tool_use: incoming.server_tool_use.nil? ? existing.server_tool_use : incoming.server_tool_use,
-            reported_cost: incoming.reported_cost.nil? ? existing.reported_cost : incoming.reported_cost
-          )
-        end
       end
     end
   end

data/lib/ruby_llm/active_record/acts_as.rb

@@ -104,6 +104,7 @@ module RubyLLM
                    -> { chronological },
                    as: :chat,
                    class_name: 'RubyLLM::ActiveRecord::Usage',
+                   inverse_of: :chat,
                    dependent: :destroy
         end
 
@@ -120,6 +121,7 @@ module RubyLLM
                      touch: touch_chat
 
           has_many :ruby_llm_tool_calls,
+                   -> { order(:created_at, :id) },
                    as: :message,
                    class_name: 'RubyLLM::ActiveRecord::ToolCall',
                    dependent: :destroy
@@ -133,6 +135,7 @@ module RubyLLM
                    -> { chronological },
                    as: :message,
                    class_name: 'RubyLLM::ActiveRecord::Usage',
+                   inverse_of: :message,
                    dependent: :nullify
         end
       end

data/lib/ruby_llm/active_record/attachment_helpers.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require 'active_support/core_ext/object/blank'
-require 'tempfile'
 
 module RubyLLM
   module ActiveRecord
@@ -12,7 +11,27 @@ module RubyLLM
         return unless message_record.respond_to?(:attachments)
 
         attachables = prepare_for_active_storage(attachments)
-        message_record.attachments.attach(attachables) if attachables.any?
+        return if attachables.empty?
+
+        message_record.attachments.attach(attachables)
+        Support::Utils.to_safe_array(attachments).grep(RubyLLM::Attachment).each do |attachment|
+          attachment.provider_file_store ||= ProviderFile.store { persisted_blob_key(message_record, attachment) }
+        end
+      end
+
+      # Active Storage records the size and MD5 of every blob, so the
+      # message's blob of the same size holds the attachment's bytes. Only
+      # attachments of equal size need the MD5 to tell them apart.
+      def persisted_blob_key(message_record, attachment)
+        blob = active_storage_blobs(attachment.source) if attachment.active_storage?
+        return blob.key if blob.is_a?(ActiveStorage::Blob)
+
+        content = attachment.content
+        candidates = message_record.attachments.blobs.where(byte_size: content.bytesize).to_a
+        return candidates.first&.key unless candidates.many?
+
+        checksum = OpenSSL::Digest::MD5.base64digest(content)
+        candidates.find { |candidate| candidate.checksum == checksum }&.key
       end
 
       def prepare_for_active_storage(attachments)
@@ -41,7 +60,8 @@ module RubyLLM
           {
             io: StringIO.new(attachment.content),
             filename: attachment.filename,
-            content_type: attachment.mime_type
+            content_type: attachment.mime_type,
+            metadata: { resolution: attachment.resolution&.to_s }.compact
           }
         end
       rescue StandardError => e
@@ -162,7 +182,23 @@ module RubyLLM
       end
 
       def active_storage_attachments?
-        respond_to?(:attachments) && attachments.attached?
+        return false unless respond_to?(:attachments)
+
+        preloaded = preloaded_attachments
+        preloaded ? preloaded.any? : attachments.attached?
+      end
+
+      def preloaded_attachments
+        return if pending_attachment_change || !self.class.reflect_on_association(:attachments_attachments)
+
+        preloaded_records(:attachments_attachments)
+      end
+
+      # Building an association proxy evaluates the association's scope,
+      # which costs more than reading the records a preload already loaded.
+      def preloaded_records(name)
+        association = association(name)
+        association.target if association.loaded?
       end
 
       def collect_attachments(action_text_attachments)
@@ -194,12 +230,13 @@ module RubyLLM
       end
 
       def stored_attachment(attachment, attachable)
-        if pending_upload_attachable?(attachable)
-          pending_upload_attachment(attachable)
-        else
-          tempfile = download_attachment(attachment)
-          RubyLLM::Attachment.new(tempfile, filename: attachment.filename.to_s)
-        end
+        stored = if pending_upload_attachable?(attachable)
+                   pending_upload_attachment(attachable)
+                 else
+                   RubyLLM::Attachment.new(attachment, resolution: attachment.metadata['resolution']&.to_sym)
+                 end
+        stored.provider_file_store = ProviderFile.store { attachment.blob&.key }
+        stored
       end
 
       def pending_upload_attachable?(attachable)
@@ -225,7 +262,9 @@ module RubyLLM
 
       def pending_upload_attachment(attachable)
         if attachable.is_a?(Hash)
-          RubyLLM::Attachment.new(attachment_hash_io(attachable), filename: attachment_hash_filename(attachable))
+          RubyLLM::Attachment.new(attachment_hash_io(attachable),
+                                  filename: attachment_hash_filename(attachable),
+                                  resolution: attachable.dig(:metadata, :resolution)&.to_sym)
         else
           RubyLLM::Attachment.new(attachable)
         end
@@ -245,20 +284,6 @@ module RubyLLM
       rescue NameError
         false
       end
-
-      def download_attachment(attachment)
-        ext = File.extname(attachment.filename.to_s)
-        basename = File.basename(attachment.filename.to_s, ext)
-        tempfile = Tempfile.new([basename, ext])
-        tempfile.binmode
-
-        attachment.download { |chunk| tempfile.write(chunk) }
-
-        tempfile.flush
-        tempfile.rewind
-        @_tempfiles << tempfile
-        tempfile
-      end
     end
   end
 end

data/lib/ruby_llm/active_record/batch.rb

@@ -3,7 +3,7 @@
 module RubyLLM
   module ActiveRecord
     # RubyLLM's private persistence for provider-side batches.
-    class Batch < ::ActiveRecord::Base # :nodoc:
+    class Batch < Record # :nodoc:
       self.table_name = 'ruby_llm_batches'
 
       STATUSES = %w[pending succeeded failed cancelled].freeze