chat.go

package groq

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/conneroisu/groq-go/pkg/builders"
	"github.com/conneroisu/groq-go/pkg/groqerr"
	"github.com/conneroisu/groq-go/pkg/models"
	"github.com/conneroisu/groq-go/pkg/schema"
	"github.com/conneroisu/groq-go/pkg/streams"
	"github.com/conneroisu/groq-go/pkg/tools"
)

const (
	// ChatMessageRoleSystem is the system chat message role.
	ChatMessageRoleSystem Role = "system"
	// ChatMessageRoleUser is the user chat message role.
	ChatMessageRoleUser Role = "user"
	// ChatMessageRoleAssistant is the assistant chat message role.
	ChatMessageRoleAssistant Role = "assistant"
	// ChatMessageRoleFunction is the function chat message role.
	ChatMessageRoleFunction Role = "function"
	// ChatMessageRoleTool is the tool chat message role.
	ChatMessageRoleTool Role = "tool"

	// ImageURLDetailHigh is the high image url detail.
	ImageURLDetailHigh ImageURLDetail = "high"
	// ImageURLDetailLow is the low image url detail.
	ImageURLDetailLow ImageURLDetail = "low"
	// ImageURLDetailAuto is the auto image url detail.
	ImageURLDetailAuto ImageURLDetail = "auto"

	// ChatMessagePartTypeText is the text chat message part type.
	ChatMessagePartTypeText ChatMessagePartType = "text"
	// ChatMessagePartTypeImageURL is the image url chat message part type.
	ChatMessagePartTypeImageURL ChatMessagePartType = "image_url"
)

type (
	// FinishReason is the finish reason.
	//
	// string
	FinishReason string
)

const (
	// ReasonStop is the stop finish reason for a chat completion.
	ReasonStop FinishReason = "stop"
	// ReasonLength is the length finish reason for a chat completion.
	ReasonLength FinishReason = "length"
	// ReasonFunctionCall is the function call finish reason for a chat
	// completion.
	ReasonFunctionCall FinishReason = "function_call"
	// ReasonToolCalls is the tool calls finish reason for a chat
	// completion.
	ReasonToolCalls FinishReason = "tool_calls"
	// ReasonContentFilter is the content filter finish reason for a chat
	// completion.
	ReasonContentFilter FinishReason = "content_filter"
	// ReasonNull is the null finish reason for a chat completion.
	ReasonNull FinishReason = "null"
)

