Better support for media (File / URL) content

This is the foundation for supporting Vision / etc via multi part content.

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    omniai (1.2.0)
+    omniai (1.2.1)
       event_stream_parser
       http
       zeitwerk

lib/omniai.rb

@@ -2,10 +2,12 @@
 
 require 'event_stream_parser'
 require 'http'
+require 'uri'
 require 'zeitwerk'
 
 loader = Zeitwerk::Loader.for_gem
 loader.inflector.inflect 'omniai' => 'OmniAI'
+loader.inflector.inflect 'url' => 'URL'
 loader.setup
 
 module OmniAI

lib/omniai/chat/content/file.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    module Content
+      # A file that is either audio / image / video.
+      class File < Media
+        attr_accessor :io
+
+        # @param io [IO, Pathname, String]
+        # @param type [Symbol, String] :image, :video, :audio, "audio/flac", "image/jpeg", "video/mpeg", etc.
+        def initialize(io, type)
+          super(type)
+          @io = io
+        end
+
+        # @return [String]
+        def fetch!
+          case @io
+          when IO then @io.read
+          else ::File.binread(@io)
+          end
+        end
+      end
+    end
+  end
+end

lib/omniai/chat/content/media.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    module Content
+      # An abstract class that represents audio / image / video and is used for both files and urls.
+      class Media
+        attr_accessor :type
+
+        # @param type [String] "audio/flac", "image/jpeg", "video/mpeg", etc.
+        def initialize(type)
+          @type = type
+        end
+
+        # @return [Boolean]
+        def text?
+          @type.match?(%r{^text/})
+        end
+
+        # @return [Boolean]
+        def audio?
+          @type.match?(%r{^audio/})
+        end
+
+        # @return [Boolean]
+        def image?
+          @type.match?(%r{^image/})
+        end
+
+        # @return [Boolean]
+        def video?
+          @type.match?(%r{^video/})
+        end
+
+        # @yield [io]
+        def fetch!(&)
+          raise NotImplementedError, "#{self.class}#fetch! undefined"
+        end
+
+        # e.g. "Hello" -> "SGVsbG8h"
+        #
+        # @return [String]
+        def data
+          Base64.strict_encode64(fetch!)
+        end
+
+        # e.g. "data:text/html;base64,..."
+        #
+        # @return [String]
+        def data_uri
+          "data:#{@type};base64,#{data}"
+        end
+      end
+    end
+  end
+end

lib/omniai/chat/content/text.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    module Content
+      # Just some text.
+      class Text
+        attr_accessor :text
+
+        # @param text [text]
+        def initialize(text)
+          @text = text
+        end
+      end
+    end
+  end
+end

lib/omniai/chat/content/url.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    module Content
+      # A url that is either audio / image / video.
+      class URL < Media
+        attr_accessor :url, :type
+
+        class HTTPError < OmniAI::HTTPError; end
+
+        # @param url [URI, String]
+        # @param type [Symbol, String] "audio/flac", "image/jpeg", "video/mpeg", etc.
+        def initialize(url, type)
+          super(type)
+          @url = url
+        end
+
+        # @raise [HTTPError]
+        #
+        # @return [String]
+        def fetch!
+          response = request!
+          String(response.body)
+        end
+
+        private
+
+        # @raise [HTTPError]
+        #
+        # @return [HTTP::Response]
+        def request!
+          response = HTTP.get(@url)
+          raise HTTPError, response.flush unless response.status.success?
+
+          response
+        end
+      end
+    end
+  end
+end

lib/omniai/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module OmniAI
-  VERSION = '1.2.0'
+  VERSION = '1.2.1'
 end

spec/omniai/chat/content/file_spec.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+RSpec.describe OmniAI::Chat::Content::File do
+  subject(:file) { described_class.new(io, type) }
+
+  let(:io) do
+    Tempfile.new.tap do |tempfile|
+      tempfile.write('Hello!')
+      tempfile.rewind
+    end
+  end
+
+  let(:type) { 'text/plain' }
+
+  around do |example|
+    example.call
+  ensure
+    io.close
+    io.unlink
+  end
+
+  describe '#type' do
+    it { expect(file.type).to eq('text/plain') }
+  end
+
+  describe '#io' do
+    it { expect(file.io).to eq(io) }
+  end
+
+  describe '#fetch!' do
+    it { expect(file.fetch!).to eql('Hello!') }
+  end
+
+  describe '#data' do
+    it { expect(file.data).to eq('SGVsbG8h') }
+  end
+
+  describe '#data_uri' do
+    it { expect(file.data_uri).to eq('data:text/plain;base64,SGVsbG8h') }
+  end
+end

spec/omniai/chat/content/media_spec.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+RSpec.describe OmniAI::Chat::Content::Media do
+  subject(:media) { described_class.new(type) }
+
+  let(:type) { 'text/plain' }
+
+  describe '#type' do
+    it { expect(media.type).to eq('text/plain') }
+  end
+
+  describe '#fetch!' do
+    it { expect { media.fetch! }.to raise_error(NotImplementedError) }
+  end
+
+  describe '#text?' do
+    context 'when type is text/plain' do
+      let(:type) { 'text/plain' }
+
+      it { expect(media).to be_text }
+    end
+
+    context 'when type is application/pdf' do
+      let(:type) { 'application/pdf' }
+
+      it { expect(media).not_to be_text }
+    end
+  end
+
+  describe '#audio?' do
+    context 'when type is audio/flac' do
+      let(:type) { 'audio/flac' }
+
+      it { expect(media).to be_audio }
+    end
+
+    context 'when type is application/pdf' do
+      let(:type) { 'application/pdf' }
+
+      it { expect(media).not_to be_audio }
+    end
+  end
+
+  describe '#image?' do
+    context 'when type is image/jpeg' do
+      let(:type) { 'image/jpeg' }
+
+      it { expect(media).to be_image }
+    end
+
+    context 'when type is application/pdf' do
+      let(:type) { 'application/pdf' }
+
+      it { expect(media).not_to be_image }
+    end
+  end
+
+  describe '#video?' do
+    context 'when type is video/mpeg' do
+      let(:type) { 'video/mpeg' }
+
+      it { expect(media).to be_video }
+    end
+
+    context 'when type is application/pdf' do
+      let(:type) { 'application/pdf' }
+
+      it { expect(media).not_to be_video }
+    end
+  end
+end

spec/omniai/chat/content/text_spec.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+RSpec.describe OmniAI::Chat::Content::Text do
+  subject(:text) { described_class.new('Hello!') }
+
+  describe '#text' do
+    it { expect(text.text).to eq('Hello!') }
+  end
+end

spec/omniai/chat/content/url_spec.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+RSpec.describe OmniAI::Chat::Content::URL do
+  subject(:url) { described_class.new('https://localhost/greeting.txt', 'text/plain') }
+
+  describe '#fetch!' do
+    before do
+      stub_request(:get, 'https://localhost/greeting.txt')
+        .to_return(body: 'Hello!', status: 200)
+    end
+
+    it { expect(url.fetch!).to eql('Hello!') }
+  end
+
+  describe '#data' do
+    before do
+      stub_request(:get, 'https://localhost/greeting.txt')
+        .to_return(body: 'Hello!', status: 200)
+    end
+
+    it { expect(url.data).to eq('SGVsbG8h') }
+  end
+
+  describe '#data_uri' do
+    before do
+      stub_request(:get, 'https://localhost/greeting.txt')
+        .to_return(body: 'Hello!', status: 200)
+    end
+
+    it { expect(url.data_uri).to eq('data:text/plain;base64,SGVsbG8h') }
+  end
+end