ai-chat 0.5.8 → 0.6.1

data/ai-chat.gemspec

@@ -2,7 +2,7 @@
 
 Gem::Specification.new do |spec|
   spec.name = "ai-chat"
-  spec.version = "0.5.8"
+  spec.version = "0.6.1"
   spec.authors = ["Raghu Betina", "Jelani Woods"]
   spec.email = ["raghu@firstdraft.com", "jelani@firstdraft.com"]
   spec.homepage = "https://github.com/firstdraft/ai-chat"
@@ -19,7 +19,7 @@ Gem::Specification.new do |spec|
   }
 
   spec.required_ruby_version = ">= 3.2"
-  spec.add_runtime_dependency "openai", "~> 0.43"
+  spec.add_runtime_dependency "openai", "~> 0.59"
   spec.add_runtime_dependency "marcel", "~> 1.0"
   spec.add_runtime_dependency "base64", "~> 0.1", "> 0.1.1"
   spec.add_runtime_dependency "json", "~> 2.0"

data/lib/ai/chat.rb

@@ -18,14 +18,19 @@ module AI
   # :reek:IrresponsibleModule
   class Chat
     # :reek:Attribute
-    attr_accessor :background, :code_interpreter, :conversation_id, :image_generation, :image_folder, :messages, :model, :reasoning_effort, :web_search
-    attr_reader :client, :last_response_id, :proxy, :schema, :schema_file, :verbosity
+    attr_accessor :background, :code_interpreter, :conversation_id, :image_folder, :messages, :model, :reasoning_effort, :web_search
+    attr_reader :client, :image_generation, :last_response_id, :proxy, :schema, :schema_file, :verbosity
 
     BASE_PROXY_URL = "https://prepend.me/api.openai.com/v1"
-
-    def initialize(api_key: nil, api_key_env_var: nil)
-      @api_key = self.class.resolve_api_key(api_key: api_key, api_key_env_var: api_key_env_var)
-      @proxy = ENV["AICHAT_PROXY"] == "true"
+    PROXY_ENV = "AICHAT_PROXY"
+    PROXY_KEY_ENV = "AICHAT_PROXY_KEY"
+    OPENAI_KEY_ENV = "OPENAI_API_KEY"
+
+    def initialize(api_key: nil, api_key_env_var: nil, proxy: nil)
+      @api_key_arg = api_key
+      @api_key_env_var_arg = api_key_env_var
+      @proxy = proxy.nil? ? ENV[PROXY_ENV]&.downcase == "true" : !!proxy
+      @api_key = resolve_api_key
       @messages = []
       @reasoning_effort = nil
       @model = "gpt-5.2"
@@ -40,8 +45,8 @@ module AI
     end
 
     def self.generate_schema!(description, location: "schema.json", api_key: nil, api_key_env_var: nil, proxy: nil)
-      api_key = resolve_api_key(api_key: api_key, api_key_env_var: api_key_env_var)
-      proxy = ENV["AICHAT_PROXY"] == "true" if proxy.nil?
+      proxy = proxy.nil? ? ENV[PROXY_ENV]&.downcase == "true" : !!proxy
+      api_key = api_key || ENV.fetch(api_key_env_var || (proxy ? PROXY_KEY_ENV : OPENAI_KEY_ENV))
       prompt_path = File.expand_path("../prompts/schema_generator.md", __dir__)
       system_prompt = File.read(prompt_path)
 
@@ -73,16 +78,6 @@ module AI
       content
     end
 
-    def self.resolve_api_key(api_key: nil, api_key_env_var: nil)
-      return api_key if api_key
-      return ENV.fetch(api_key_env_var) if api_key_env_var
-
-      aichat_api_key = ENV["AICHAT_API_KEY"]
-      return aichat_api_key if aichat_api_key && !aichat_api_key.empty?
-
-      ENV.fetch("OPENAI_API_KEY")
-    end
-
     # :reek:TooManyStatements
     # :reek:NilCheck
     def add(content, role: "user", response: nil, status: nil, image: nil, images: nil, file: nil, files: nil)
@@ -169,15 +164,20 @@ module AI
     end
 
     def proxy=(value)
-      @proxy = value
-      @client = if value
-        OpenAI::Client.new(
-          api_key: @api_key,
-          base_url: BASE_PROXY_URL
-        )
-      else
-        OpenAI::Client.new(api_key: @api_key)
-      end
+      new_proxy = !!value
+      previous_proxy = @proxy
+      @proxy = new_proxy
+      new_key = resolve_api_key
+      client_options = {api_key: new_key}
+      client_options[:base_url] = BASE_PROXY_URL if new_proxy
+      new_client = OpenAI::Client.new(**client_options)
+
+      @api_key = new_key
+      @api_key_validated = false
+      @client = new_client
+    rescue => error
+      @proxy = previous_proxy
+      raise
     end
 
     def schema=(value)