type (
	// ImageURLDetail is the detail of the image at the URL.
	//
	// string
	ImageURLDetail string
	// ChatMessagePartType is the chat message part type.
	//
	// string
	ChatMessagePartType string
	// Role is the role of the chat completion message.
	//
	// string
	Role string
	// PromptAnnotation represents the prompt annotation.
	PromptAnnotation struct {
		PromptIndex int `json:"prompt_index,omitempty"`
	}
	// ChatMessageImageURL represents the chat message image url.
	ChatMessageImageURL struct {
		// URL is the url of the image.
		URL string `json:"url,omitempty"`
		// Detail is the detail of the image url.
		Detail ImageURLDetail `json:"detail,omitempty"`
	}
	// ChatMessagePart represents the chat message part of a chat completion
	// message.
	ChatMessagePart struct {
		// Text is the text of the chat message part.
		Text string `json:"text,omitempty"`
		// Type is the type of the chat message part.
		Type ChatMessagePartType `json:"type,omitempty"`
		// ImageURL is the image url of the chat message part.
		ImageURL *ChatMessageImageURL `json:"image_url,omitempty"`
	}
	// ChatCompletionMessage represents the chat completion message.
	ChatCompletionMessage struct {
		// Name is the name of the chat completion message.
		Name string `json:"name"`
		// Role is the role of the chat completion message.
		Role Role `json:"role"`
		// Content is the content of the chat completion message.
		Content string `json:"content"`
		// MultiContent is the multi content of the chat completion
		// message.
		MultiContent []ChatMessagePart `json:"-"`
		// FunctionCall setting for Role=assistant prompts this may be
		// set to the function call generated by the model.
		FunctionCall *tools.FunctionCall `json:"function_call,omitempty"`
		// ToolCalls setting for Role=assistant prompts this may be set
		// to the tool calls generated by the model, such as function
		// calls.
		ToolCalls []tools.ToolCall `json:"tool_calls,omitempty"`
		// ToolCallID is setting for Role=tool prompts this should be
		// set to the ID given in the assistant's prior request to call
		// a tool.
		ToolCallID string `json:"tool_call_id,omitempty"`
	}
	// ChatCompletionResponseFormat is the chat completion response format.
	ChatCompletionResponseFormat struct {
		// Type is the type of the chat completion response format.
		Type Format `json:"type,omitempty"`
		// JSONSchema is the json schema of the chat completion response
		// format.
		JSONSchema *ChatCompletionResponseFormatJSONSchema `json:"json_schema,omitempty"`
	}
	// ChatCompletionResponseFormatJSONSchema is the chat completion
	// response format json schema.
	ChatCompletionResponseFormatJSONSchema struct {
		// Name is the name of the chat completion response format json
		// schema.
		//
		// it is used to further identify the schema in the response.
		Name string `json:"name"`
		// Description is the description of the chat completion
		// response format json schema.
		Description string `json:"description,omitempty"`
		// Schema is the schema of the chat completion response format
		// json schema.
		Schema schema.Schema `json:"schema"`
		// Strict determines whether to enforce the schema upon the
		// generated content.
		Strict bool `json:"strict"`
	}
	// ChatCompletionRequest represents a request structure for the chat
	// completion API.
	ChatCompletionRequest struct {
		// Model is the model of the chat completion request.
		Model models.ChatModel `json:"model"`
		// Messages is the messages of the chat completion request.
		//
		// These act as the prompt for the model.
		Messages []ChatCompletionMessage `json:"messages"`
		// MaxTokens is the max tokens of the chat completion request.
		MaxTokens int `json:"max_tokens,omitempty"`
		// Temperature is the temperature of the chat completion
		// request.
		Temperature float32 `json:"temperature,omitempty"`
		// TopP is the top p of the chat completion request.
		TopP float32 `json:"top_p,omitempty"`
		// N is the n of the chat completion request.
		N int `json:"n,omitempty"`
		// Stream is the stream of the chat completion request.
		Stream bool `json:"stream,omitempty"`
		// Stop is the stop of the chat completion request.
		Stop []string `json:"stop,omitempty"`
		// PresencePenalty is the presence penalty of the chat
		// completion request.
		PresencePenalty float32 `json:"presence_penalty,omitempty"`
		// ResponseFormat is the response format of the chat completion
		// request.
		ResponseFormat *ChatCompletionResponseFormat `json:"response_format,omitempty"`
		// Seed is the seed of the chat completion request.
		Seed *int `json:"seed,omitempty"`
		// FrequencyPenalty is the frequency penalty of the chat
		// completion request.
		FrequencyPenalty float32 `json:"frequency_penalty,omitempty"`
		// LogitBias is must be a token id string (specified by their
		// token ID in the tokenizer), not a word string. incorrect: `"logit_bias":{ "You": 6}`, correct: `"logit_bias":{"1639": 6}` refs: https://platform.openai.com/docs/api-reference/chat/create#chat/create-logit_bias
		LogitBias map[string]int `json:"logit_bias,omitempty"`
		// LogProbs indicates whether to return log probabilities of the
		// output tokens or not. If true, returns the log probabilities
		// of each output token returned in the content of message.
		//
		// This option is currently not available on the
		// gpt-4-vision-preview model.
		LogProbs bool `json:"logprobs,omitempty"`
		// TopLogProbs is an integer between 0 and 5 specifying the
		// number of most likely tokens to return at each token
		// position, each with an associated log probability. Logprobs
		// must be set to true if this parameter is used.
		TopLogProbs int `json:"top_logprobs,omitempty"`
		// User is the user of the chat completion request.
		User string `json:"user,omitempty"`
		// Tools is the tools of the chat completion request.
		Tools []tools.Tool `json:"tools,omitempty"`
		// This can be either a string or an ToolChoice object.
		ToolChoice any `json:"tool_choice,omitempty"`
		// Options for streaming response. Only set this when you set stream: true.
		StreamOptions *StreamOptions `json:"stream_options,omitempty"`
		// Disable the default behavior of parallel tool calls by setting it: false.
		ParallelToolCalls any `json:"parallel_tool_calls,omitempty"`
		// RetryDelay is the delay between retries.
		RetryDelay time.Duration `json:"-"`
	}
	// LogProbs is the top-level structure containing the log probability information.
	LogProbs struct {
		// Content is a list of message content tokens with log
		// probability information.
		Content []struct {
			// Token is the token of the log prob.
			Token string `json:"token"`
			// LogProb is the log prob of the log prob.
			LogProb float64 `json:"logprob"`
			// Omitting the field if it is null
			Bytes []byte `json:"bytes,omitempty"`
			// TopLogProbs is a list of the most likely tokens and
			// their log probability, at this token position. In
			// rare cases, there may be fewer than the number of
			// requested top_logprobs returned.
			TopLogProbs []TopLogProbs `json:"top_logprobs"`
		} `json:"content"`
	}
	// TopLogProbs represents the top log probs.
	TopLogProbs struct {
		// Token is the token of the top log probs.
		Token string `json:"token"`
		// LogProb is the log prob of the top log probs.
		LogProb float64 `json:"logprob"`
		// Bytes is the bytes of the top log probs.
		Bytes []byte `json:"bytes,omitempty"`
	}
	// ChatCompletionChoice represents the chat completion choice.
	ChatCompletionChoice struct {
		Index int `json:"index"` // Index is the index of the choice.
		// Message is the chat completion message of the choice.
		Message ChatCompletionMessage `json:"message"`
		// FinishReason is the finish reason of the choice.
		FinishReason FinishReason `json:"finish_reason"`
		// LogProbs is the log probs of the choice.
		//
		// This is basically the probability of the model choosing the
		// token.
		LogProbs *LogProbs `json:"logprobs,omitempty"`
	}
	// ChatCompletionResponse represents a response structure for chat
	// completion API.
	ChatCompletionResponse struct {
		// ID is the id of the response.
		ID string `json:"id"`
		// Object is the object of the response.
		Object string `json:"object"`
		// Created is the created time of the response.
		Created int64 `json:"created"`
		// Model is the model of the response.
		Model models.ChatModel `json:"model"`
		// Choices is the choices of the response.
		Choices []ChatCompletionChoice `json:"choices"`
		// Usage is the usage of the response.
		Usage Usage `json:"usage"`
		// SystemFingerprint is the system fingerprint of the response.
		SystemFingerprint string `json:"system_fingerprint"`
		// Header is the header of the response.
		http.Header
	}
	// ChatCompletionStreamChoiceDelta represents a response structure for
	// chat completion API.
	ChatCompletionStreamChoiceDelta struct {
		// Content is the content of the response.
		Content string `json:"content,omitempty"`
		// Role is the role of the creator of the completion.
		Role string `json:"role,omitempty"`
		// FunctionCall is the function call of the response.
		FunctionCall *tools.FunctionCall `json:"function_call,omitempty"`
		// ToolCalls are the tool calls of the response.
		ToolCalls []tools.ToolCall `json:"tool_calls,omitempty"`
	}
	// ChatCompletionStreamChoice represents a response structure for chat
	// completion API.
	ChatCompletionStreamChoice struct {
		// Index is the index of the choice.
		Index int `json:"index"`
		// Delta is the delta of the choice.
		Delta ChatCompletionStreamChoiceDelta `json:"delta"`
		// FinishReason is the finish reason of the choice.
		FinishReason FinishReason `json:"finish_reason"`
	}
	// StreamOptions represents the stream options.
	StreamOptions struct {
		// IncludeUsage is the include usage option of the stream
		// options.
		//
		// If set, an additional chunk will be streamed before the data:
		// [DONE] message.
		// The usage field on this chunk shows the token usage
		// statistics for the entire request, and the choices field will
		// always be an empty array.
		//
		// All other chunks will also include a usage field, but with a
		// null value.
		IncludeUsage bool `json:"include_usage,omitempty"`
	}
	// ChatCompletionStreamResponse represents a response structure for chat
	// completion API.
	ChatCompletionStreamResponse struct {
		// ID is the identifier for the chat completion stream response.
		ID string `json:"id"`
		// Object is the object type of the chat completion stream
		// response.
		Object string `json:"object"`
		// Created is the creation time of the chat completion stream
		// response.
		Created int64 `json:"created"`
		// Model is the model used for the chat completion stream
		// response.
		Model models.ChatModel `json:"model"`
		// Choices is the choices for the chat completion stream
		// response.
		Choices []ChatCompletionStreamChoice `json:"choices"`
		// SystemFingerprint is the system fingerprint for the chat
		// completion stream response.
		SystemFingerprint string `json:"system_fingerprint"`
		// PromptAnnotations is the prompt annotations for the chat
		// completion stream response.
		PromptAnnotations []PromptAnnotation `json:"prompt_annotations,omitempty"`
		// PromptFilterResults is the prompt filter results for the chat
		// completion stream response.
		PromptFilterResults []struct {
			Index int `json:"index"`
		} `json:"prompt_filter_results,omitempty"`
		// Usage is an optional field that will only be present when you
		// set stream_options: {"include_usage": true} in your request.
		//
		// When present, it contains a null value except for the last
		// chunk which contains the token usage statistics for the
		// entire request.
		Usage *Usage `json:"usage,omitempty"`
	}
	// ChatCompletionStream is a stream of ChatCompletionStreamResponse.
	ChatCompletionStream struct {
		*streams.StreamReader[*ChatCompletionStreamResponse]
	}
)

