gpt-storyteller/src/main.cr

require "option_parser"
require "pretty_print"
require "openai"
require "spinner"

require "./zone"
require "./parsers/prompt_string"
require "./builders/prompt_string"
require "./builders/openai_chat"
require "./builders/openai_insert"

class Storyteller

  enum OpenAIMode
    Chat
    Complete
    Insert

    def self.from_s(str_mode)
      return OpenAIMode.values.find do |openai_mode|
        openai_mode.to_s.downcase == str_mode.downcase
      end
    end
  end

  def initialize()
  end

  def self.start(argv)
    input_file = STDIN
    input_file_path = ""
    output_file = STDOUT
    output_file_path = ""
    past_characters_limit = 1000
    use_color = true
    make_request = true
    gpt_mode = OpenAIMode::Chat
    verbose = false
    gpt_temperature = 0.82
    gpt_presence_penalty = 1
    gpt_frequency_penalty = 1
    gpt_max_tokens = 256

    parser = OptionParser.new do |parser|
      parser.banner = "Usage: storyteller [options]"

      parser.separator("Options:")


      parser.on("-i FILE", "--input=FILE", "Path to input file") do |file|
        input_file_path = file
      end

      parser.on("-h", "--help", "Show this help") do
        puts parser
        exit
      end


      parser.on("-n", "--no-color", "Disable color output") do
        use_color = false
      end

      parser.on("-o FILE", "--output=FILE", "Path to output file") do |file|
        use_color = false
        output_file_path = file
      end

      parser.on("-v", "--verbose", "Be verbose (cumulative)") do
        verbose = true
      end

      parser.on("--dry-run", "Don't call the API") do
        make_request = false
      end


      parser.separator("GPT options")

      parser.on("--gpt-mode MODE", "GPT mode (chat,insert,complete) (default: chat)") do |chosen_mode|
        result_mode = OpenAIMode.from_s(chosen_mode.downcase)
        if result_mode.nil?
          STDERR.puts "ERROR: unknown mode #{chosen_mode}"
          exit 1
        end
        gpt_mode = result_mode unless result_mode.nil?
      end

      parser.on("--gpt-temperature TEMPERATURE", "GPT Temperature") do |temperature|
        gpt_temperature = temperature
      end
      parser.on("--gpt-presence-penalty PENALTY", "GPT Presence Penalty") do |presence_penalty|
        gpt_presence_penalty = presence_penalty
      end
      parser.on("--gpt-frequency-penalty PENALTY", "GPT Frequency Penalty") do |frequency_penalty|
        gpt_frequency_penalty = frequency_penalty
      end
      parser.on("--gpt-max-tokens TOKENS", "GPT Max Tokens") do |max_tokens|
        gpt_max_tokens = max_tokens
      end

    end
    parser.parse(ARGV)

    # Create Storyteller instance
    storyteller = Storyteller.new()

    # Read file and initialize zones
    if !input_file_path.empty?
      # puts "d: Using input file #{input_file_path}"
      input_file = File.open(input_file_path)
    end
    prompt = storyteller.read_file(input_file)
    input_file.close

    # Build GPT-3 request
    prompt = storyteller.complete(prompt, make_request, verbose)
    pp prompt if verbose
    exit 0 if !make_request

    if !output_file_path.empty?
      # puts "d: Using output file #{input_file_path}"
      output_file = File.open(output_file_path, "w")
    end
    storyteller.write_file(output_file, prompt, use_color)
    output_file.close
  rescue ex : OptionParser::InvalidOption
    STDERR.puts parser
    STDERR.puts "\nERROR: #{ex.message}"
    exit 1
  end

  def complete(prompt : Prompt, make_request : Bool, verbose : Bool)
    builder = Builder::OpenAIChat.new(verbose: verbose)
    messages = builder.build(prompt)

    STDERR.puts messages if verbose

    return prompt if !make_request

    channel_ready = Channel(Bool).new
    channel_tick = Channel(Bool).new
    # sp = Spin.new(0.5, Spinner::Charset[:progress])

    spawn do
      tick = 0
      loop do
        tick += 1
        print "."
        if tick > (3 * 60)
          print "(timeout)"
          exit 1
        end
        channel_tick.send(true) 
        sleep 1.seconds
      end
    end

    spawn do
      openai = OpenAI::Client.new(access_token: ENV.fetch("OPENAI_API_KEY"))
      result = openai.chat(
        "gpt-3.5-turbo", 
        messages,
        {
          "temperature" => 0.82,
          "presence_penalty" => 1,
          "frequency_penalty" => 1,
          "max_tokens" => 256
        }
      )
      prompt.past_zone.content << "\n" + result.choices.first["message"]["content"] + "\n"
      channel_ready.send(true)
    end

    # sp.start
    channel_ready.receive
    channel_ready.close
    # sp.stop

    prompt
  end

  def read_file(input_file : IO::FileDescriptor)
    content = input_file.gets_to_end

    # puts "d: building parser"
    parser = Parser::PromptString.new
    # puts "d: parsing"
    prompt = parser.parse(content)
    # pp prompt
  end

  def write_file(output_file : IO::FileDescriptor, prompt : Prompt, use_color : Bool)
    # STDERR.puts "d: building builder"
    builder = Builder::PromptString.new(use_color)
    # STDERR.puts "d: building"
    text = builder.build(prompt)
    output_file.write_string(text.to_slice)
  end

  def display_completion(completion : String)
    # Code pour afficher la complétion
  end
end

Storyteller.start(ARGV)