ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/active_record/message_methods.rb

@@ -17,9 +17,23 @@ module RubyLLM
         send(chat_association_name)
       end
 
-      # Converts this record to a RubyLLM::Message.
+      # Returns an assistant RubyLLM::Chunk with +content+ at position
+      # +index+ that streams into this message, as the chunks a persisted
+      # chat yields do, so it answers +message+ and +broadcast+. Use it to
+      # stand in for a streamed reply in tests:
+      #
+      #   chunk = message.chunk("Hello", index: 0)
+      #   chunk.broadcast
+      #
+      def chunk(content, index:)
+        ChunkMethods.attach(RubyLLM::Chunk.new(role: :assistant, content:, index:), self)
+      end
+
+      # Converts this record to a RubyLLM::Message. The message carries the
+      # recorded cost, so RubyLLM::Message#to_h exports it.
       def to_llm
         entries = ruby_llm_usage_entries
+        succeeded = entries.reverse.find(&:succeeded?)
         RubyLLM::Message.new(
           role: role.to_sym,
           content: extract_content,
@@ -29,17 +43,29 @@ module RubyLLM
           server_tool_calls: server_tool_calls,
           raw_content: optional_column(:raw_content),
           raw_reasoning: optional_column(:raw_reasoning),
+          phase: optional_column(:phase),
+          continuation: optional_column(:continuation),
           usage_entries: entries,
+          cost: (cost if entries.any?),
           tool_calls: tool_calls,
           tool_call_id: parent_tool_call&.id,
+          provider_call_id: parent_tool_call&.provider_call_id,
           mcp_result: mcp_result,
           finish_reason: optional_column(:finish_reason),
-          model: entries.reverse.find(&:succeeded?)&.model,
+          model: succeeded&.model || stand_in_column(:model_id),
+          provider: (stand_in_column(:provider) unless succeeded),
           cache_until_here: cache_until_here?,
           cache_ttl: cache_ttl
         )
       end
 
+      # The fallback a reply with no successful attempt, such as a cancelled
+      # one, came from. On other roles these columns are the application's.
+      def stand_in_column(name)
+        optional_column(name) if role.to_s == 'assistant'
+      end
+      private :stand_in_column
+
       def ruby_llm_usage_entries # :nodoc:
         ruby_llm_usage_records.map(&:to_entry)
       end
@@ -53,6 +79,20 @@ module RubyLLM
         to_llm.model
       end
 
+      # Returns the model ID from the last successful attempt, as #model
+      # does, or a +model_id+ column the application keeps, or +nil+.
+      def model_id
+        to_llm.model_id || optional_column(:model_id)
+      end
+
+      # Returns the slug of the provider that produced the message, from its
+      # last successful attempt, or a +provider+ column the application keeps,
+      # or +nil+. RubyLLM writes +model_id+ and +provider+ columns, when the
+      # table has them, on replies a fallback model wrote.
+      def provider
+        to_llm.provider || optional_column(:provider)
+      end
+
       # Returns the model from the registry, or +nil+ if it is unknown.
       def model_info
         to_llm.model_info
@@ -60,11 +100,13 @@ module RubyLLM
 
       # Marks the message as a prompt cache boundary and returns it. Pass
       # +ttl:+ to give the boundary its own cache lifetime, as
-      # RubyLLM::Message#cache_until_here describes.
+      # RubyLLM::Message#cache_until_here describes. The chat that returned
+      # the message sends the boundary with its next request.
       def cache_until_here(ttl: nil)
         attributes = { cache_until_here: true }
         attributes[:cache_ttl] = ttl if ttl || has_attribute?(:cache_ttl)
         update!(attributes)
+        association(chat_association_name).target&.refresh_messages
         self
       end
 
@@ -73,6 +115,12 @@ module RubyLLM
         optional_column(:cache_until_here) || false
       end
 
+      # Serialization reads each column through its method, and
+      # #cache_until_here marks a boundary instead of reading one.
+      def read_attribute_for_serialization(name) # :nodoc:
+        name.to_s == 'cache_until_here' ? cache_until_here? : super
+      end
+
       # Returns the cache lifetime given to the message's boundary, or +nil+
       # when the chat's caching policy applies.
       def cache_ttl
@@ -97,6 +145,12 @@ module RubyLLM
         Array(optional_column(:server_tool_calls)).map { |call| RubyLLM::ServerToolCall.from_h(call) }
       end
 
+      # Returns the pages a hosted web search returned, as an array of
+      # RubyLLM::WebSource. See RubyLLM::Message#web_sources.
+      def web_sources
+        to_llm.web_sources
+      end
+
       # Returns the token usage across every attempt that produced the message.
       def tokens
         RubyLLM::Tokens.aggregate(ruby_llm_usages.map(&:tokens))
@@ -104,14 +158,15 @@ module RubyLLM
 
       # Returns the cost across every attempt that produced the message, as a RubyLLM::Cost.
       def cost
