Class: Langchain::LLM::AI21

Inherits:
Base
  • Object
show all
Defined in:
lib/langchain/llm/ai21.rb

Overview

Wrapper around AI21 Studio APIs.

Gem requirements:

gem "ai21", "~> 0.2.1"

Usage:

llm = Langchain::LLM::AI21.new(api_key: ENV["AI21_API_KEY"])

Constant Summary collapse

DEFAULTS = 
{
  temperature: 0.0,
  model: "j2-ultra"
}.freeze

Instance Attribute Summary

Attributes inherited from Base

#client, #defaults

Instance Method Summary collapse

Methods inherited from Base

#chat, #chat_parameters, #default_dimension, #default_dimensions, #embed

Methods included from DependencyHelper

#depends_on

Constructor Details

#initialize(api_key:, default_options: {}) ⇒ AI21

Returns a new instance of AI21.



19
20
21
22
23
24
 
# File 'lib/langchain/llm/ai21.rb', line 19

def initialize(api_key:, default_options: {})
  depends_on "ai21"

  @client = ::AI21::Client.new(api_key)
  @defaults = DEFAULTS.merge(default_options)
end

Instance Method Details

#complete(prompt:, **params) ⇒ Langchain::LLM::AI21Response

Generate a completion for a given prompt

Parameters:

  • prompt (String) —

    The prompt to generate a completion for

  • params (Hash) —

    The parameters to pass to the API

Returns:



33
34
35
36
37
38
 
# File 'lib/langchain/llm/ai21.rb', line 33

def complete(prompt:, **params)
  parameters = complete_parameters params

  response = client.complete(prompt, parameters)
  Langchain::LLM::AI21Response.new response, model: parameters[:model]
end

#summarize(text:, **params) ⇒ String

Generate a summary for a given text

Parameters:

  • text (String) —

    The text to generate a summary for

  • params (Hash) —

    The parameters to pass to the API

Returns:

  • (String) —

    The summary



47
48
49
50
51
 
# File 'lib/langchain/llm/ai21.rb', line 47

def summarize(text:, **params)
  response = client.summarize(text, "TEXT", params)
  response.dig(:summary)
  # Should we update this to also return a Langchain::LLM::AI21Response?
end