From 6f76aeccf9f457d71547d1ffde44c26cdd7f0c58 Mon Sep 17 00:00:00 2001
From: pekopekopekopayo <dlehddus1285@gmail.com>
Date: Mon, 3 Nov 2025 18:45:42 +0900
Subject: [PATCH] Refactor Attachment#initialize to extract source handling
 methods

---
 lib/ruby_llm/attachment.rb         |  45 +++++++++----
 spec/ruby_llm/chat_content_spec.rb | 100 ++++++++++++++++++++++++++++-
 2 files changed, 133 insertions(+), 12 deletions(-)

diff --git a/lib/ruby_llm/attachment.rb b/lib/ruby_llm/attachment.rb
index 1e4a8e931..af983d065 100644
--- a/lib/ruby_llm/attachment.rb
+++ b/lib/ruby_llm/attachment.rb
@@ -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
@@ -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)
 
diff --git a/spec/ruby_llm/chat_content_spec.rb b/spec/ruby_llm/chat_content_spec.rb
index 21b8a8112..b43fa3486 100644
--- a/spec/ruby_llm/chat_content_spec.rb
+++ b/spec/ruby_llm/chat_content_spec.rb
@@ -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'
 
@@ -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