lmcli/pkg/cli/openai.go

78 lines
1.6 KiB
Go

package cli
import (
"context"
"errors"
"fmt"
"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
}
func CreateChatCompletionStream(system string, messages []Message, output io.Writer) error {
client := openai.NewClient(config.OpenAI.APIKey)
ctx := context.Background()
req := CreateChatCompletionRequest(system, messages)
req.Stream = true
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 {
//fmt.Printf("\nStream error: %v\n", err)
return err
}
fmt.Fprint(output, response.Choices[0].Delta.Content)
}
}