Monkeypatch RubyLLM to work with Shrine files.

ruby_llm_shrine.rb

# frozen_string_literal: true

# Monkey patch to add Shrine support to RubyLLM::Attachment
# - Works with Shrine::UploadedFile, Shrine::Attacher (via .file), or Array of UploadedFile (takes first)
# - Reads bytes from storage (open/download), avoiding expiring URLs
# - Pulls filename and mime_type from Shrine metadata when present

module RubyLLM
  class Attachment
    # ---- Shrine detection helpers ------------------------------------------

    def shrine?
      return false unless defined?(Shrine)

      uf = shrine_uploaded_file
      !!uf
    end

    def shrine_uploaded_file
      return nil unless defined?(Shrine)

      src = @source

      # If it's an array of files, take the first to maintain single-attachment semantics
      src = src.first if src.is_a?(Array)

      # If it's an attacher, grab the underlying file
      if src.respond_to?(:file) && src.file.is_a?(Shrine::UploadedFile)
        src = src.file
      end

      return src if src.is_a?(Shrine::UploadedFile)

      nil
    end
    private :shrine_uploaded_file

    # ---- Override a few accessors/behaviors --------------------------------

    # Provide a filename lazily if we weren't able to set it in initialize
    def filename
      return @filename if defined?(@filename) && @filename
      return extract_filename_from_shrine if shrine?

      nil
    end

    # Compute mime_type lazily with Shrine hints (and keep original fallbacks)
    def mime_type
      # If we already have a specific type, use it
      if defined?(@mime_type) && @mime_type && @mime_type != "application/octet-stream"
        return @mime_type
      end

      # Prefer Shrine metadata when available
      if shrine?
        @mime_type = shrine_content_type
      end

      # Fall back to RubyLLM's detection on path/source and then content sniffing
      @mime_type ||= RubyLLM::MimeType.for(url? ? nil : @source, name: filename)
      @mime_type = RubyLLM::MimeType.for(content) if @mime_type == "application/octet-stream"
      @mime_type = "audio/wav" if @mime_type == "audio/x-wav" # normalize

      @mime_type
    end

    # Add a Shrine branch to the content loader while preserving original logic
    def content
      return @content if defined?(@content) && !@content.nil?

      if url?
        fetch_content
      elsif path?
        load_content_from_path
      elsif active_storage?
        load_content_from_active_storage
      elsif shrine?
        load_content_from_shrine
      elsif io_like?
        load_content_from_io
      else
        RubyLLM.logger.warn "Source is neither a URL, path, ActiveStorage, Shrine, nor IO-like: #{@source.class}"
        nil
      end

      @content
    end

    private

    # ---- Shrine-specific loaders/derivers -----------------------------------

    def load_content_from_shrine
      uploaded = shrine_uploaded_file
      return unless uploaded

      # Prefer a streaming read when available
      io = uploaded.open
      begin
        @content = io.read
      ensure
        io.close if io.respond_to?(:close)
      end

      # Fallback to download (Tempfile) if needed
      if @content.nil? || @content.empty?
        tmp = uploaded.download
        begin
          @content = tmp.read
        ensure
          tmp.close! if tmp.respond_to?(:close!)
        end
      end
    rescue => e
      RubyLLM.logger.error "Shrine content load failed: #{e.class}: #{e.message}"
      @content = nil
    end

    def extract_filename_from_shrine
      uploaded = shrine_uploaded_file
      return "attachment" unless uploaded

      (uploaded.respond_to?(:original_filename) && uploaded.original_filename.presence) ||
        (uploaded.respond_to?(:metadata) && uploaded.metadata&.dig("filename").presence) ||
        begin
          # last resort: try derive from URL path if present
          u = begin
            uploaded.url
          rescue
            nil
          end
          u ? File.basename(URI.parse(u).path.to_s) : nil
        rescue
          nil
        end ||
        "attachment"
    end

    def shrine_content_type
      uploaded = shrine_uploaded_file
      return unless uploaded

      # Many apps use the :determine_mime_type plugin, which stores metadata["mime_type"]
      if uploaded.respond_to?(:mime_type) && uploaded.mime_type
        uploaded.mime_type
      else
        uploaded.respond_to?(:metadata) ? uploaded.metadata&.dig("mime_type") : nil
      end
    end
  end
end

# Example usage in your Rails app:
#
# 1. Place this file in config/initializers/ruby_llm_shrine.rb
#
# 2. Set up your model with Shrine:
#
# class Message < ApplicationRecord
#   include RubyLLM::ActiveRecord::ActsAs
#   include ImageUploader::Attachment(:image)
#   include DocumentUploader::Attachment(:document)
#
#   acts_as_llm_message
# end
#
# 3. Use Shrine attachments in your chat:
#
# # With uploaded file
# user_message = Message.create!(
#   role: "user",
#   content: "What's in this image?"
# )
# user_message.image = params[:file]  # Shrine will handle this
# user_message.save!
#
# # RubyLLM will automatically detect and use the Shrine attachment
# chat = RubyLLM::Chat.new(provider: :openai, model: "gpt-4o")
# chat.messages << user_message.to_llm_message
# response = chat.completion
#
# # With existing Shrine attachment
# existing_message = Message.find(1)
# if existing_message.image.present?
#   content = RubyLLM::Content.new("Analyze this image", [existing_message.image])
#   chat.ask(content)
# end