-        records = ruby_llm_usages.to_a
+        records = ruby_llm_usage_records
         RubyLLM::Cost.aggregate(records.map(&:cost), complete: records.all?(&:cost_available?))
       end
 
-      # Returns the tool calls as RubyLLM::ToolCall values keyed by call id.
+      # Returns the tool calls, as RubyLLM::Message#tool_calls does: an
+      # Enumerable of RubyLLM::ToolCall objects, also readable by call ID.
       def tool_calls
         records = preloaded_records(:ruby_llm_tool_calls) || ruby_llm_tool_calls
-        records.to_h { |record| [record.tool_call_id, record.to_llm] }
+        RubyLLM::ToolCall::Collection.new(records.to_h { |record| [record.tool_call_id, record.to_llm] })
       end
 
       # Returns the tool call this message answers, or +nil+.
@@ -162,6 +217,24 @@ module RubyLLM
         payload_error_message(extract_content)
       end
 
+      # Returns whether a tool ended the turn with this result. See
+      # RubyLLM::Message#halted?.
+      def halted?
+        to_llm.halted?
+      end
+
+      # Returns the path of the file holding the whole of a tool result too
+      # long to send whole, or +nil+. See RubyLLM::Message#spilled_path.
+      def spilled_path
+        to_llm.spilled_path
+      end
+
+      # Returns whether this is the part of a streamed reply the chat kept
+      # when it was cancelled. See RubyLLM::Message#cancelled?.
+      def cancelled?
+        to_llm.cancelled?
+      end
+
       # Returns why the model stopped, as the Symbol RubyLLM::Message#finish_reason reports.
       def finish_reason
         optional_column(:finish_reason)&.to_sym
@@ -176,6 +249,7 @@ module RubyLLM
       def max_tokens?
         to_llm.max_tokens?
       end
+      alias truncated? max_tokens?
 
       # Returns +true+ if the model stopped to request tool calls.
       def tool_call_stop?
@@ -187,12 +261,10 @@ module RubyLLM
         to_llm.content_filtered?
       end
 
-      # Returns the content parsed as JSON, or +nil+ when the message has no text.
+      # Returns the content parsed as JSON, or +nil+ when the message has no
+      # text. Reads only the final answer, as RubyLLM::Message#parsed does.
       def parsed
-        text = extract_content
-        return if text.nil? || text.empty?
-
-        JSON.parse(text)
+        to_llm.parsed
       end
 
       private

data/lib/ruby_llm/active_record/model.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'active_support/core_ext/module/delegation'
+require 'concurrent/set'
 
 module RubyLLM
   module ActiveRecord
@@ -18,6 +19,7 @@ module RubyLLM
       scope :unlisted, -> { where.not(unlisted_at: nil) }
 
       UNLISTED_WARNING_LIMIT = 5 # :nodoc:
+      STORED_REGISTRIES = Concurrent::Set.new # :nodoc:
 
       class << self
         def read # :nodoc:
@@ -42,16 +44,28 @@ module RubyLLM
         end
 
         def save_to_database(registry = RubyLLM.models) # :nodoc:
-          transaction do
-            kept = registry.all.map do |model_info|
-              model = find_or_initialize_by(model_id: model_info.id, provider: model_info.provider)
-              model.update!(attributes_from_llm(model_info))
-              model.id
-            end
-            unlist(kept)
+          rows = registry.all.index_by { |model_info| [model_info.provider, model_info.id] }
+                         .transform_values { |model_info| attributes_from_llm(model_info) }
+          validate_incomplete(rows.values)
+
+          transaction(requires_new: true) do
+            upsert_all(rows.values, unique_by: registry_key) if rows.any?
+            unlist(pluck(:id, :provider, :model_id).reject { |_id, *key| rows.key?(key) }.map(&:first))
           end
         end
 
+        # A database stays marked once a chat that loaded the registry into it
+        # has committed. A transactional test rolls that load back after every
+        # example, and loading it again on each empty table would repeat the
+        # whole catalog in every test.
+        def registry_stored? # :nodoc:
+          STORED_REGISTRIES.include?(connection_db_config)
+        end
+
+        def registry_stored!(db_config) # :nodoc:
+          STORED_REGISTRIES << db_config
+        end
+
         def from_llm(model_info) # :nodoc:
           new(attributes_from_llm(model_info))
         end
@@ -62,9 +76,9 @@ module RubyLLM
         # away. A row an application record points at cannot: deleting it would
         # dangle the reference. That row is stamped unlisted instead, keeping
         # the chats that use it resolvable while the application migrates them.
-        def unlist(kept_ids)
+        def unlist(dropped_ids)
           stayed = []
-          where.not(id: kept_ids).find_each do |model|
+          where(id: dropped_ids).find_each do |model|
             transaction(requires_new: true) { model.destroy! }
           rescue ::ActiveRecord::InvalidForeignKey
             model.update!(unlisted_at: Time.current) unless model.unlisted_at
@@ -90,6 +104,20 @@ module RubyLLM
           extra.positive? ? "#{listing}, and #{extra} more" : listing
         end
 
