ai-chat 0.2.3 → 0.3.0

data/lib/ai/chat.rb

@@ -4,9 +4,15 @@ require "base64"
 require "json"
 require "marcel"
 require "openai"
+require "ostruct"
 require "pathname"
 require "stringio"
 require "fileutils"
+require "tty-spinner"
+require "timeout"
+
+require_relative "http"
+include AI::Http
 
 module AI
   # :reek:MissingSafeMethod { exclude: [ generate! ] }
@@ -16,34 +22,72 @@ module AI
   # :reek:IrresponsibleModule
   class Chat
     # :reek:Attribute
-    attr_accessor :messages, :model, :web_search, :previous_response_id, :image_generation, :image_folder
+    attr_accessor :background, :code_interpreter, :conversation_id, :image_generation, :image_folder, :messages, :model, :proxy, :previous_response_id, :web_search
     attr_reader :reasoning_effort, :client, :schema
 
     VALID_REASONING_EFFORTS = [:low, :medium, :high].freeze
+    PROXY_URL = "https://prepend.me/".freeze
 
     def initialize(api_key: nil, api_key_env_var: "OPENAI_API_KEY")
-      api_key ||= ENV.fetch(api_key_env_var)
+      @api_key = api_key || ENV.fetch(api_key_env_var)
       @messages = []
       @reasoning_effort = nil
       @model = "gpt-4.1-nano"
-      @client = OpenAI::Client.new(api_key: api_key)
+      @client = OpenAI::Client.new(api_key: @api_key)
       @previous_response_id = nil
+      @proxy = false
       @image_generation = false
       @image_folder = "./images"
     end
 
+    def self.generate_schema!(description, api_key: nil, api_key_env_var: "OPENAI_API_KEY", proxy: false)
+      @api_key ||= ENV.fetch(api_key_env_var)
+      prompt_path = File.expand_path("../prompts/schema_generator.md", __dir__)
+      system_prompt = File.open(prompt_path).read
+
+      json = if proxy
+        uri = URI(PROXY_URL + "api.openai.com/v1/responses")
+        parameters = {
+          model: "o4-mini",
+          input: [
+            {role: :system, content: system_prompt},
+            {role: :user, content: description},
+          ],
+          text: {format: {type: "json_object"}},
+          reasoning: {effort: "high"}
+        }
+
+        send_request(uri, content_type: "json", parameters: parameters, method: "post")
+      else
+        client = OpenAI::Client.new(api_key: api_key)
+        response = client.responses.create(
+          model: "o4-mini",
+          input: [
+            {role: :system, content: system_prompt},
+            {role: :user, content: description}
+          ],
+          text: {format: {type: "json_object"}},
+          reasoning: {effort: "high"}
+        )
+
+        output_text = response.output_text
+        JSON.parse(output_text)
+      end
+      JSON.pretty_generate(json)
+    end
+
     # :reek:TooManyStatements
     # :reek:NilCheck
-    def add(content, role: "user", response: nil, image: nil, images: nil, file: nil, files: nil)
+    def add(content, role: "user", response: nil, status: nil, image: nil, images: nil, file: nil, files: nil)
       if image.nil? && images.nil? && file.nil? && files.nil?
