llm_cost_tracker 0.14.2 → 0.15.0

data/lib/llm_cost_tracker/integrations/ruby_llm.rb

@@ -1,420 +1,19 @@
 # frozen_string_literal: true
 
-require_relative "base"
+require_relative "ruby_llm/v1"
+require_relative "ruby_llm/v2"
 
 module LlmCostTracker
   module Integrations
     module RubyLlm
-      extend Base
-
-      minimum_version "1.15.0"
-      maximum_version "3.0.0"
-
       class << self
-        def patch_targets
-          [
-            patch_target("RubyLLM::Provider", with: ProviderPatch),
-            patch_target("RubyLLM::Providers::Gemini::Transcription", with: GeminiTranscriptionPatch, optional: true),
-            patch_target("RubyLLM::Protocols::Gemini",
-                         with: GeminiImagesPatch,
-                         optional: true,
-                         skip_when_methods_missing: true),
-            patch_target("RubyLLM::Protocols::Gemini",
-                         with: GeminiCachePatch,
-                         optional: true,
-                         skip_when_methods_missing: true),
-            patch_target("RubyLLM::Streaming", with: StreamPatch, optional: true),
-            patch_target("RubyLLM::Protocol::Streaming", with: StreamPatch, optional: true),
-            *%w[RubyLLM::Providers::OpenAI RubyLLM::Providers::Gemini
-                RubyLLM::Protocols::ChatCompletions RubyLLM::Protocols::Gemini].map do |name|
-              patch_target(name, with: ResponseBodyPatch, optional: true, skip_when_methods_missing: true)
-            end
-          ]
-        end
-
-        def record_completion(provider, response, request:, latency_ms:, has_block:)
-          model = response_model_id(response) || model_id_from_request(request[:model])
-          record_usage(
-            provider: provider,
-            model: model,
-            response: response,
-            latency_ms: latency_ms,
-            stream: has_block || request[:stream] == true,
-            cache_ttl: request[:caching].try(:[], :ttl),
-            service_line_items: service_line_items(provider.slug.to_s, response, model)
-          )
-        end
-
-        def service_line_items(provider, response, model)
-          body = raw_body(response)
-          case provider
-          when "anthropic"
-            counts = response.try(:tokens).try(:server_tool_use) || body.dig("usage", "server_tool_use")
-            Providers::Anthropic::UsageExtractor.service_line_items(server_tool_use: counts&.symbolize_keys)
-          when "openai" then Providers::Openai::ServiceCharges.service_line_items_for(body, model: model)
-          when "gemini" then Providers::Gemini::Parser.new.service_line_items_for(gemini_body(response), model: model)
-          when "openrouter" then Providers::Openai::ServiceCharges.billed_line_items(usage_hash(body))
-          else []
-          end
-        end
-
-        def record_embedding(provider, response, request:, latency_ms:)
-          record_usage(
-            provider: provider,
-            model: response_model_id(response) || model_id_from_request(request[:model]),
-            response: response,
-            latency_ms: latency_ms,
-            stream: false,
-            output_tokens: 0,
-            service_line_items: gemini_video_line_items(provider, response)
-          )
-        end
-
-        def gemini_video_line_items(provider, response)
-          details = raw_body(response).dig("usageMetadata", "promptTokenDetails") if provider.slug.to_s == "gemini"
-          video = Providers::Gemini::UsageExtractor.modality_tokens(details, "VIDEO")
-          video.positive? ? [Charges::LineItem.build(dimension_key: "video_input", quantity: video)] : []
-        end
-
-        def record_transcription(provider, response, request:, latency_ms:, stream: false)
-          model = response_model_id(response) || model_id_from_request(request[:model])
-          match = LlmCostTracker::Pricing::Matcher.lookup(provider: provider.slug.to_s, model: model)
-          counts = token_counts(response)
-          usage = usage_hash(raw_body(response))
-          no_tokens = counts[:input].to_i.zero? && counts[:output].to_i.zero?
-          duration = billed_duration(usage, response, no_tokens)
-          line_items = Providers::Openai::ServiceCharges.transcription_line_items(duration)
-          audio_input = Providers::Openai::UsageExtractor.audio_input_tokens(usage)
-          audio_input = counts[:input].to_i if audio_input.zero? && match&.prices&.key?("audio_input")
-          record_usage(
-            provider: provider,
-            model: model,
-            response: response,
-            latency_ms: latency_ms,
-            stream: stream,
-            audio_input_tokens: audio_input,
-            service_line_items: line_items,
-            usage_source: (Usage::Source::UNKNOWN if no_tokens && line_items.empty?)
-          )
-        end
-
-        def billed_duration(usage, response, no_tokens)
-          return usage if usage[:type].to_s == "duration"
-
-          { type: "duration", seconds: response.duration&.ceil } if no_tokens
-        end
-
-        def record_image(provider, response, request:, latency_ms:)
-          image = response.is_a?(Array) ? response.first : response
-          model = response_model_id(image) || model_id_from_request(request[:model])
-          usage = image_usage(image)
-          extractor = Providers::Openai::UsageExtractor
-          image_input = extractor.image_input_tokens(usage)
-          image_output, text_output = extractor.split_output(
-            output_tokens: usage[:output_tokens].to_i,
-            image_output_details: gemini_image_output_tokens(image) || extractor.image_output_tokens(usage),
-            text_output_details: extractor.text_output_tokens(usage),
-            audio_output: 0,
-            default_to_image: model.to_s.match?(/\A(gpt-image-|gemini-.*-image)/)
-          )
-          record_passthrough(
-            provider: provider.slug.to_s,
-            model: model,
-            response: image,
-            latency_ms: latency_ms,
-            input_tokens: [usage[:input_tokens].to_i - image_input, 0].max,
-            image_input_tokens: image_input,
-            output_tokens: text_output,
-            image_output_tokens: image_output
-          )
-        end
-
-        def record_moderation(provider, response, request:, latency_ms:)
-          record_passthrough(
-            provider: provider.slug.to_s,
-            model: response_model_id(response) || model_id_from_request(request[:model]),
-            response: response,
-            latency_ms: latency_ms,
-            input_tokens: 0,
-            output_tokens: 0
-          )
-        end
-
-        def record_cache_storage(cache)
-          return unless active? && cache.provider.to_s == "gemini"
-
-          record_safely do
-            event = Providers::Gemini::Parser.new.cache_storage_event(cache.metadata)
-            LlmCostTracker::Tracker.record(event: event) if event
-          end
-        end
-
-        def image_usage(image)
-          usage = image.try(:usage)
-          usage = image.send(:raw_usage) if !usage.is_a?(Hash) && image.respond_to?(:raw_usage, true)
-          (usage.is_a?(Hash) ? usage : {}).with_indifferent_access
-        end
-
-        def gemini_image_output_tokens(image)
-          metadata = image.instance_variable_get(:@llm_cost_tracker_usage_metadata)
-          Providers::Gemini::UsageExtractor.modality_tokens(metadata["candidatesTokensDetails"], "IMAGE") if metadata
-        end
-
-        def record_usage(provider:,
-                         model:,
-                         response:,
-                         latency_ms:,
-                         stream:,
-                         output_tokens: nil,
-                         audio_input_tokens: 0,
-                         cache_ttl: nil,
-                         service_line_items: [],
-                         usage_source: nil)
-          return unless active?
-
-          record_safely do
-            counts = token_counts(response, provider.slug.to_s)
-            output_tokens = counts[:output] if output_tokens.nil?
-            next if counts[:input].nil? && output_tokens.nil? && service_line_items.empty? && usage_source.nil?
-
-            cache_write_5m, cache_write_1h = cache_write_split(provider, response, counts[:cache_write], cache_ttl)
-            LlmCostTracker::Tracker.record(
-              event: Event.build(
-                provider: provider.slug.to_s,
-                model: model,
-                pricing_mode: pricing_mode_for(provider: provider, model: model, response: response),
-                token_usage: gemini_token_usage(provider, response) || Usage::TokenUsage.build(
-                  input_tokens: counts[:input].to_i - audio_input_tokens,
-                  audio_input_tokens: audio_input_tokens,
-                  output_tokens: output_tokens.to_i,
-                  cache_read_input_tokens: counts[:cache_read].to_i,
-                  cache_write_input_tokens: cache_write_5m,
-                  cache_write_extended_input_tokens: cache_write_1h,
-                  hidden_output_tokens: counts[:thinking].to_i
-                ),
-                service_line_items: service_line_items + gemini_cache_read_line_items(provider, response),
-                stream: stream,
-                usage_source: usage_source || LlmCostTracker::Usage::Source::SDK_RESPONSE,
-                provider_response_id: provider_response_id_for(response)
-              ),
-              latency_ms: latency_ms
-            )
-          end
-        end
-
-        def gemini_token_usage(provider, response)
-          usage = gemini_usage_metadata(response) if provider.slug.to_s == "gemini"
-          return unless usage.is_a?(Hash)
-
-          usage = usage.merge("promptTokensDetails" => usage["promptTokenDetails"]) if usage.key?("promptTokenDetails")
-          Providers::Gemini::UsageExtractor.token_usage(usage)
-        end
-
-        def gemini_cache_read_line_items(provider, response)
-          usage = gemini_usage_metadata(response) if provider.slug.to_s == "gemini"
-          usage.is_a?(Hash) ? Providers::Gemini::UsageExtractor.cache_read_line_items(usage) : []
-        end
-
-        def gemini_body(response)
-          content = response.try(:raw_content)
-          interaction = content["response"] if content.is_a?(Hash)
-          interaction.is_a?(Hash) && interaction["usage"].is_a?(Hash) ? interaction : raw_body(response)
-        end
-
-        def gemini_usage_metadata(response)
-          body = gemini_body(response)
-          return body["usageMetadata"] unless body["usage"].is_a?(Hash)
-
-          Providers::Gemini::Parser.new.interaction_usage_metadata(body["usage"], body["service_tier"])
-        end
-
-        def token_counts(response, provider = nil)
-          tokens = response.try(:tokens)
-          return { input: response.try(:input_tokens), output: response.try(:output_tokens) } unless tokens
-
-          usage = raw_body(response)["usage"] || {}
-          input = usage["inputTokens"] || tokens.input
-          input = [input, usage["input_tokens"]].compact.max if provider == "anthropic"
-          output = tokens.output
-          thinking = tokens.thinking.to_i
-          raw_input = (usage["input_tokens"] || usage["prompt_tokens"]).to_i
-          output += thinking if output && usage["total_tokens"] == raw_input + output + thinking
-          {
-            input: input,
-            output: output,
-            cache_read: tokens.cache_read,
-            cache_write: tokens.cache_write,
-            thinking: tokens.thinking
-          }
-        end
-
-        def cache_write_split(provider, response, cache_write, cache_ttl)
-          usage = raw_body(response)["usage"] || {}
-          cache = case provider.slug.to_s
-                  when "anthropic" then usage["cache_creation"]
-                  when "bedrock"
-                    Array(usage["cacheDetails"]).to_h { |d| ["ephemeral_#{d['ttl']}_input_tokens", d["inputTokens"]] }
-                  end
-          return [cache_write.to_i, 0] unless cache.is_a?(Hash)
-
-          five_minute = cache["ephemeral_5m_input_tokens"].to_i
-          one_hour = cache["ephemeral_1h_input_tokens"].to_i
-          # RubyLLM sums the writes of every pause_turn segment, but the raw body is only the last segment's.
-          earlier = [cache_write.to_i - five_minute - one_hour, 0].max
-          cache_ttl.to_s == "1h" ? [five_minute, one_hour + earlier] : [five_minute + earlier, one_hour]
-        end
-
-        def model_id_from_request(value)
-          return nil if value.nil?
-          return value.to_s if value.is_a?(String) || value.is_a?(Symbol)
+        def install = implementation.install
 
