lmcli/pkg/cli/openai.go

72 lines
1.6 KiB
Go
Raw Normal View History

package cli
2023-10-30 15:23:07 -06:00
import (
"context"
2023-10-30 15:45:21 -06:00
"errors"
2023-10-30 15:23:07 -06:00
"fmt"
2023-10-30 15:45:21 -06:00
"io"
2023-10-30 15:23:07 -06:00
openai "github.com/sashabaranov/go-openai"
)
func CreateChatCompletionRequest(messages []Message) (openai.ChatCompletionRequest) {
var chatCompletionMessages []openai.ChatCompletionMessage
2023-10-30 15:23:07 -06:00
for _, m := range(messages) {
chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{
2023-10-30 15:23:07 -06:00
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)
2023-10-30 15:23:07 -06:00
resp, err := client.CreateChatCompletion(
context.Background(),
CreateChatCompletionRequest(messages),
2023-10-30 15:23:07 -06:00
)
if err != nil {
return "", fmt.Errorf("ChatCompletion error: %v\n", err)
}
return resp.Choices[0].Message.Content, nil
}
2023-10-30 15:45:21 -06:00
func CreateChatCompletionStream(system string, messages []Message, output io.Writer) (error) {
client := openai.NewClient(config.OpenAI.APIKey)
2023-10-30 15:45:21 -06:00
ctx := context.Background()
req := CreateChatCompletionRequest(messages)
req.Stream = true
2023-10-30 15:45:21 -06:00
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)
}
}