ai21

Documentation

Overview

Package ai21 provides a client for interacting with the ai21 text completion API.

Index

• Variables
• type Client
  • func New(apiKey string, optFns ...func(o *Options)) *Client
  • func (c *Client) CreateCompletion(ctx context.Context, model string, req *CompleteRequest) (*CompleteResponse, error)
• type CompleteRequest
• type CompleteResponse
• type Completion
• type Data
• type FinishReason
• type GeneratedToken
• type HTTPClient
• type Options
• type Penalty
• type Prompt
• type TextRange
• type Tokens

Variables

var SupportedModels = []string{"j2-light", "j2-mid", "j2-ultra"}

SupportedModels is a list of supported AI21 models.

Types

type Client

type Client struct {
	// contains filtered or unexported fields
}

Client is a client for interacting with the ai21 API.

func New

func New(apiKey string, optFns ...func(o *Options)) *Client

New creates a new ai21 Client instance with the provided API key.

func (*Client) CreateCompletion

func (c *Client) CreateCompletion(ctx context.Context, model string, req *CompleteRequest) (*CompleteResponse, error)

CreateCompletion sends a text completion request to the ai21 API and returns the response.

type CompleteRequest

type CompleteRequest struct {
	Prompt           string   `json:"prompt"`
	NumResults       int      `json:"numResults,omitempty"`
	MaxTokens        int      `json:"maxTokens,omitempty"`
	MinTokens        int      `json:"minTokens,omitempty"`
	Temperature      float64  `json:"temperature,omitempty"`
	TopP             float64  `json:"topP,omitempty"`
	StopSequences    []string `json:"stopSequences,omitempty"`
	TopKReturn       int      `json:"topKReturn,omitempty"`
	FrequencyPenalty Penalty  `json:"frequencyPenalty,omitempty"`
	PresencePenalty  Penalty  `json:"presencePenalty,omitempty"`
	CountPenalty     Penalty  `json:"countPenalty,omitempty"`
}

CompleteRequest represents a request for text completion.

type CompleteResponse

type CompleteResponse struct {
	ID          string       `json:"id"`
	Prompt      Prompt       `json:"prompt"`
	Completions []Completion `json:"completions"`
}

CompleteResponse represents the response from a text completion request.

type Completion

type Completion struct {
	Data         Data         `json:"data"`
	FinishReason FinishReason `json:"finishReason"`
}

Completion represents a completion result.

type Data

type Data struct {
	Text   string   `json:"text"`
	Tokens []Tokens `json:"tokens"`
}

Data represents the completion text and tokens.

type FinishReason

type FinishReason struct {
	Reason string `json:"reason"`
	Length int    `json:"length"`
}

FinishReason represents the reason and length for completion finishing.

type GeneratedToken

type GeneratedToken struct {
	Token      string  `json:"token"`
	Logprob    float64 `json:"logprob"`
	RawLogprob float64 `json:"raw_logprob"`
}

GeneratedToken represents a generated token with log probabilities.

type HTTPClient

type HTTPClient interface {
	Do(req *http.Request) (*http.Response, error)
}

HTTPClient is an interface for making HTTP requests.

type Options

type Options struct {
	// The base URL of the ai21 API.
	APIUrl string

	// The HTTP client to use for making API requests.
	HTTPClient HTTPClient
}

Options represents configuration options for the ai21 client.

type Penalty

type Penalty struct {
	Scale               int  `json:"scale"`
	ApplyToNumbers      bool `json:"applyToNumbers"`
	ApplyToPunctuations bool `json:"applyToPunctuations"`
	ApplyToStopwords    bool `json:"applyToStopwords"`
	ApplyToWhitespaces  bool `json:"applyToWhitespaces"`
	ApplyToEmojis       bool `json:"applyToEmojis"`
}

Penalty represents penalty options for text completion.

type Prompt

type Prompt struct {
	Text   string   `json:"text"`
	Tokens []Tokens `json:"tokens"`
}

Prompt represents the prompt text and tokens.

type TextRange

type TextRange struct {
	Start int `json:"start"`
	End   int `json:"end"`
}

TextRange represents a range of text in the prompt.

type Tokens

type Tokens struct {
	GeneratedToken GeneratedToken `json:"generatedToken"`
	TopTokens      interface{}    `json:"topTokens"`
	TextRange      TextRange      `json:"textRange"`
}

Tokens represents generated tokens and top tokens.