// MarshalJSON method implements the json.Marshaler interface.
func (m ChatCompletionMessage) MarshalJSON() ([]byte, error) {
	if m.Content != "" && m.MultiContent != nil {
		return nil, &groqerr.ErrContentFieldsMisused{}
	}
	if len(m.MultiContent) > 0 {
		msg := struct {
			Name         string              `json:"name,omitempty"`
			Role         Role                `json:"role"`
			Content      string              `json:"-"`
			MultiContent []ChatMessagePart   `json:"content,omitempty"`
			FunctionCall *tools.FunctionCall `json:"function_call,omitempty"`
			ToolCalls    []tools.ToolCall    `json:"tool_calls,omitempty"`
			ToolCallID   string              `json:"tool_call_id,omitempty"`
		}(m)
		return json.Marshal(msg)
	}
	msg := struct {
		Name         string              `json:"name,omitempty"`
		Role         Role                `json:"role"`
		Content      string              `json:"content"`
		MultiContent []ChatMessagePart   `json:"-"`
		FunctionCall *tools.FunctionCall `json:"function_call,omitempty"`
		ToolCalls    []tools.ToolCall    `json:"tool_calls,omitempty"`
		ToolCallID   string              `json:"tool_call_id,omitempty"`
	}(m)
	return json.Marshal(msg)
}

