chat_completions.go

package openai

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/fabiustech/openai/models"
	"github.com/fabiustech/openai/objects"
	"github.com/fabiustech/openai/params"
	"github.com/fabiustech/openai/routes"
)

// ChatRole is an enum of the various message roles.
type ChatRole string

const (
	// System represents a system message, which helps set the behavior of the assistant.
	System ChatRole = "system"
	// User represents a message from a user, which helps instruct the assistant. They can be generated by the end users
	// of an application, or set by a developer as an instruction.
	User ChatRole = "user"
	// Assistant represents a message from the assistant. Assistant messages help store prior responses. They can also
	// be written by a developer to help give examples of desired behavior.
	Assistant ChatRole = "assistant"
	// RoleFunction represents a call to a function.
	RoleFunction ChatRole = "function"
)

// FunctionCallResponse represents a response from a function call.
type FunctionCallResponse struct {
	// Name is the name of the function to be called.
	Name string `json:"name"`
	// Arguments are the arguments to the function call, encoded as a JSON string.
	Arguments string `json:"arguments"`
}

// ChatMessage represents a message in a chat completion.
type ChatMessage struct {
	Role         ChatRole              `json:"role"`
	Content      string                `json:"content,omitempty"`
	FunctionCall *FunctionCallResponse `json:"function_call,omitempty"`
}

// Unmarshal unmarshals the content of the message into the provided value.
func (c *FunctionCallResponse) Unmarshal(val any) error {
	return json.Unmarshal([]byte(c.Arguments), val)
}

// Function is a list of functions the model may generate JSON inputs for.
type Function struct {
	// Name is the name of the function to be called. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a
	// maximum length of 64.
	Name string `json:"name"`
	// Description is the description of what the function does.
	Description string `json:"description"`
	// Parameters are the parameters the functions accepts, described as a JSON Schema object. See the guide
	// (https://platform.openai.com/docs/guides/gpt/function-calling) for examples, and the JSON Schema reference
	// (https://json-schema.org/understanding-json-schema/) for documentation about the format.
	// TODO: Is there a stronger typed representation of this w/o adding a dependency?
	Parameters json.RawMessage `json:"parameters"`
}

// TODO(Andy): Support streaming.

const (
	functionCallNone = "none"
	functionCallAuto = "auto"
)

// FunctionCall specifies that the model should explicitly call the |Name|d function.
// To specify that the model should not call a function, use FunctionCallNone.
// To specify the default behavior, use FunctionCallAuto.
type FunctionCall struct {
	Name         string `json:"name"`
	returnString *string
}

// MarshalJSON implements json.Marshaler.
func (f *FunctionCall) MarshalJSON() ([]byte, error) {
	if f.returnString != nil {
		return []byte(fmt.Sprintf(`"%s"`, *f.returnString)), nil
	}
	return json.Marshal(struct {
		Name string `json:"name"`
	}{
		Name: f.Name,
	})
}

// FunctionCallNone returns a FunctionCall which will specify that no function should be called.
func FunctionCallNone() *FunctionCall {
	return &FunctionCall{returnString: params.Optional(functionCallNone)}
}

// FunctionCallAuto returns a FunctionCall which will specify that the model should pick between an end-user or
// calling a function.
func FunctionCallAuto() *FunctionCall {
	return &FunctionCall{returnString: params.Optional(functionCallAuto)}
}

// ChatCompletionRequest contains all relevant fields for requests to the chat completions endpoint.
type ChatCompletionRequest struct {
	// Model specifies the ID of the model to use.
	// See more here: https://platform.openai.com/docs/models/overview.
	Model models.ChatCompletion `json:"model"`
	// Messages are the messages to generate chat completions for, in the chat format.
	Messages []*ChatMessage `json:"messages"`
	// Functions are a list of functions the model may generate JSON inputs for.
	Functions []*Function `json:"functions,omitempty"`
	// FunctionCall controls how the model responds to function calls. "none" means the model does not call a function,
	// and responds to the end-user. "auto" means the model can pick between an end-user or calling a function.
	// Specifying a particular function via {"name":\ "my_function"} forces the model to call that function.
	// "none" is the default when no functions are present. "auto" is the default if functions are present.
	// Use FunctionCallByName to generate a FunctionCall value which will explicitly call a function named |name|.
	FunctionCall *FunctionCall `json:"function_call,omitempty"`
	// Temperature specifies what sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the
	// output more random, while lower values like 0.2 will make it more focused and deterministic. OpenAI generally
	// recommends altering this or top_p but not both.
	// Defaults to 1.
	Temperature *float64 `json:"temperature,omitempty"`
	// TopP specifies an alternative to sampling with temperature, called nucleus sampling, where the model considers
	// the results of the tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10%
	// probability mass are considered. OpenAI generally recommends altering this or temperature but not both.
	// Defaults to 1.
	TopP *float64 `json:"top_p,omitempty"`
	// N specifies many chat completion choices to generate for each input message.
	// Defaults to 1.
	N int `json:"n,omitempty"`
	// Stop specifies up to 4 sequences where the API will stop generating further tokens.
	Stop []string `json:"stop,omitempty"`
	// MaxTokens specifies the maximum number of tokens to generate in the chat completion. The total length of input
	// tokens and generated tokens is limited by the model's context length.
	// Defaults to Infinity.
	MaxTokens int `json:"max_tokens,omitempty"`
	// PresencePenalty can be a number between -2.0 and 2.0. Positive values penalize new tokens based on whether they
	// appear in the text so far, increasing the model's likelihood to talk about new topics.
	// Defaults to 0.
	PresencePenalty float32 `json:"presence_penalty,omitempty"`
	// FrequencyPenalty can be a number between -2.0 and 2.0. Positive values penalize new tokens based on their
	// existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.
	// Defaults to 0.
	FrequencyPenalty float32 `json:"frequency_penalty,omitempty"`
	// LogitBias modifies the likelihood of specified tokens appearing in the completion. Accepts a json object that
	// maps tokens (specified by their token ID in the tokenizer) to an associated bias value from -100 to 100.
	// Mathematically, the bias is added to the logits generated by the model prior to sampling. The exact effect will
	// vary per model, but values between -1 and 1 should decrease or increase likelihood of selection; values like -100
	// or 100 should result in a ban or exclusive selection of the relevant token.
	// Defaults to null.
	LogitBias map[string]int `json:"logit_bias,omitempty"`
	// User is a unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
	// See more here: https://beta.openai.com/docs/guides/safety-best-practices/end-user-ids
	User string `json:"user,omitempty"`
}

// CreateChatCompletion creates a chat completion for the provided prompt and parameters.
func (c *Client) CreateChatCompletion(ctx context.Context, cr *ChatCompletionRequest) (*ChatCompletionResponse, error) {
	var b, err = c.post(ctx, routes.ChatCompletions, cr)
	if err != nil {
		return nil, err
	}

	var resp = &ChatCompletionResponse{}
	if err = json.Unmarshal(b, resp); err != nil {
		return nil, err
	}

	return resp, nil
}

// ChatCompletionResponse is the response from the chat completions endpoint.
type ChatCompletionResponse struct {
	ID      string                  `json:"id"`
	Object  objects.Object          `json:"object"`
	Created uint64                  `json:"created"`
	Choices []*ChatCompletionChoice `json:"choices"`
	Usage   *Usage                  `json:"usage"`
}

// ChatCompletionChoice represents one of possible chat completions.
type ChatCompletionChoice struct {
	Index        int          `json:"index"`
	Message      *ChatMessage `json:"message"`
	FinishReason string       `json:"finish_reason"`
}