-          (value.try(:id) || value.try(:model_id) || value.try(:model))&.to_s
-        end
-
-        def provider_response_id_for(response)
-          body = gemini_body(response)
-          body["id"] || body["responseId"]
-        end
-
-        def raw_body(response)
-          raw = response.try(:raw)
-          body = raw.respond_to?(:body) ? raw.body : raw
-          body = (raw || response).instance_variable_get(:@llm_cost_tracker_body) unless body.is_a?(Hash)
-          body.is_a?(Hash) ? body : {}
-        end
-
-        def usage_hash(body) = (body["usage"] || {}).deep_symbolize_keys
-
-        def keep_usage(result, body)
-          usage = body.slice("usage", "usageMetadata") if body.is_a?(Hash)
-          result.instance_variable_set(:@llm_cost_tracker_body, usage) if usage
-          result
-        end
-
-        def response_model_id(response)
-          (response.try(:model_id) || response.try(:model))&.to_s
-        end
-
-        def pricing_mode_for(provider:, model:, response:)
-          body = raw_body(response)
-          case provider.slug.to_s
-          when "anthropic", "bedrock"
-            Providers::Anthropic::UsageExtractor.pricing_mode(request: { model: model },
-                                                              usage: body["usage"]&.deep_symbolize_keys)
-          when "gemini" then gemini_usage_metadata(response).try(:[], "serviceTier")
-          when "openai", "xai", "mistral"
-            Providers::Openai::ResponseParser.combined_pricing_mode(
-              provider: provider.slug.to_s,
-              host: URI(provider.api_base).host,
-              model: model,
-              service_tier: body["service_tier"] || body.dig("usage", "service_tier")
-            )
-          else body["service_tier"]
-          end
-        end
-
-        def request_params(args, kwargs)
-          input = args.first
-          if input.is_a?(Array)
-            input = input.map { |msg| msg.try(:content).then { |content| content.try(:text) || content } || msg }
-          end
-          kwargs.merge(input: input, model: model_id_from_request(kwargs[:model])).with_indifferent_access
-        end
-
-        def blocking_seam(resource, record_method, **extras)
-          {
-            provider: resource.slug.to_s,
-            record: lambda do |response, request, latency_ms|
-              public_send(record_method, resource, response, request: request, latency_ms: latency_ms, **extras)
-            end
-          }
-        end
-      end
-
-      module ProviderPatch
-        def complete(*args, **kwargs, &)
-          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_completion, has_block: block_given?)
-          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
-        end
-
-        def embed(*args, **kwargs)
-          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_embedding)
-          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
-        end
-
-        def transcribe(*args, **kwargs, &)
-          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_transcription, stream: block_given?)
-          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
-        end
-
-        def paint(*args, **kwargs)
-          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_image)
-          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
-        end
-
-        def moderate(*args, **kwargs)
-          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_moderation)
-          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
-        end
-      end
-
-      module GeminiImagesPatch
-        def parse_image_responses(response, *, **)
-          images = super
-          metadata = response.body["usageMetadata"]
-          Array(images).each { |image| image.instance_variable_set(:@llm_cost_tracker_usage_metadata, metadata) }
-          images
-        end
-      end
-
-      module GeminiCachePatch
-        def cache_content(*, **)
-          super.tap { |cache| LlmCostTracker::Integrations::RubyLlm.record_cache_storage(cache) }
-        end
-      end
-
-      module GeminiTranscriptionPatch
-        def transcribe(*args, **kwargs)
-          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_transcription)
-          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
-        end
-      end
-
-      module ResponseBodyPatch
-        def parse_transcription_response(response, **)
-          LlmCostTracker::Integrations::RubyLlm.keep_usage(super, response.body)
-        end
-
-        def parse_embedding_response(response, **)
-          LlmCostTracker::Integrations::RubyLlm.keep_usage(super, response.body)
-        end
-
-        def build_streamed_transcription(chunks, **)
-          LlmCostTracker::Integrations::RubyLlm.keep_usage(super, chunks.reverse.find(&:done?)&.raw)
-        end
-      end
-
-      module StreamPatch
-        private
-
-        def stream_response(...)
-          body = @llm_cost_tracker_stream_body = {}
-          super.tap do |message|
-            message.raw.instance_variable_set(:@llm_cost_tracker_body, body)
-            input = body.dig("usage", "input_tokens")
-            input = body["usage"]["input_tokens"] = input + @llm_cost_tracker_paused_input.to_i if input
-            @llm_cost_tracker_paused_input = (input if message.try(:finish_reason) == :pause_turn)
-          end
-        end
+        def status = implementation.status
 