// UnmarshalJSON method implements the json.Unmarshaler interface.
func (m *ChatCompletionMessage) UnmarshalJSON(bs []byte) (err error) {
	msg := struct {
		Name         string `json:"name,omitempty"`
		Role         Role   `json:"role"`
		Content      string `json:"content"`
		MultiContent []ChatMessagePart
		FunctionCall *tools.FunctionCall `json:"function_call,omitempty"`
		ToolCalls    []tools.ToolCall    `json:"tool_calls,omitempty"`
		ToolCallID   string              `json:"tool_call_id,omitempty"`
	}{}
	err = json.Unmarshal(bs, &msg)
	if err == nil {
		*m = ChatCompletionMessage(msg)
		return nil
	}
	multiMsg := struct {
		Name         string `json:"name,omitempty"`
		Role         Role   `json:"role"`
		Content      string
		MultiContent []ChatMessagePart   `json:"content"`
		FunctionCall *tools.FunctionCall `json:"function_call,omitempty"`
		ToolCalls    []tools.ToolCall    `json:"tool_calls,omitempty"`
		ToolCallID   string              `json:"tool_call_id,omitempty"`
	}{}
	err = json.Unmarshal(bs, &multiMsg)
	if err != nil {
		return err
	}
	*m = ChatCompletionMessage(multiMsg)
	return nil
}

