Matt Low
2c64ab501b
Removed the system parameter on ChatCopmletion functions, and persist it in conversations as well.
75 lines
1.6 KiB
Go
75 lines
1.6 KiB
Go
package cli
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
|
|
openai "github.com/sashabaranov/go-openai"
|
|
)
|
|
|
|
func CreateChatCompletionRequest(messages []Message) *openai.ChatCompletionRequest {
|
|
chatCompletionMessages := []openai.ChatCompletionMessage{}
|
|
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(messages []Message) (string, error) {
|
|
client := openai.NewClient(config.OpenAI.APIKey)
|
|
resp, err := client.CreateChatCompletion(
|
|
context.Background(),
|
|
*CreateChatCompletionRequest(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(messages []Message, output chan string) error {
|
|
client := openai.NewClient(config.OpenAI.APIKey)
|
|
ctx := context.Background()
|
|
|
|
req := CreateChatCompletionRequest(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
|
|
}
|
|
}
|