Sha256: e347b419a53d76f17d07d291af35b67d9ca1bd764d7fcced1d6a0d7dd41c081b

Contents?: true

Size: 1.57 KB

Versions: 2

Compression:

Stored size: 1.57 KB

Contents

# frozen_string_literal: true

module GenAI
  class Language
    class OpenAI < Base
      include GenAI::Api::Format::OpenAI

      EMBEDDING_MODEL = 'text-embedding-ada-002'
      COMPLETION_MODEL = 'gpt-3.5-turbo-1106'

      def initialize(token:, options: {})
        depends_on 'ruby-openai'

        @client = ::OpenAI::Client.new(access_token: token)
      end

      def embed(input, model: nil)
        parameters = { input: input, model: model || EMBEDDING_MODEL }

        response = handle_errors { client.embeddings(parameters: parameters) }

        build_result(model: parameters[:model], raw: response, parsed: extract_embeddings(response))
      end

      def complete(prompt, options = {})
        parameters = build_completion_options(prompt, options)

        response = handle_errors { client.chat(parameters: parameters) }

        build_result(model: parameters[:model], raw: response, parsed: extract_completions(response))
      end

      def chat(messages, options = {})
        parameters = {
          messages: messages.map(&:deep_symbolize_keys),
          model: options.delete(:model) || COMPLETION_MODEL
        }.merge(options)

        response = handle_errors { client.chat(parameters: parameters) }

        build_result(model: parameters[:model], raw: response, parsed: extract_completions(response))
      end

      private

      def build_completion_options(prompt, options)
        {
          messages: [{ role: DEFAULT_ROLE, content: prompt }],
          model: options.delete(:model) || COMPLETION_MODEL
        }.merge(options)
      end
    end
  end
end

Version data entries

2 entries across 2 versions & 1 rubygems

Version Path
gen-ai-0.4.1 lib/gen_ai/language/open_ai.rb
gen-ai-0.4.0 lib/gen_ai/language/open_ai.rb