// MarshalJSON implements the json.Marshaler interface.
func (r FinishReason) MarshalJSON() ([]byte, error) {
	if r == ReasonNull || r == "" {
		return []byte("null"), nil
	}
	return []byte(
		`"` + string(r) + `"`,
	), nil // best effort to not break future API changes
}

// SetHeader sets the header of the response.
func (r *ChatCompletionResponse) SetHeader(h http.Header) { r.Header = h }

// CreateChatCompletion method is an API call to create a chat completion.
//
// Example:
//
//	func run(
//	        ctx context.Context,
//	) error {
//	        key := os.Getenv("GROQ_KEY")
//	        client, err := groq.NewClient(key)
//	        if err != nil {
//	                return err
//	        }
//	        response, err := client.CreateChatCompletion(
//	                ctx,
//	                groq.ChatCompletionRequest{
//	                        Model: models.ModelLlavaV157B4096Preview,
//	                        Messages: []groq.ChatCompletionMessage{
//	                                {
//	                                        Role: groq.ChatMessageRoleUser,
//	                                        MultiContent: []groq.ChatMessagePart{
//	                                                {
//	                                                        Type: groq.ChatMessagePartTypeText,
//	                                                        Text: "What is the contents of the image?",
//	                                                },
//	                                                {
//	                                                        Type: groq.ChatMessagePartTypeImageURL,
//	                                                        ImageURL: &groq.ChatMessageImageURL{
//	                                                                URL:    "https://cdnimg.webstaurantstore.com/images/products/large/87539/251494.jpg",
//	                                                                Detail: "auto",
//	                                                        },
//	                                                }},
//	                                },
//	                        },
//	                        MaxTokens: 2000,
//	                },
//	        )
//	        if err != nil {
//	                return err
//	        }
//	        fmt.Println(response.Choices[0].Message.Content)
//	        return nil
//	}
func (c *Client) CreateChatCompletion(
	ctx context.Context,
	request ChatCompletionRequest,
) (response ChatCompletionResponse, err error) {
	request.Stream = false
	req, err := builders.NewRequest(
		ctx,
		c.header,
		http.MethodPost,
		c.fullURL(chatCompletionsSuffix, withModel(request.Model)),
		builders.WithBody(request))
	if err != nil {
		return
	}
	err = c.sendRequest(req, &response)
	reqErr, ok := err.(*groqerr.APIError)
	if ok && (reqErr.HTTPStatusCode == http.StatusServiceUnavailable ||
		reqErr.HTTPStatusCode == http.StatusInternalServerError) {
		time.Sleep(request.RetryDelay)
		return c.CreateChatCompletion(ctx, request)
	}
	return
}

// CreateChatCompletionStream method is an API call to create a chat completion
// w/ streaming support.
//
// If set, tokens will be sent as data-only server-sent events as they become
// available, with the stream terminated by a data: [DONE] message.
//
// Example:
//
//	func run(
//	        ctx context.Context,
//	        r io.Reader,
//	        w io.Writer,
//	) error {
//	        key := os.Getenv("GROQ_KEY")
//	        client, err := groq.NewClient(key)
//	        if err != nil {
//	                return err
//	        }
//	        for {
//	                err = input(ctx, client, r, w)
//	                if err != nil {
//	                        return err
//	                }
//	        }
//	}
//	func input(
//	        ctx context.Context,
//	        client *groq.Client,
//	        r io.Reader,
//	        w io.Writer,
//	) error {
//	        fmt.Println("")
//	        fmt.Print("->")
//	        reader := bufio.NewReader(r)
//	        writer := w
//	        var lines []string
//	        select {
//	        case <-ctx.Done():
//	                return ctx.Err()
//	        default:
//	                line, err := reader.ReadString('\n')
//	                if err != nil {
//	                        return err
//	                }
//	                if len(strings.TrimSpace(line)) == 0 {
//	                        break
//	                }
//	                lines = append(lines, line)
//	                break
//	        }
//	        history = append(history, groq.ChatCompletionMessage{
//	                Role:    groq.ChatMessageRoleUser,
//	                Content: strings.Join(lines, "\n"),
//	        })
//	        output, err := client.CreateChatCompletionStream(
//	                ctx,
//	                groq.ChatCompletionRequest{
//	                        Model:     models.ModelGemma29BIt,
//	                        Messages:  history,
//	                        MaxTokens: 2000,
//	                },
//	        )
//	        if err != nil {
//	                return err
//	        }
//	        fmt.Fprintln(writer, "\nai: ")
//	        for {
//	                response, err := output.Recv()
//	                if err != nil {
//	                        return err
//	                }
//	                if response.Choices[0].FinishReason == groq.ReasonStop {
//	                        break
//	                }
//	                fmt.Fprint(writer, response.Choices[0].Delta.Content)
//	        }
//	        return nil
//	}
func (c *Client) CreateChatCompletionStream(
	ctx context.Context,
	request ChatCompletionRequest,
) (stream *ChatCompletionStream, err error) {
	request.Stream = true
	req, err := builders.NewRequest(
		ctx,
		c.header,
		http.MethodPost,
		c.fullURL(
			chatCompletionsSuffix,
			withModel(request.Model)),
		builders.WithBody(request),
	)
	if err != nil {
		return nil, err
	}
	resp, err := sendRequestStream(c, req)
	if err != nil {
		return
	}
	return &ChatCompletionStream{
		StreamReader: resp,
	}, nil
}