+        # upsert_all skips validations. An entry missing a required value goes
+        # through them on its own, so it raises RecordInvalid as a save would.
+        def validate_incomplete(rows)
+          rows.each do |attributes|
+            new(attributes).validate! if attributes.values_at(:model_id, :provider, :name).any?(&:blank?)
+          end
+        end
+
+        # Active Record's MySQL adapter rejects a named conflict target;
+        # MySQL matches every unique index on its own.
+        def registry_key
+          %i[provider model_id] if connection_pool.with_connection(&:supports_insert_conflict_target?)
+        end
+
         def attributes_from_llm(model_info)
           {
             model_id: model_info.id,

data/lib/ruby_llm/active_record/provider_file.rb

@@ -7,20 +7,37 @@ module RubyLLM
     # already uploaded instead of uploading the blob again. Rows name the
     # blob by its key, which Active Storage never hands out twice, so a row
     # left by a deleted blob can never match a newer one.
-    class ProviderFile < Record # :nodoc:
+    #
+    #   RubyLLM::ActiveRecord::ProviderFile.where(blob_key: document.file.key).pluck(:provider, :file_id)
+    #
+    # Its columns:
+    #
+    # [+blob_key+] The key of the Active Storage blob.
+    # [+provider+] The provider's slug.
+    # [+account+] A digest that identifies the provider account the file was uploaded to.
+    # [+file_id+] The provider's id for the uploaded file.
+    # [+expires_at+] When the provider deletes the file, or +nil+.
+    #
+    # RubyLLM writes these rows. Read and query them, and add scopes,
+    # columns, or encryption of your own, but do not create or update them.
+    # Encrypt +file_id+ deterministically, since RubyLLM queries it.
+    class ProviderFile < Record
       self.table_name = 'ruby_llm_provider_files'
 
       validates :blob_key, :provider, :account, :file_id, presence: true
 
       # Returns a store for the blob whose key the block returns, or +nil+
       # while the table has not been migrated.
-      def self.store(&)
+      def self.store(&) # :nodoc:
         Store.new(&) if table_exists?
       end
 
       # Deletes the uploads recorded for +blob+ as Active Storage destroys it.
-      def self.forget_blob(blob)
-        where(blob_key: blob.key).delete_all if table_exists?
+      # Reads the stored key: Blob#key generates one when it is missing, and
+      # a destroyed blob is frozen.
+      def self.forget_blob(blob) # :nodoc:
+        key = blob[:key]
+        where(blob_key: key).delete_all if key && table_exists?
       end
 
       # The uploads of one blob. The block finds the blob the first time a

data/lib/ruby_llm/active_record/queued_message.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module ActiveRecord
+    # A user message queued on a persisted chat with ChatMethods#queue,
+    # waiting in RubyLLM's +ruby_llm_queued_messages+ table until the
+    # conversation loop takes it. Any process can queue one, and the
+    # process that runs the chat takes it, so a request can add to a reply
+    # a job is running. Its attachments move to the message it becomes.
+    #
+    #   chat.queued_messages.each { |queued| render queued }
+    #   chat.queued_messages.find(params[:id]).withdraw
+    #
+    # It answers +content+, +attachments+, +sender+, +steer?+, and
+    # +withdraw+, like RubyLLM::QueuedMessage.
+    class QueuedMessage < Record
+      self.table_name = 'ruby_llm_queued_messages'
+
+      belongs_to :chat, polymorphic: true
+      belongs_to :sender, polymorphic: true, optional: true
+
+      # Declares the attachments when Active Storage is loaded. Called by
+      # +acts_as_chat+, since Active Storage loads after RubyLLM's records.
+      def self.declare_attachments # :nodoc:
+        return unless respond_to?(:has_many_attached) && !reflect_on_attachment(:attachments)
+
+        has_many_attached :attachments
+      end
+
+      # Returns when the message was queued.
+      def queued_at
+        created_at
+      end
+
+      # Deletes the record and its attachments before the conversation
+      # loop takes it, from any process. Returns +true+, or +false+ when it
+      # is no longer queued: the loop took it into the conversation, or it
+      # was withdrawn already.
+      #
+      #   chat.queued_messages.find(params[:id]).withdraw # => true
+      #
+      def withdraw
+        return false unless self.class.where(id:).delete_all == 1
+
+        attachments.purge_later if respond_to?(:attachments) && attachments.attached?
+        true
+      end
+    end
+  end
+end

data/lib/ruby_llm/active_record/record.rb

@@ -2,7 +2,9 @@
 
 module RubyLLM
   module ActiveRecord
-    # Abstract base for RubyLLM's model, tool-call, usage, and batch records.
+    # Abstract base for the records of the tables RubyLLM owns: Model,
+    # ToolCall, Usage, Batch, MCPCredential, ProviderFile, QueuedMessage,
+    # and ChatFile.
     # Inherits the application's default connection unless configured otherwise.
     # Keep these records on the same database and connection pool as chats and
     # messages to preserve foreign keys and shared transactions.