-        def build_on_data_handler(*, &handler)
-          body = @llm_cost_tracker_stream_body
-          super do |data|
-            if body && data.is_a?(Hash)
-              body.deep_merge!(data.values_at("message", "response").find { |part| part.is_a?(Hash) } || data)
-            end
-            handler.call(data)
-          end
+        def implementation
+          version = V2.gem_version
+          version && version >= Gem::Version.new("2.0.0") ? V2 : V1
         end
       end
     end

data/lib/llm_cost_tracker/ledger/isolation.rb

@@ -6,8 +6,21 @@ module LlmCostTracker
   module Ledger
     module Isolation
       class << self
-        def guard(model = LlmCostTracker::Call, &)
-          connection = model.connection
+        def guard(model = LlmCostTracker::Call, &block)
+          leased(model) { |connection| isolate(connection, model, &block) }
+        end
+
+        def leased(model = LlmCostTracker::Call, &block)
+          writing { model.with_connection(prevent_permanent_checkout: true, &block) }
+        end
+
+        def writing(&)
+          ActiveRecord::Base.connected_to(role: ActiveRecord.writing_role, prevent_writes: false, &)
+        end
+
+        private
+
+        def isolate(connection, model, &)
           nested = connection.transaction_open?
           return yield unless nested
 

data/lib/llm_cost_tracker/ledger/period/totals.rb