// CreateChatCompletionJSON method is an API call to create a chat completion
// w/ object output.
//
// Example:
//
//	// Responses is a response from the models endpoint.
//	type Responses []struct {
//	        Title string `json:"title" jsonschema:"title=Poem Title,description=Title of the poem, minLength=1, maxLength=20"`
//	        Text  string `json:"text" jsonschema:"title=Poem Text,description=Text of the poem, minLength=10, maxLength=200"`
//	}
//
//	func run(
//	        ctx context.Context,
//	) error {
//	        client, err := groq.NewClient(os.Getenv("GROQ_KEY"))
//	        if err != nil {
//	                return err
//	        }
//	        resp := &Responses{}
//	        err = client.CreateChatCompletionJSON(ctx, groq.ChatCompletionRequest{
//	                Model: models.ModelLlama3Groq70B8192ToolUsePreview,
//	                Messages: []groq.ChatCompletionMessage{
//	                        {
//	                                Role:    groq.ChatMessageRoleUser,
//	                                Content: "Create 5 short poems in json format with title and text.",
//	                        },
//	                },
//	                MaxTokens: 2000,
//	        }, resp)
//	        if err != nil {
//	                return err
//	        }
//
//	        jsValue, err := json.MarshalIndent(resp, "", "  ")
//	        if err != nil {
//	                return err
//	        }
//	        fmt.Println(string(jsValue))
//
//	        return nil
//	}
func (c *Client) CreateChatCompletionJSON(
	ctx context.Context,
	request ChatCompletionRequest,
	output any,
) (err error) {
	schema, err := schema.ReflectSchema(reflect.TypeOf(output))
	if err != nil {
		return err
	}
	request.ResponseFormat = &ChatCompletionResponseFormat{}
	request.ResponseFormat.JSONSchema = &ChatCompletionResponseFormatJSONSchema{
		Name:        schema.Title,
		Description: schema.Description,
		Schema:      *schema,
		Strict:      true,
	}
	response, err := c.CreateChatCompletion(ctx, request)
	if err != nil {
		reqErr, ok := err.(*groqerr.APIError)
		if ok && (reqErr.HTTPStatusCode == http.StatusServiceUnavailable ||
			reqErr.HTTPStatusCode == http.StatusInternalServerError) {
			time.Sleep(request.RetryDelay)
			return c.CreateChatCompletionJSON(ctx, request, output)
		}
	}
	content := response.Choices[0].Message.Content
	split := strings.Split(content, "```")
	if len(split) > 1 {
		content = split[1]
	}
	err = json.Unmarshal([]byte(content), &output)
	if err != nil {
		return fmt.Errorf(
			"error unmarshalling response (%s) to output: %v",
			response.Choices[0].Message.Content,
			err,
		)
	}
	return
}