ai_prompt_response.rb

# frozen_string_literal: true

# == Schema Information
#
# Table name: ai_prompt_responses
#
#  id             :bigint           not null, primary key
#  ai_model       :text
#  completed_at   :datetime
#  error_message  :text
#  prompt_options :jsonb
#  prompt_text    :text
#  response_text  :text
#  started_at     :datetime
#  created_at     :datetime         not null
#  updated_at     :datetime         not null
#
module Ai
  class PromptResponse < ApplicationRecord
    FLUSH_INTERVAL = 100

    NOVA_LITE = 'amazon.nova-lite-v1:0'
    NOVA_PRO = 'us.amazon.nova-pro-v1:0'
    NOVA_MICRO = 'us.amazon.nova-micro-v1:0'

    AI_MODELS = {
      'Nova Lite' => NOVA_LITE,
      'Nova Pro' => NOVA_PRO,
      'Nova Micro' => NOVA_MICRO,
    }.freeze

    DEFAULT_INFERENCE_CONFIG = {
      max_tokens: 512,
      temperature: 0.5,
      top_p: 0.9,
    }.freeze

    self.table_name = 'ai_prompt_responses'

    has_one_attached :prompt_image

    attribute :ai_model, :string, default: NOVA_LITE
    attribute :prompt_options, :jsonb, default: DEFAULT_INFERENCE_CONFIG

    validates :ai_model, presence: true
    validates :prompt_text, presence: true
    validates :prompt_options, presence: true

    after_update_commit lambda {
      broadcast_replace_later_to self,
                                 partial: "admin/prompt_responses/prompt_response",
                                 locals: { prompt_response: self }

      broadcast_replace_later_to "prompt_responses",
                                 target: ActionView::RecordIdentifier.dom_id(self, :row),
                                 partial: "admin/prompt_responses/prompt_response_row",
                                 locals: { prompt_response: self }
    }

    def invoke
      raise ArgumentError, "Must have ai_model, prompt_text, and prompt_options" if ai_model.blank? || prompt_text.blank? || prompt_options.empty?
      raise ArgumentError, "Cannot reinvoke a completed prompt" if completed_at.present?

      aws_credentials = Aws::Credentials.new(ENV.fetch('AWS_BEDROCK_ACCESS_KEY_ID', nil), ENV.fetch('AWS_BEDROCK_SECRET_ACCESS_KEY', nil))
      bedrock_client = Aws::BedrockRuntime::Client.new(region: 'us-east-1', credentials: aws_credentials)

      params = {
        model_id: ai_model,
        inference_config: prompt_options,
        messages: [
          {
            role: 'user',
            content: build_content,
          },
        ],
      }

      bedrock_client.converse_stream(params) do |stream|
        buffer = +""
        deltas = 0

        stream.on_content_block_delta_event do |event|
          deltas += 1
          buffer << event.dig(:delta, :text)

          if (deltas % FLUSH_INTERVAL).zero?
            update!(response_text: buffer)
            yield self if block_given?
          end
        end

        stream.on_message_start_event do |_event|
          update!(started_at: Time.current)
          yield self if block_given?
        end

        stream.on_message_stop_event do |_event|
          update!(response_text: buffer, completed_at: Time.current)
          yield self if block_given?
        end

        stream.on_internal_server_exception_event do |event|
          update!(response_text: buffer, error_message: event.dig(:message))
          yield self if block_given?
        end
      end
    end

    def max_tokens
      prompt_options["max_tokens"]
    end

    def max_tokens=(value)
      prompt_options["max_tokens"] = value
    end

    def temperature
      prompt_options["temperature"]
    end

    def temperature=(value)
      prompt_options["temperature"] = value
    end

    def top_p
      prompt_options["top_p"]
    end

    def top_p=(value)
      prompt_options["top_p"] = value
    end

    private

    def build_content
      content = [{ text: prompt_text }]

      if prompt_image.attached?
        content << {
          image: {
            format: prompt_image.content_type.split('/').last,
            source: {
              bytes: prompt_image.download,
            },
          },
        }
      end

      content
    end
  end
end