-        messages.push(
-          {
-            role: role,
-            content: content,
-            response: response
-          }.compact
-        )
-
+        message = {
+          role: role,
+          content: content,
+          response: response
+        }
+        message[:content] = content if content
+        message[:status] = status if status
+        messages.push(message)
       else
         text_and_files_array = [
           {
@@ -76,7 +120,8 @@ module AI
         messages.push(
           {
             role: role,
-            content: text_and_files_array
+            content: text_and_files_array,
+            status: status
           }
         )
       end
@@ -90,53 +135,32 @@ module AI
       add(message, role: "user", image: image, images: images, file: file, files: files)
     end
 
-    def assistant(message, response: nil)
-      add(message, role: "assistant", response: response)
+    def assistant(message, response: nil, status: nil)
+      add(message, role: "assistant", response: response, status: status)
     end
 
     # :reek:NilCheck
     # :reek:TooManyStatements
     def generate!
+      validate_api_key
       response = create_response
+      parse_response(response)
 
-      text_response = extract_text_from_response(response)
-
-      image_filenames = extract_and_save_images(response)
-      response_usage = response.usage.to_h.slice(:input_tokens, :output_tokens, :total_tokens)
-
-      chat_response = {
-        id: response.id,
-        model: response.model,
-        usage: response_usage,
-        total_tokens: response_usage[:total_tokens],
-        images: image_filenames
-      }
-
-      message = if schema
-        if text_response.nil? || text_response.empty?
-          raise ArgumentError, "No text content in response to parse as JSON for schema: #{schema.inspect}"
-        end
-        JSON.parse(text_response, symbolize_names: true)
-      else
-        text_response
-      end
+      self.previous_response_id = last.dig(:response, :id) unless (conversation_id && !background)
+      last
+    end
 
-      if image_filenames.empty?
-        assistant(message, response: chat_response)
+    # :reek:BooleanParameter
+    # :reek:ControlParameter
+    # :reek:DuplicateMethodCall
+    # :reek:TooManyStatements
+    def get_response(wait: false, timeout: 600)
+      response = if wait
+        wait_for_response(timeout)
       else
-        messages.push(
-          {
-            role: "assistant",
-            content: message,
-            images: image_filenames,
-            response: chat_response
-          }.compact
-        )
+        retrieve_response(previous_response_id)
       end
-
-      self.previous_response_id = chat_response[:id]
-
-      message
+      parse_response(response)
     end
 
     # :reek:NilCheck
@@ -172,15 +196,64 @@ module AI
       messages.last
     end
 
+    def items(order: :asc)
+      raise "No conversation_id set. Call generate! first to create a conversation." unless conversation_id
+
+      if proxy
+        uri = URI(PROXY_URL + "api.openai.com/v1/conversations/#{conversation_id}/items?order=#{order.to_s}")
+        response_hash = send_request(uri, content_type: "json", method: "get")
+
+        if response_hash.key?(:data)
+          response_hash.dig(:data).map do |hash|
+            # Transform values to allow expected symbols that non-proxied request returns 
+
+            hash.transform_values! do |value|
+              if hash.key(value) == :type
+                value.to_sym
+              else
+                value
+              end
+            end
+          end
+          response_hash
+        end
+        # Convert to Struct to allow same interface as non-proxied request
+        create_deep_struct(response_hash)
+      else
+        client.conversations.items.list(conversation_id, order: order)
+      end
+    end
+
+    def verbose
+      page = items
+
+      box_width = 78
+      inner_width = box_width - 4
+
+      puts
+      puts "┌#{"─" * (box_width - 2)}┐"
+      puts "│ Conversation: #{conversation_id.ljust(inner_width - 14)} │"
+      puts "│ Items: #{page.data.length.to_s.ljust(inner_width - 7)} │"
+      puts "└#{"─" * (box_width - 2)}┘"
+      puts
+
+      ap page.data, limit: 10, indent: 2
+    end
+
     def inspect
       "#<#{self.class.name} @messages=#{messages.inspect} @model=#{@model.inspect} @schema=#{@schema.inspect} @reasoning_effort=#{@reasoning_effort.inspect}>"
     end
 
     # Support for Ruby's pp (pretty print)
+    # :reek:TooManyStatements
+    # :reek:NilCheck
+    # :reek:FeatureEnvy
+    # :reek:DuplicateMethodCall
+    # :reek:UncommunicativeParameterName
     def pretty_print(q)
-      q.group(1, "#<#{self.class}", '>') do
+      q.group(1, "#<#{self.class}", ">") do
         q.breakable
-        
+
         # Show messages with truncation
         q.text "@messages="
         truncated_messages = @messages.map do |msg|
@@ -191,7 +264,7 @@ module AI
           truncated_msg
         end
         q.pp truncated_messages
-        
+
         # Show other instance variables (except sensitive ones)
         skip_vars = [:@messages, :@api_key, :@client]
         instance_variables.sort.each do |var|
@@ -207,10 +280,10 @@ module AI
       end
     end
 
-
     private
 
     class InputClassificationError < StandardError; end
+    class WrongAPITokenUsedError < StandardError; end
 
     # :reek:FeatureEnvy
     # :reek:ManualDispatch
@@ -225,21 +298,128 @@ module AI
       end
     end
 
+    def create_conversation
+      self.conversation_id = if proxy
+        uri = URI(PROXY_URL + "api.openai.com/v1/conversations")
+        response = send_request(uri, content_type: "json", method: "post")
+        response.dig(:id)
+      else
+        conversation = client.conversations.create
+        conversation.id
+      end
+    end
+
     # :reek:TooManyStatements
     def create_response
       parameters = {
         model: model
       }
 
+      parameters[:background] = background if background
       parameters[:tools] = tools unless tools.empty?
       parameters[:text] = schema if schema
       parameters[:reasoning] = {effort: reasoning_effort} if reasoning_effort
-      parameters[:previous_response_id] = previous_response_id if previous_response_id
+
+      if previous_response_id && conversation_id
+        warn "Both conversation_id and previous_response_id are set. Using previous_response_id for forking. Only set one."
+        parameters[:previous_response_id] = previous_response_id
+      elsif previous_response_id
+        parameters[:previous_response_id] = previous_response_id
+      elsif conversation_id
+        parameters[:conversation] = conversation_id
+      else
+        create_conversation
+      end
 
       messages_to_send = prepare_messages_for_api
       parameters[:input] = strip_responses(messages_to_send) unless messages_to_send.empty?
 
-      client.responses.create(**parameters)
+      if proxy
+        uri = URI(PROXY_URL + "api.openai.com/v1/responses")
+        send_request(uri, content_type: "json", parameters: parameters, method: "post")
+      else
+        client.responses.create(**parameters)
+      end
+    end
+
+    # :reek:NilCheck
+    # :reek:TooManyStatements
+    def parse_response(response)
+      if proxy && response.is_a?(Hash)
+        response_messages = response.dig(:output).select do |output|
+          output.dig(:type) == "message"
+        end
+
+        message_contents = response_messages.map do |message|
+          message.dig(:content)
+        end.flatten
+
+        output_texts = message_contents.select do |content|
+          content[:type] == "output_text"
+        end
+
+        text_response = output_texts.map { |output| output[:text] }.join
+        response_id = response.dig(:id)
+        response_status = response.dig(:status).to_sym
+        response_model = response.dig(:model)
+        response_usage = response.dig(:usage)&.slice(:input_tokens, :output_tokens, :total_tokens)
+
+        if response.key?(:conversation)
+          self.conversation_id = response.dig(:conversation, :id)
+        end
+      else        
+        text_response = response.output_text
+        response_id = response.id
+        response_status = response.status
+        response_model = response.model
+        response_usage = response.usage.to_h.slice(:input_tokens, :output_tokens, :total_tokens)
+
+        if response.conversation
+          self.conversation_id = response.conversation.id
+        end
+      end
+      image_filenames = extract_and_save_images(response) + extract_and_save_files(response)
+
+      chat_response = {
+        id: response_id,
+        model: response_model,
+        usage: response_usage || {},
+        total_tokens: response_usage&.fetch(:total_tokens, 0),
+        images: image_filenames
+      }.compact
+
+      response_content = if schema
+        if text_response.nil? || text_response.empty?
+          raise ArgumentError, "No text content in response to parse as JSON for schema: #{schema.inspect}"
+        end
+        JSON.parse(text_response, symbolize_names: true)
+      else
+        text_response
+      end
+
+      existing_message_position = messages.find_index do |message|
+        message.dig(:response, :id) == response_id
+      end
+
+      message = {
+        role: "assistant",
+        content: response_content,
+        response: chat_response,
+        status: response_status
+      }
+
+      message.store(:images, image_filenames) unless image_filenames.empty?
+
+      if existing_message_position
+        messages[existing_message_position] = message
+      else
+        messages.push(message)
+        message
+      end
+    end
+
+    def cancel_request
+      client.responses.cancel(previous_response_id)
     end
 
     def prepare_messages_for_api
@@ -389,19 +569,12 @@ module AI
       if image_generation
         tools_list << {type: "image_generation"}
       end
+      if code_interpreter
+        tools_list << {type: "code_interpreter", container: {type: "auto"}}
+      end
       tools_list
     end
 
-    # :reek:UtilityFunction
-    # :reek:ManualDispatch
-    def extract_text_from_response(response)
-      response.output.flat_map { |output|
-        output.respond_to?(:content) ? output.content : []
-      }.compact.find { |content|
-        content.is_a?(OpenAI::Models::Responses::ResponseOutputText)
-      }&.text
-    end
-
     # :reek:FeatureEnvy
     # :reek:UtilityFunction
     def wrap_schema_if_needed(schema)
@@ -432,37 +605,235 @@ module AI
     def extract_and_save_images(response)
       image_filenames = []
 
-      image_outputs = response.output.select { |output|
-        output.respond_to?(:type) && output.type == :image_generation_call
-      }
+      if proxy
+        image_outputs = response.dig(:output).select { |output|
+          output.dig(:type) == "image_generation_call"
+        }
+      else       
+        image_outputs = response.output.select { |output|
+          output.respond_to?(:type) && output.type == :image_generation_call
+        }
+      end
 
       return image_filenames if image_outputs.empty?
 
+      response_id = proxy ? response.dig(:id) : response.id
+      subfolder_path = create_images_folder(response_id)
+
+      image_outputs.each_with_index do |output, index|
+        if proxy
+          next unless output.key?(:result) && output.dig(:result)
+        else
+          next unless output.respond_to?(:result) && output.result
+        end
+
+        warn_if_file_fails_to_save do
+          result = proxy ? output.dig(:result) : output.result
+          image_data = Base64.strict_decode64(result)
+
+          filename = "#{(index + 1).to_s.rjust(3, "0")}.png"
+          file_path = File.join(subfolder_path, filename)
+
+          File.binwrite(file_path, image_data)
+
+          image_filenames << file_path
+        end
+      end
+
+      image_filenames
+    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")
 
-      subfolder_name = "#{timestamp}_#{response.id}"
+      subfolder_name = "#{timestamp}_#{response_id}"
       subfolder_path = File.join(image_folder || "./images", subfolder_name)
       FileUtils.mkdir_p(subfolder_path)
+      subfolder_path
+    end
 
-      image_outputs.each_with_index do |output, index|
-        next unless output.respond_to?(:result) && output.result
+    def warn_if_file_fails_to_save
+      yield
+    rescue => error
+      warn "Failed to save image: #{error.message}"
+    end
 
-        begin
-          image_data = Base64.strict_decode64(output.result)
+    def validate_api_key
+      openai_api_key_used = @api_key.start_with?("sk-proj")
+      proxy_api_key_used = !openai_api_key_used
+      proxy_enabled = proxy
+      proxy_disabled = !proxy
+
+      if openai_api_key_used && proxy_enabled
+        raise WrongAPITokenUsedError, <<~STRING
+          It looks like you're using an official API key from OpenAI with proxying enabled. When proxying is enabled you must use an OpenAI API key from prepend.me. Please disable proxy or update your API key before generating a response.
+        STRING
+      elsif proxy_api_key_used && proxy_disabled
+        raise WrongAPITokenUsedError, <<~STRING
+          It looks like you're using an unofficial OpenAI API key from prepend.me. When using an unofficial API key you must enable proxy before generating a response. Proxying is currently disabled, please enable it before generating a response.
+
+          Example:
+
+            chat = AI::Chat.new
+            chat.proxy = true
+            chat.user(...)
+            chat.generate!
+        STRING
+      end
+    end
 
-          filename = "#{(index + 1).to_s.rjust(3, "0")}.png"
-          filepath = File.join(subfolder_path, filename)
+    # :reek:FeatureEnvy
+    # :reek:ManualDispatch
+    # :reek:NestedIterators
+    # :reek:TooManyStatements
+    def extract_and_save_files(response)
+      filenames = []
+
+      if proxy
+        message_outputs = response.dig(:output).select do |output|
+          output.dig(:type) == "message"
+        end
+  
+        outputs_with_annotations = message_outputs.map do |message|
+          message.dig(:content).find do |content|
+            content.dig(:annotations).length.positive?
+          end
+        end.compact
+      else
+        message_outputs = response.output.select do |output|
+          output.respond_to?(:type) && output.type == :message
+        end
+  
+        outputs_with_annotations = message_outputs.map do |message|
+          message.content.find do |content|
+            content.respond_to?(:annotations) && content.annotations.length.positive?
+          end
+        end.compact
+      end
 
-          File.binwrite(filepath, image_data)
+      return filenames if outputs_with_annotations.empty?
 
-          image_filenames << filepath
-        rescue => error
-          warn "Failed to save image: #{error.message}"
+      response_id = proxy ? response.dig(:id) : response.id
+      subfolder_path = create_images_folder(response_id)
+
+      if proxy
+        annotations = outputs_with_annotations.map do |output|
+          output.dig(:annotations).find do |annotation|
+            annotation.key?(:filename)
+          end
+        end.compact
+  
+        annotations.each do |annotation|
+          container_id = annotation.dig(:container_id)
+          file_id = annotation.dig(:file_id)
+          filename = annotation.dig(:filename)
+  
+          warn_if_file_fails_to_save do
+            file_content = retrieve_file(file_id, container_id: container_id)
+            file_path = File.join(subfolder_path, filename)
+            File.binwrite(file_path, file_content)
+            filenames << file_path
+          end
+        end
+      else
+        annotations = outputs_with_annotations.map do |output|
+          output.annotations.find do |annotation|
+            annotation.respond_to?(:filename)
+          end
+        end.compact
+  
+        annotations.each do |annotation|
+          container_id = annotation.container_id
+          file_id = annotation.file_id
+          filename = annotation.filename
+  
+          warn_if_file_fails_to_save do
+            file_content = retrieve_file(file_id, container_id: container_id)
+            file_path = File.join(subfolder_path, filename)
+            File.open(file_path, "wb") do |file|
+              file.write(file_content.read)
+            end
+            filenames << file_path
+          end
         end
       end
+      filenames
+    end
 
-      image_filenames
+    # This is similar to ActiveJob's :polynomially_longer retry option
+    # :reek:DuplicateMethodCall
+    # :reek:UtilityFunction
+    def calculate_wait(executions)
+      # cap the maximum wait time to ~110 seconds
+      executions = executions.clamp(1..10)
+      jitter = 0.15
+      ((executions**2) + (Kernel.rand * (executions**2) * jitter)) + 2
+    end
+
+    def timeout_request(duration)
+      Timeout.timeout(duration) do
+        yield
+      end
+    rescue Timeout::Error
+      client.responses.cancel(previous_response_id)
+    end
+
+    # :reek:DuplicateMethodCall
+    # :reek:TooManyStatements
+    def wait_for_response(timeout)
+        spinner = TTY::Spinner.new("[:spinner] Thinking ...", format: :dots)
+        spinner.auto_spin
+        api_response = retrieve_response(previous_response_id)
+        number_of_times_polled = 0
+        response = timeout_request(timeout) do
+          status = if api_response.respond_to?(:status)
+            api_response.status
+          else 
+            api_response.dig(:status)&.to_sym
+          end
+
+          while status != :completed
+            some_amount_of_seconds = calculate_wait(number_of_times_polled)
+            sleep some_amount_of_seconds
+            number_of_times_polled += 1
+            api_response = retrieve_response(previous_response_id)
+            status = if api_response.respond_to?(:status)
+              api_response.status
+            else
+              api_response.dig(:status)&.to_sym
+            end
+          end
+          api_response
+        end
+        
+        status = if api_response.respond_to?(:status)
+          api_response.status
+        else 
+          api_response.dig(:status).to_sym
+        end
+        exit_message = status == :cancelled ? "request timed out" : "done!"
+        spinner.stop(exit_message)
+        response
+    end
+
+    def retrieve_response(previous_response_id)
+      if proxy
+        uri = URI(PROXY_URL + "api.openai.com/v1/responses/#{previous_response_id}")
+        send_request(uri, content_type: "json", method: "get")
+      else
+        client.responses.retrieve(previous_response_id)
+      end
+    end
+
+    def retrieve_file(file_id, container_id: nil)
+      if proxy
+        uri = URI(PROXY_URL + "api.openai.com/v1/containers/#{container_id}/files/#{file_id}/content")
+        send_request(uri, method: "get")
+      else
+        container_content = client.containers.files.content
+        file_content = container_content.retrieve(file_id, container_id: container_id)
+      end
     end
   end
 end