78 lines
1.6 KiB
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)
|
|
}
|
|
}
|