lmcli/pkg/cli/openai.go

package cli

import (
	"context"
	"errors"
	"io"

	openai "github.com/sashabaranov/go-openai"
)

func CreateChatCompletionRequest(system string, messages []Message) *openai.ChatCompletionRequest {
	chatCompletionMessages := []openai.ChatCompletionMessage{
		{
			Role:    "system",
			Content: system,
		},
	}

	for _, m := range messages {
		chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{
			Role:    m.Role,
			Content: m.OriginalContent,
		})
	}

	return &openai.ChatCompletionRequest{
		Model:     openai.GPT4,
		MaxTokens: 256,
		Messages:  chatCompletionMessages,
		Stream:    true,
	}
}

// CreateChatCompletion accepts a slice of Message and returns the response
// of the Large Language Model.
func CreateChatCompletion(system string, messages []Message) (string, error) {
	client := openai.NewClient(config.OpenAI.APIKey)
	resp, err := client.CreateChatCompletion(
		context.Background(),
		*CreateChatCompletionRequest(system, messages),
	)

	if err != nil {
		return "", err
	}

	return resp.Choices[0].Message.Content, nil
}

// CreateChatCompletionStream submits an streaming Chat Completion  API request
// and sends the received data to the output channel.
func CreateChatCompletionStream(system string, messages []Message, output chan string) error {
	client := openai.NewClient(config.OpenAI.APIKey)
	ctx := context.Background()

	req := CreateChatCompletionRequest(system, messages)
	req.Stream = true

	defer close(output)

	stream, err := client.CreateChatCompletionStream(ctx, *req)
	if err != nil {
		return err
	}

	defer stream.Close()

	for {
		response, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			return nil
		}

		if err != nil {
			return err
		}

		output <- response.Choices[0].Delta.Content
	}
}