ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/message.rb

@@ -34,11 +34,37 @@ module RubyLLM
 
     # The ID of the model that produced the message, +nil+ on user messages.
     attr_reader :model
+    alias model_id model
 
     attr_writer :model_info # :nodoc:
 
-    # The tool calls the assistant requested, as a Hash of ToolCall objects
-    # keyed by call ID, or +nil+.
+    # Returns the slug of the provider that produced the message, from its
+    # last successful attempt, or +nil+ when none is recorded, as on user
+    # messages. A reply from a model given to Chat#with_fallbacks names that
+    # model's provider.
+    def provider
+      @provider || ruby_llm_usage_entries.reverse.find(&:succeeded?)&.provider
+    end
+
+    # Marks the message as the work of a fallback model.
+    def stood_in(model:, provider:) # :nodoc:
+      @model ||= model
+      @provider = provider
+      @stood_in = true
+      self
+    end
+
+    def stood_in? # :nodoc:
+      @stood_in || false
+    end
+
+    # The tool calls the assistant requested, as an Enumerable of ToolCall
+    # objects in provider order, also readable by call ID. Empty when there
+    # are none.
+    #
+    #   response.tool_calls.map(&:name) # => ["weather"]
+    #   response.tool_calls["call_1"]   # => #<RubyLLM::ToolCall ...>
+    #
     attr_reader :tool_calls
 
     # The ID of the tool call this message answers. Set only on tool result
@@ -59,9 +85,16 @@ module RubyLLM
 
     # Why the model stopped: +:stop+, +:max_tokens+, +:tool_calls+, or
     # +:content_filter+. Any other reason comes through as the provider
-    # spelled it, such as Anthropic's +:pause_turn+.
+    # named it, as a lowercase Symbol, such as Anthropic's +:pause_turn+ or
+    # Gemini's +:malformed_function_call+.
     attr_reader :finish_reason
 
+    # The ContentFilter verdicts the provider reported on the request and
+    # this response, or an empty Array. Read them to see which filter ended
+    # a response that is #content_filtered?. Only the live response carries
+    # them: #to_h leaves them out, so a persisted chat never stores them.
+    attr_reader :content_filters
+
     # The provider-executed tool steps in this response, as an array of
     # ServerToolCall objects. Empty unless the chat enabled tools with
     # Chat#with_provider_tools and the model used one.
@@ -82,6 +115,16 @@ module RubyLLM
     # +nil+ when the provider returned none.
     attr_reader :raw_reasoning # :nodoc:
 
+    # The provider's label for this assistant text within its turn, such as
+    # OpenAI's +commentary+ or +final_answer+, kept so later requests can
+    # send it back. +nil+ when the provider gave none.
+    attr_reader :phase # :nodoc:
+
+    # The provider's opaque state for an assistant message it has not
+    # finished, which a later request sends back to continue it. +nil+ once
+    # the message is complete.
+    attr_reader :continuation # :nodoc:
+
     # The Chat this message belongs to, set when it is added to a
     # conversation. Backs #tool_results.
     attr_accessor :conversation # :nodoc:
@@ -93,8 +136,10 @@ module RubyLLM
       @config = options[:config]
       @attachments = Attachment.wrap(options[:attachments], config: @config)
       @model = options[:model]
+      @provider = options[:provider]
       @supplied_cost = coerce_value(options[:cost], Cost)
       @tool_call_id = options[:tool_call_id]
