Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 34 additions & 11 deletions lib/ruby_llm/attachment.rb
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,8 @@ class Attachment

def initialize(source, filename: nil)
@source = source
if url?
@source = URI source
@filename = filename || File.basename(@source.path).to_s
elsif path?
@source = Pathname.new source
@filename = filename || @source.basename.to_s
elsif active_storage?
@filename = filename || extract_filename_from_active_storage
else
@filename = filename
end
@source = source_type_cast
@filename = filename || source_filename

determine_mime_type
end
Expand Down Expand Up @@ -166,6 +157,38 @@ def load_content_from_active_storage
end
end

def source_type_cast
if url?
URI(@source)
elsif path?
Pathname.new(@source)
else
@source
end
end

def source_filename
if url?
File.basename(@source.path).to_s
elsif path?
@source.basename.to_s
elsif io_like?
extract_filename_from_io
elsif active_storage?
extract_filename_from_active_storage
end
end

def extract_filename_from_io
if defined?(ActionDispatch::Http::UploadedFile) && @source.is_a?(ActionDispatch::Http::UploadedFile)
@source.original_filename.to_s
elsif @source.respond_to?(:path)
File.basename(@source.path).to_s
else
'attachment'
end
end

def extract_filename_from_active_storage # rubocop:disable Metrics/PerceivedComplexity
return 'attachment' unless defined?(ActiveStorage)

Expand Down
100 changes: 99 additions & 1 deletion spec/ruby_llm/chat_content_spec.rb
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# frozen_string_literal: true

require 'spec_helper'

require 'action_dispatch/http/upload'
RSpec.describe RubyLLM::Chat do # rubocop:disable RSpec/MultipleMemoizedHelpers
include_context 'with configured RubyLLM'

Expand Down Expand Up @@ -226,4 +226,102 @@
expect(attachment.send(:url?)).to be true
end
end

describe 'IO attachment handling' do # rubocop:disable RSpec/MultipleMemoizedHelpers
it 'handles StringIO objects' do
require 'stringio'
text_content = 'Hello, this is a test file'
string_io = StringIO.new(text_content)

attachment = RubyLLM::Attachment.new(string_io)

expect(attachment.io_like?).to be true
expect(attachment.content).to eq(text_content)
expect(attachment.filename).to eq('attachment')
expect(attachment.mime_type).to eq('application/octet-stream')
end

it 'handles StringIO objects with filename' do
require 'stringio'
text_content = 'Hello, this is a test file'
string_io = StringIO.new(text_content)

attachment = RubyLLM::Attachment.new(string_io, filename: 'test.txt')

expect(attachment.io_like?).to be true
expect(attachment.content).to eq(text_content)
expect(attachment.filename).to eq('test.txt')
expect(attachment.mime_type).to eq('text/plain')
end

it 'handles Tempfile objects' do
tempfile = Tempfile.new(['test', '.txt'])
tempfile.write('Tempfile content')
tempfile.rewind

attachment = RubyLLM::Attachment.new(tempfile)

expect(attachment.io_like?).to be true
expect(attachment.content).to eq('Tempfile content')
expect(attachment.filename).to be_present
expect(attachment.mime_type).to eq('text/plain')
end

it 'handles File objects' do
file = File.open(text_path, 'r')

attachment = RubyLLM::Attachment.new(file)

expect(attachment.io_like?).to be true
expect(attachment.content).to be_present
expect(attachment.filename).to eq('ruby.txt')
expect(attachment.mime_type).to eq('text/plain')

file.close
end

it 'handles ActionDispatch::Http::UploadedFile' do
tempfile = Tempfile.new(['ruby', '.png'])
tempfile.binmode
File.open(image_path, 'rb') { |f| tempfile.write(f.read) }
tempfile.rewind

uploaded_file = ActionDispatch::Http::UploadedFile.new(
tempfile: tempfile,
filename: 'ruby.png',
type: 'image/png'
)

attachment = RubyLLM::Attachment.new(uploaded_file)

expect(attachment.io_like?).to be true
expect(attachment.content).to be_present
expect(attachment.filename).to eq('ruby.png')
expect(attachment.mime_type).to eq('image/png')
expect(attachment.type).to eq(:image)
end

it 'rewinds IO objects before reading' do
require 'stringio'
string_io = StringIO.new('Initial content')
string_io.read # Move position to end

attachment = RubyLLM::Attachment.new(string_io, filename: 'test.txt')

expect(attachment.content).to eq('Initial content')
end

it 'creates content with IO attachments' do
require 'stringio'
string_io = StringIO.new('Test content')
content = RubyLLM::Content.new('Check this')
content.add_attachment(string_io, filename: 'test.txt')

expect(content.attachments).not_to be_empty
expect(content.attachments.first).to be_a(RubyLLM::Attachment)
expect(content.attachments.first.io_like?).to be true
expect(content.attachments.first.filename).to eq('test.txt')
expect(content.attachments.first.mime_type).to eq('text/plain')
end
end
end