@@ -69,7 +69,7 @@ module LlmCostTracker
         end
 
         def quote(value)
-          LlmCostTracker::Call.connection.quote(value)
+          LlmCostTracker::Call.lease_connection.quote(value)
         end
       end
     end

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -4,11 +4,13 @@ require "bigdecimal/util"
 
 require_relative "isolation"
 require_relative "period"
+require_relative "schema/adapter"
 
 module LlmCostTracker
   module Ledger
     module Rollups
       SOURCE_COLUMNS = %i[tracked_at total_cost pricing_snapshot provider].freeze
+      BUCKET = %i[period period_start currency provider].freeze
 
       class << self
         def cache_active?
@@ -17,7 +19,7 @@ module LlmCostTracker
           rollup = LlmCostTracker::CallRollup
           return true if rollup.table_exists?
 
-          rollup.connection.schema_cache.clear_data_source_cache!(rollup.table_name)
+          rollup.with_connection { |connection| connection.schema_cache.clear_data_source_cache!(rollup.table_name) }
           return true if rollup.table_exists?
 
           warn_missing_table
@@ -25,10 +27,7 @@ module LlmCostTracker
         end
 
         def increment!(events)
-          return unless cache_active?
-
-          rows = rows_from_buckets(bucket_totals(Array(events)))
-          LlmCostTracker::CallRollup.increment_all(rows) if rows.any?
+          write(rollup_rows(events)) if cache_active?
         end
 
         ROLLUP_INCREMENT_ATTEMPTS = 3