+      @provider_call_id = options[:provider_call_id] || options[:provider_tool_call_id]
       @tokens = options[:tokens] || Tokens.new(
         input: options[:input_tokens],
         output: options[:output_tokens],
@@ -111,8 +156,11 @@ module RubyLLM
       @server_tool_calls = Array(options[:server_tool_calls]).map { |call| coerce_value(call, ServerToolCall) }
       @raw_content = options[:raw_content]
       @raw_reasoning = options[:raw_reasoning]
+      @phase = options[:phase]
+      @continuation = options[:continuation]
       @mcp_result = coerce_mcp_result(options[:mcp_result])
       @finish_reason = options[:finish_reason]&.to_sym
+      @content_filters = coerce_list(options[:content_filters], ContentFilter)
       self.ruby_llm_usage_entries = options[:usage_entries] if options[:usage_entries]
       @cache_until_here = options.fetch(:cache_until_here, false)
       @cache_ttl = options[:cache_ttl] if @cache_until_here
@@ -121,15 +169,99 @@ module RubyLLM
     end
 
     # Returns #content parsed as JSON, memoized after the first call.
-    # Useful for reading structured output responses.
+    # Useful for reading structured output responses. When the provider
+    # marks part of the turn as commentary before the final answer, only
+    # the final answer is parsed.
     #
     #   response = chat.with_schema(PersonSchema).ask "Generate a person"
     #   response.parsed # => {"name" => "Alice", "age" => 30}
     #
     def parsed
-      return if content.nil? || content.empty?
+      text = answer_text
+      return if text.nil? || text.empty?
+
+      @parsed ||= JSON.parse(text)
+    end
+
+    # Returns the pages a hosted web search returned for this message, as an
+    # array of WebSource objects with each address once. Reads what the message
+    # keeps, so a restored message answers the same. Empty when the message
+    # ran no search or its provider reports no sources.
+    #
+    #   response = chat.with_provider_tools(:web_search).ask "What changed in Ruby 4.0?"
+    #   response.web_sources.map(&:url)
+    #
+    # Sources come from the outside world and are not a safe allowlist on their
+    # own: a result can be attacker-controlled, and OpenAI can list a page the
+    # model chose to open beside what its search found.
+    #
+    def web_sources
+      return [] if server_tool_calls.empty? || provider.nil?
+
+      Array(Provider.resolve(provider)&.web_sources(self)).uniq(&:url)
+    end
+
+    # Returns a copy of the message with the given fields changed, like
+    # Data#with. The message itself is left alone. Change +content+ or
+    # +attachments+:
+    #
+    #   message.with(content: "[left out]")
+    #
+    # New content also drops what holds the old text: on an assistant
+    # message its thinking, citations, server tool calls, raw response, and
+    # native response blocks, and on a tool result its MCP result. Raises
+    # ArgumentError for another field.
+    def with(**changes)
+      unknown = changes.keys - %i[content attachments]
+      raise ArgumentError, "Message#with cannot change #{unknown.join(', ')}" if unknown.any?
+
+      changes.reduce(self) do |message, (field, value)|
+        field == :content ? message.with_content(value) : message.with_attachments(value)
+      end
+    end
+
+    def with_content(content) # :nodoc:
+      (role == :assistant ? without_native_content : dup).tap do |message|
+        message.instance_variable_set(:@content, normalize_content(content))
+        message.instance_variable_set(:@mcp_result, nil)
+        if role == :assistant
+          message.instance_variable_set(:@citations, [])
+          message.instance_variable_set(:@server_tool_calls, [])
+          message.instance_variable_set(:@raw, nil)
+        end
+        message.remove_instance_variable(:@parsed) if message.instance_variable_defined?(:@parsed)
+      end
+    end
+
+    # Returns whether a tool ended the turn with this result, with
+    # Tool#halt. Its finish reason is +:halted+.
+    def halted?
+      finish_reason == :halted
+    end
 
-      @parsed ||= JSON.parse(content)
+    # Returns whether this is the part of a streamed reply the chat kept
+    # when it was cancelled (see Chat#cancel), or the result that closed a
+    # tool call it left unfinished. Its finish reason is +:cancelled+.
+    def cancelled?
+      finish_reason == :cancelled
+    end
+
+    # Returns the path, in the chat's files, of the whole of a tool result
+    # too long to send whole, or +nil+ when the result went whole. A chat
+    # with a file store keeps such a result in a file and sends the model
+    # its start and this path (see Chat#with_file_store).
+    #
+    #   if (path = message.spilled_path)
+    #     chat.file_store.read(path)
+    #   end
+    #
+    def spilled_path
+      Chat::FileReferences.spilled_path(content) if tool_result?
+    end
+
+    def mark_cancelled # :nodoc:
+      @finish_reason = :cancelled
+      self
     end
 
     def with_attachments(attachments) # :nodoc:
@@ -137,19 +269,29 @@ module RubyLLM
       dup.tap { |message| message.instance_variable_set(:@attachments, wrapped) }
     end
 
+    def with_tool_calls(tool_calls) # :nodoc:
+      dup.tap { |message| message.instance_variable_set(:@tool_calls, coerce_tool_calls(tool_calls)) }
+    end
+
+    def with_tool_call_id(tool_call_id) # :nodoc:
+      dup.tap { |message| message.instance_variable_set(:@tool_call_id, tool_call_id) }
+    end
+
     def without_native_content(raw_content: nil) # :nodoc:
       dup.tap do |message|
         message.instance_variable_set(:@thinking, nil)
         message.instance_variable_set(:@raw_reasoning, nil)
+        message.instance_variable_set(:@phase, nil)
+        message.instance_variable_set(:@continuation, nil)
         message.instance_variable_set(:@raw_content, raw_content)
-        message.instance_variable_set(:@tool_calls, tool_calls_without_thought_signatures)
+        message.instance_variable_set(:@tool_calls, coerce_tool_calls(tool_calls_without_thought_signatures))
       end
     end
 
     # Returns +true+ if the assistant requested one or more tool calls,
     # +false+ otherwise.
     def tool_call?
-      !tool_calls.nil? && !tool_calls.empty?
+      !tool_calls.empty?
     end
 
     # Returns +true+ if the message carries the result of a tool call,
@@ -158,6 +300,17 @@ module RubyLLM
       !tool_call_id.nil? && !tool_call_id.empty?
     end
 
+    def provider_call_id # :nodoc:
+      return @provider_call_id if @provider_call_id
+      return tool_call_id unless tool_result? && conversation
+
+      conversation.messages.reverse_each do |message|
+        call = message.tool_calls[tool_call_id]
+        return call.provider_call_id if call
+      end
+      tool_call_id
+    end
+
     # Returns the tool result messages answering this message's tool calls,
     # or an empty array when it made none. Mirrors the +tool_results+
     # association on acts_as_message records.
@@ -176,11 +329,17 @@ module RubyLLM
       finish_reason == :stop && !tool_call?
     end
 
+    # Returns the error text of a failed tool result, or +nil+.
+    def tool_error_message
+      tool_result_payload&.[]('error')
+    end
+
     # Returns +true+ if the response was cut off by a token limit,
     # +false+ otherwise.
     def max_tokens?
       finish_reason == :max_tokens
     end
+    alias truncated? max_tokens?
 
     # Returns +true+ if the model stopped to request tool calls,
     # +false+ otherwise.
@@ -213,18 +372,21 @@ module RubyLLM
       return ruby_llm_usage_cost if model.nil? && ruby_llm_usage_entries.any?
       return @supplied_cost if model.nil? && @supplied_cost
 
-      Cost.new(tokens:, model: model || model_info)
+      Cost.new(tokens: model ? tokens.without_reported_cost : tokens, model: model || model_info)
     end
 
     # Marks this message as an explicit prompt cache boundary. Providers
     # with boundary controls use the conversation up to and including this
-    # message as the cacheable prefix. Pass +ttl:+ to give this boundary its
-    # own cache lifetime instead of the +ttl:+ from Chat#with_caching, on
+    # message as the cacheable prefix. Pass +ttl:+, in any form
+    # Chat#with_caching takes, to give this boundary its own cache lifetime
+    # instead of the +ttl:+ from Chat#with_caching, on
     # providers that set a lifetime per boundary. Returns +self+.
     #
     #   chat.add_message(role: :user, content: long_context).cache_until_here
     #   chat.add_message(role: :user, content: handbook).cache_until_here(ttl: "1h")
     #
+    # To mark system instructions, pass +cache_until_here:+ to
+    # Chat#with_instructions.
     def cache_until_here(ttl: nil)
       @cache_until_here = true
       @cache_ttl = ttl
@@ -261,16 +423,20 @@ module RubyLLM
         content: content,
         attachments: list_to_h(attachments),
         model: model,
+        provider: provider,
         cost: @supplied_cost && cost.to_h,
         reported_cost: tokens.reported_cost,
-        tool_calls: tool_calls&.transform_values(&:to_h),
+        tool_calls: tool_calls.empty? ? nil : tool_calls.transform_values(&:to_h),
         tool_call_id: tool_call_id,
+        provider_call_id: @provider_call_id,
         thinking: thinking&.text,
         thinking_signature: thinking&.signature,
         citations: list_to_h(citations),
         server_tool_calls: list_to_h(server_tool_calls),
         raw_content: raw_content,
         raw_reasoning: raw_reasoning,
+        phase: phase,
+        continuation: continuation,
         mcp_result: mcp_result&.dump,
         finish_reason: finish_reason,
         cache_until_here: cache_until_here? || nil,
@@ -297,25 +463,49 @@ module RubyLLM
 
     private
 
+    def answer_text
+      slug = ruby_llm_usage_entries.reverse.find(&:succeeded?)&.provider if raw_content || phase
+      (Provider.resolve(slug)&.answer_text(self) if slug) || content
+    end
+
+    def tool_result_payload
+      return unless tool_result? && content
+
+      payload = JSON.parse(content)
+      payload if payload.is_a?(Hash)
+    rescue JSON::ParserError
+      nil
+    end
+
     def list_to_h(list)
       list.empty? ? nil : list.map(&:to_h)
     end
 
     def coerce_tool_calls(tool_calls)
-      return tool_calls unless tool_calls.is_a?(Hash)
+      return tool_calls if tool_calls.is_a?(ToolCall::Collection)
+      if tool_calls.is_a?(Array) && tool_calls.all?(ToolCall)
+        return ToolCall::Collection.new(ToolCall.index(tool_calls))
+      end
 
-      tool_calls.to_h do |id, call|
-        next [id, call] unless call.is_a?(Hash)
+      ToolCall::Collection.new(
+        tool_calls.to_h do |id, call|
+          [id, call.is_a?(Hash) ? tool_call_from_h(id, call) : call]
+        end
+      )
+    end
 
-        attributes = call.transform_keys(&:to_sym)
-        [id, ToolCall.new(id: attributes[:id] || id, name: attributes[:name],
-                          arguments: attributes[:arguments] || {},
-                          thought_signature: attributes[:thought_signature], remote: attributes.fetch(:remote, false))]
-      end
+    def tool_call_from_h(id, call)
+      attributes = call.transform_keys(&:to_sym)
+      ToolCall.new(id: attributes[:id] || id, name: attributes[:name],
+                   provider_call_id: attributes[:provider_call_id] || attributes[:provider_id],
+                   arguments: attributes[:arguments] || {},
+                   malformed_arguments: attributes[:malformed_arguments],
+                   namespace: attributes[:namespace],
+                   thought_signature: attributes[:thought_signature], remote: attributes.fetch(:remote, false))
     end
 
     def tool_calls_without_thought_signatures
-      tool_calls&.transform_values do |call|
+      tool_calls.transform_values do |call|
         call.dup.tap { |copy| copy.thought_signature = nil }
       end
     end
@@ -333,12 +523,16 @@ module RubyLLM
       value.is_a?(Hash) ? klass.from_h(value) : value
     end
 
+    def coerce_list(values, klass)
+      Array(values).map { |value| coerce_value(value, klass) }
+    end
+
     def coerce_mcp_result(value)
       value.is_a?(Hash) ? MCP::Result.load(value.transform_keys(&:to_s)) : value
     end
 
     def normalize_content(content)
-      return '' if role == :assistant && content.nil? && tool_calls && !tool_calls.empty?
+      return '' if role == :assistant && content.nil? && tool_call?
       return content if content.nil? || content.is_a?(String)
 
       raise ArgumentError,
@@ -354,7 +548,7 @@ module RubyLLM
       {
         role: role,
         content: content,
-        tool_calls: tool_calls&.values&.map(&:name),
+        tool_calls: tool_calls.map(&:name),
         tool_call_id: tool_call_id,
         model: model,
         finish_reason: finish_reason

data/lib/ruby_llm/model/modalities.rb

@@ -5,7 +5,7 @@ module RubyLLM
     # A Model::Modalities lists the kinds of content a model accepts and
     # produces, as arrays of Strings. Instances come from Model#modalities.
     #
-    #   model = RubyLLM.models.find('gpt-5.6')
+    #   model = RubyLLM.models.find('gpt-6.1-sol')
     #   model.modalities.input   # => ["text", "image", "pdf"]
     #   model.modalities.output  # => ["text"]
     #

data/lib/ruby_llm/model/pricing.rb

@@ -3,9 +3,10 @@
 module RubyLLM
   class Model
     # A Pricing groups a model's prices by usage category: text tokens,
-    # images, audio tokens, and embeddings. Each category is a
-    # PricingCategory. Prices are in USD per million tokens. Instances come
-    # from Model#pricing.
+    # images, audio tokens, and embeddings, plus the characters, audio
+    # seconds, and search units some operations bill instead of tokens.
+    # Each category is a PricingCategory. Prices are in USD per million
+    # units of the category. Instances come from Model#pricing.
     #
     #   model = RubyLLM.models.find "claude-sonnet-5"
     #   model.pricing.text_tokens.input   # => 3
@@ -13,7 +14,7 @@ module RubyLLM
     #
     class Pricing
       # The pricing categories a model may define.
-      CATEGORIES = %i[text_tokens images audio_tokens embeddings].freeze
+      CATEGORIES = %i[text_tokens images audio_tokens embeddings characters audio_seconds search_units].freeze
 
       def initialize(data) # :nodoc:
         @data = {}
@@ -47,6 +48,27 @@ module RubyLLM
         category(:embeddings)
       end
 
+      # Returns the PricingCategory for prices per million input
+      # characters, such as a text-to-speech model's, or an empty category
+      # if the model has none.
+      def characters
+        category(:characters)
+      end
+
+      # Returns the PricingCategory for prices per million seconds of
+      # input audio, such as a transcription model's, or an empty category
+      # if the model has none.
+      def audio_seconds
+        category(:audio_seconds)
+      end
+
+      # Returns the PricingCategory for prices per million search units,
+      # such as a rerank model's, or an empty category if the model has
+      # none.
+      def search_units
+        category(:search_units)
+      end
+
       # Returns the pricing data as a nested Hash keyed by category.
       # Categories without prices are omitted.
       def to_h

data/lib/ruby_llm/model/pricing_category.rb

@@ -5,7 +5,9 @@ module RubyLLM
     # A PricingCategory holds the standard, batch, and long-context pricing
     # tiers for one kind of model usage, such as text tokens or images.
     # Model#pricing returns a Pricing collection whose categories are
-    # PricingCategory instances. Prices are in USD per million tokens.
+    # PricingCategory instances. Prices are in USD per million tokens, or
+    # per million units for categories billed in characters, audio seconds,
+    # or search units.
     #
     #   model    = RubyLLM.models.find "claude-sonnet-5"
     #   category = model.pricing.text_tokens
@@ -39,6 +41,7 @@ module RubyLLM
         @batch = tier_from(data[:batch])
         @long_context = tier_from(data[:long_context])
         @long_context_threshold = Integer(data[:long_context_threshold], exception: false)
+        @context_tiers = context_tiers_from(data[:context_tiers])
       end
 
       # Returns the standard-tier input price in USD per million tokens,
@@ -72,15 +75,11 @@ module RubyLLM
       end
 
       # Returns the PricingTier that applies for a prompt of the given size.
-      # Uses #long_context when that tier exists and +prompt_tokens+ is
-      # greater than #long_context_threshold; otherwise returns #standard.
+      # Uses the largest context tier +prompt_tokens+ is greater than, such
+      # as #long_context above #long_context_threshold; otherwise returns
+      # #standard.
       def tier_for(prompt_tokens)
-        if long_context && long_context_threshold &&
-           prompt_tokens.to_i > long_context_threshold
-          long_context
-        else
-          standard
-        end
+        @context_tiers.find { |size, _| prompt_tokens.to_i > size }&.last || standard
       end
 
       # Returns a Hash with present tier hashes and optional
@@ -105,20 +104,38 @@ module RubyLLM
         entry, threshold = context_cost(cost)
         return [nil, nil] unless entry
 
-        rates = {
+        rates = context_rates(entry)
+        rates.empty? ? [nil, nil] : [rates, threshold]
+      end
+
+      # Returns every context tier in +cost+ as <tt>[size, rates]</tt>. The
+      # registry schema holds one long-context tier, so a model with more
+      # keeps the rest in its metadata.
+      def self.context_tiers_from_cost(cost) # :nodoc:
+        cost = RubyLLM::Support::Utils.deep_symbolize_keys(cost || {})
+        context_entries(cost).filter_map do |entry|
+          size = Integer(entry.dig(:tier, :size), exception: false)
+          rates = context_rates(entry)
+          [size, rates] if size && rates.any?
+        end
+      end
+
+      def self.context_rates(entry) # :nodoc:
+        {
           input_per_million: entry[:input],
           output_per_million: entry[:output],
           cache_read_input_per_million: entry[:cache_read],
           cache_write_input_per_million: entry[:cache_write],
           reasoning_output_per_million: entry[:reasoning]
         }.compact
-        rates.empty? ? [nil, nil] : [rates, threshold]
+      end
+
+      def self.context_entries(cost) # :nodoc:
+        Array(cost[:tiers]).select { |entry| entry.is_a?(Hash) && entry.dig(:tier, :type).to_s == 'context' }
       end
 
       def self.context_cost(cost) # :nodoc:
-        context_tier = Array(cost[:tiers]).find do |entry|
-          entry.is_a?(Hash) && entry.dig(:tier, :type).to_s == 'context'
-        end
+        context_tier = context_entries(cost).first
         if context_tier
           [context_tier, Integer(context_tier.dig(:tier, :size), exception: false)]
         elsif cost[:context_over_200k].is_a?(Hash)
@@ -126,10 +143,16 @@ module RubyLLM
         end
       end
 
-      private_class_method :context_cost
+      private_class_method :context_cost, :context_rates, :context_entries
 
       private
 
+      def context_tiers_from(tiers)
+        tiers = Array(tiers).to_h { |size, rates| [size, tier_from(rates)] }
+        tiers[long_context_threshold] = long_context if long_context && long_context_threshold
+        tiers.sort_by { |size, _| -size }
+      end
+
       def tier_from(tier_data)
         return nil if empty_tier?(tier_data)
 

data/lib/ruby_llm/model/pricing_tier.rb

@@ -4,8 +4,9 @@ module RubyLLM
   class Model
     # A PricingTier holds the prices a model charges within one billing tier,
     # such as standard, batch, or long_context. Each price is in USD per one
-    # million tokens. Prices missing from the registry read as +nil+. A zero
-    # price means the usage is free.
+    # million tokens, or per million units of a category billed in other
+    # units, such as Pricing#characters. Prices missing from the registry
+    # read as +nil+. A zero price means the usage is free.
     #
     # Instances come from PricingCategory#standard, #batch, and #long_context:
     #

data/lib/ruby_llm/model.rb

@@ -5,7 +5,7 @@ module RubyLLM
   # capabilities, modalities, pricing, and provider metadata. Instances come
   # from the registry through RubyLLM.models and from Chat#model.
   #
-  #   model = RubyLLM.models.find('gpt-5.6')
+  #   model = RubyLLM.models.find('gpt-6.1-sol')
   #   model.name              # => "GPT-5.6"
   #   model.provider          # => "openai"
   #   model.context_window    # => 1050000
@@ -14,7 +14,7 @@ module RubyLLM
   class Model
     include Support::Inspectable
 
-    # The provider's identifier for the model, e.g. <tt>"gpt-5.6"</tt>.
+    # The provider's identifier for the model, e.g. <tt>"gpt-6.1-sol"</tt>.
     attr_reader :id
 
     # The human-readable model name, e.g. <tt>"GPT-5.6"</tt>.
@@ -150,16 +150,24 @@ module RubyLLM
     end
 
     # Builds a Cost for +tokens+ (a Tokens object, or anything responding
-    # to +tokens+) using this model's pricing.
+    # to +tokens+) using this model's pricing. A Chat or Message is priced
+    # one provider request at a time, so each request pays the context
+    # tier its own prompt reached.
     #
     #   cost = model.cost_for(response.tokens)
     #   puts cost.total
     #
     # Pass <tt>tier: :batch</tt> to use the model's explicit batch rates.
+    # Amounts the model has no batch price for are +nil+.
     #
     def cost_for(tokens, tier: :standard)
-      tokens = tokens.tokens if tokens.respond_to?(:tokens)
+      requests = usage_entries_of(tokens)
+      if requests.any?
+        complete = requests.all? { |entry| entry.usage_available? || entry.tokens.reported_cost }
+        return Cost.aggregate(requests.map { |entry| cost_for(entry.tokens, tier:) }, complete:)
+      end
 
+      tokens = tokens.tokens if tokens.respond_to?(:tokens)
       Cost.new(tokens:, model: self, tier:)
     end
 
@@ -207,15 +215,25 @@ module RubyLLM
 
     private
 
+    def usage_entries_of(source)
+      return source.usage_entries if source.respond_to?(:usage_entries)
+      return source.ruby_llm_usage_entries if source.respond_to?(:ruby_llm_usage_entries)
+
+      []
+    end
+
     def pricing_data_with_long_context(pricing)
       pricing = RubyLLM::Support::Utils.deep_symbolize_keys(RubyLLM::Support::Utils.deep_dup(pricing))
+      cost = metadata[:cost] || metadata['cost']
       text = pricing[:text_tokens] || {}
+      tiers = PricingCategory.context_tiers_from_cost(cost)
+      pricing[:text_tokens] = text.merge(context_tiers: tiers) if tiers.size > 1
       return pricing if text[:long_context]
 
-      rates, threshold = PricingCategory.long_context_from_cost(metadata[:cost] || metadata['cost'])
+      rates, threshold = PricingCategory.long_context_from_cost(cost)
       return pricing unless rates
 
-      pricing[:text_tokens] = text.merge(long_context: rates)
+      pricing[:text_tokens] = (pricing[:text_tokens] || {}).merge(long_context: rates)
       pricing[:text_tokens][:long_context_threshold] = threshold if threshold
       pricing
     end