@@ -205,6 +205,19 @@ module AI
       end
     end
 
+    def image_generation=(value)
+      case value
+      when true
+        @image_generation = true
+      when false, nil
+        @image_generation = false
+      when Hash
+        @image_generation = value.transform_keys(&:to_sym)
+      else
+        raise ArgumentError, "Invalid image_generation value: #{value.inspect}. Must be true, false, or a Hash of tool options (e.g., { size: \"1536x1024\", quality: \"low\" })."
+      end
+    end
+
     def last
       messages.last
     end
@@ -262,6 +275,21 @@ module AI
 
     private
 
+    def resolve_api_key
+      env_var = @api_key_env_var_arg || (@proxy ? PROXY_KEY_ENV : OPENAI_KEY_ENV)
+      @api_key_arg || ENV.fetch(env_var) {
+        if @proxy
+          raise KeyError, "Proxy mode is enabled but #{PROXY_KEY_ENV} is not set. " \
+            "Create an environment variable called #{PROXY_KEY_ENV} " \
+            "with your API key from Prepend.me."
+        else
+          raise KeyError, "#{OPENAI_KEY_ENV} is not set. " \
+            "Create an environment variable called #{OPENAI_KEY_ENV} " \
+            "with your API key from https://platform.openai.com/api-keys."
+        end
+      }
+    end
+
     class InputClassificationError < StandardError; end
 
     class WrongAPITokenUsedError < StandardError; end
@@ -508,7 +536,8 @@ module AI
         tools_list << {type: "web_search"}
       end
       if image_generation
-        tools_list << {type: "image_generation"}
+        options = image_generation.is_a?(Hash) ? image_generation : {}
+        tools_list << options.merge(type: "image_generation")
       end
       if code_interpreter
         tools_list << {type: "code_interpreter", container: {type: "auto"}}
@@ -562,7 +591,8 @@ module AI
           result = output.result
           image_data = Base64.strict_decode64(result)
 
-          filename = "#{(index + 1).to_s.rjust(3, "0")}.png"
+          extension = image_extension_for(image_data)
+          filename = "#{(index + 1).to_s.rjust(3, "0")}.#{extension}"
           file_path = File.join(subfolder_path, filename)
 
           File.binwrite(file_path, image_data)
@@ -574,6 +604,18 @@ module AI
       image_filenames
     end
 
+    IMAGE_EXTENSIONS_BY_MIME_TYPE = {
+      "image/png" => "png",
+      "image/jpeg" => "jpg",
+      "image/webp" => "webp"
+    }.freeze
+
+    # :reek:UtilityFunction
+    def image_extension_for(bytes)
+      mime_type = Marcel::MimeType.for(StringIO.new(bytes))
+      IMAGE_EXTENSIONS_BY_MIME_TYPE.fetch(mime_type, "png")
+    end
+
     def create_images_folder(response_id)
       # ISO 8601 basic format with centisecond precision
       timestamp = Time.now.strftime("%Y%m%dT%H%M%S%2N")
@@ -597,11 +639,11 @@ module AI
     rescue OpenAI::Errors::AuthenticationError
       message = if proxy
         <<~STRING
-          It looks like you're using an invalid API key. Proxying is enabled, so you must use an OpenAI API key from prepend.me. Please disable proxy or update your API key before generating a response.
+          Your API key was not accepted by Prepend.me. Since proxy mode is enabled, you need a valid API key from Prepend.me in the #{PROXY_KEY_ENV} environment variable.
         STRING
       else
         <<~STRING
-          It looks like you're using an invalid API key. Check to make sure your API key is valid before generating a response.
+          Your API key was not accepted by OpenAI. Make sure the #{OPENAI_KEY_ENV} environment variable contains a valid key from https://platform.openai.com/api-keys.
         STRING
       end
       raise WrongAPITokenUsedError, message, cause: nil

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ai-chat
 version: !ruby/object:Gem::Version
-  version: 0.5.8
+  version: 0.6.1
 platform: ruby
 authors:
 - Raghu Betina
@@ -16,14 +16,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.43'
+        version: '0.59'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.43'
+        version: '0.59'
 - !ruby/object:Gem::Dependency
   name: marcel
   requirement: !ruby/object:Gem::Requirement