@@ -41,15 +40,18 @@ module LlmCostTracker
         def increment_safely!(events)
           return unless LlmCostTracker.configuration.budgets.totals_source == :cache
 
-          LlmCostTracker::Call.current_transaction.after_commit { increment_with_retries(events) }
+          Isolation.leased do
+            LlmCostTracker::Call.current_transaction.after_commit { increment_with_retries(events) }
+          end
         end
 
         def increment_with_retries(events)
-          retryable = !LlmCostTracker::Call.connection.transaction_open?
+          retryable = LlmCostTracker::Call.with_connection { |connection| !connection.transaction_open? }
+          rows = nil
           attempt = 0
           begin
             attempt += 1
-            Isolation.guard { increment!(events) }
+            Isolation.guard { write(rows ||= cache_active? ? rollup_rows(events) : []) }
           rescue LlmCostTracker::Error
             raise
           rescue StandardError => e
@@ -66,6 +68,20 @@ module LlmCostTracker
           end
         end
 
+        def rollup_rows(events)
+          rows_from_buckets(bucket_totals(Array(events))).sort_by { |row| row.values_at(*BUCKET).map(&:to_s) }
+        end
+
+        def write(rows)
+          return if rows.empty?
+
+          size = Schema::Adapter.mysql?(LlmCostTracker::CallRollup.lease_connection) ? 1 : rows.size
+          until rows.empty?
+            LlmCostTracker::CallRollup.increment_all(rows.first(size))
+            rows.shift(size)
+          end
+        end
+
         def rolled_back?(error)
           error.is_a?(ActiveRecord::Deadlocked) || error.is_a?(ActiveRecord::LockWaitTimeout)
         end

data/lib/llm_cost_tracker/ledger/schema/adapter.rb

@@ -30,6 +30,15 @@ module LlmCostTracker
             raise Error, "Unsupported database adapter: #{adapter_name(value)}. Use PostgreSQL or MySQL."
           end
 
+          def drop_invalid_index(connection, name)
+            return unless postgresql?(connection)
+
+            invalid = connection.select_value(
+              "SELECT 1 FROM pg_index WHERE indexrelid = to_regclass(#{connection.quote(name.to_s)}) AND NOT indisvalid"
+            )
+            connection.execute("DROP INDEX CONCURRENTLY IF EXISTS #{connection.quote_table_name(name)}") if invalid
+          end
+
           PG_PERIOD_FORMATS = { day: "YYYY-MM-DD", month: "YYYY-MM" }.freeze
           MYSQL_PERIOD_FORMATS = { day: "%Y-%m-%d", month: "%Y-%m" }.freeze
           private_constant :PG_PERIOD_FORMATS, :MYSQL_PERIOD_FORMATS
@@ -38,13 +47,9 @@ module LlmCostTracker
             period = period.to_sym
             zone = time_zone&.tzinfo&.name
             if postgresql?(connection)
-              if zone && postgresql_zone?(connection, zone)
-                column = "(#{column}::timestamp AT TIME ZONE 'UTC') AT TIME ZONE '#{zone}'"
-              end
-              "TO_CHAR(DATE_TRUNC('#{period}', #{column}), '#{PG_PERIOD_FORMATS.fetch(period)}')"
+              postgresql_period_sql(connection, period, column, zone)
             elsif mysql?(connection)
-              column = "COALESCE(CONVERT_TZ(#{column}, '+00:00', '#{zone}'), #{column})" if zone
-              "DATE_FORMAT(#{column}, '#{MYSQL_PERIOD_FORMATS.fetch(period)}')"
+              mysql_period_sql(period, column, zone)
             else
               ensure_supported!(connection)
             end
@@ -54,6 +59,18 @@ module LlmCostTracker
 
           private
 
+          def postgresql_period_sql(connection, period, column, zone)
+            if zone && postgresql_zone?(connection, zone)
+              column = "(#{column}::timestamp AT TIME ZONE 'UTC') AT TIME ZONE '#{zone}'"
+            end
+            "TO_CHAR(DATE_TRUNC('#{period}', #{column}), '#{PG_PERIOD_FORMATS.fetch(period)}')"
+          end
+
+          def mysql_period_sql(period, column, zone)
+            column = "COALESCE(CONVERT_TZ(#{column}, '+00:00', '#{zone}'), #{column})" if zone
+            "DATE_FORMAT(#{column}, '#{MYSQL_PERIOD_FORMATS.fetch(period)}')"
+          end
+
           def postgresql_zone?(connection, zone)
             @postgresql_zones ||= {}
             @postgresql_zones.fetch(zone) do

data/lib/llm_cost_tracker/ledger/schema/base.rb

@@ -21,11 +21,8 @@ module LlmCostTracker
         end
 
         def current_schema_errors
-          Adapter.ensure_supported!(model.connection)
-          unless model.table_exists?
-            model.connection.schema_cache.clear_data_source_cache!(model.table_name)
-            return ["#{model.table_name} table is missing"] unless model.table_exists?
-          end
+          model.with_connection { |connection| Adapter.ensure_supported!(connection) }
+          return ["#{model.table_name} table is missing"] unless table_exists?
 
           columns_hash = model.columns_hash
           cache = @schema_capabilities
@@ -38,6 +35,13 @@ module LlmCostTracker
 
         private
 
+        def table_exists?
+          return true if model.table_exists?
+
+          model.with_connection { |connection| connection.schema_cache.clear_data_source_cache!(model.table_name) }
+          model.table_exists?
+        end
+
         def column_errors(columns_hash)
           missing = required_columns - columns_hash.keys
           missing.any? ? ["missing columns: #{missing.join(', ')}"] : []