2023-11-04 12:20:13 -06:00
|
|
|
package cli
|
2023-10-30 15:23:07 -06:00
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
2023-10-30 15:45:21 -06:00
|
|
|
"errors"
|
|
|
|
"io"
|
2023-11-23 20:45:43 -07:00
|
|
|
"strings"
|
2023-11-04 16:56:22 -06:00
|
|
|
|
2023-10-30 15:23:07 -06:00
|
|
|
openai "github.com/sashabaranov/go-openai"
|
|
|
|
)
|
|
|
|
|
2023-11-08 23:07:52 -07:00
|
|
|
func CreateChatCompletionRequest(model string, messages []Message, maxTokens int) openai.ChatCompletionRequest {
|
2023-11-05 01:54:12 -06:00
|
|
|
chatCompletionMessages := []openai.ChatCompletionMessage{}
|
2023-11-04 16:56:22 -06:00
|
|
|
for _, m := range messages {
|
2023-10-30 16:23:27 -06:00
|
|
|
chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{
|
2023-11-04 16:56:22 -06:00
|
|
|
Role: m.Role,
|
2023-10-30 15:23:07 -06:00
|
|
|
Content: m.OriginalContent,
|
|
|
|
})
|
|
|
|
}
|
|
|
|
|
2023-11-05 11:19:30 -07:00
|
|
|
return openai.ChatCompletionRequest{
|
2023-11-08 23:07:52 -07:00
|
|
|
Model: model,
|
2023-11-04 16:56:22 -06:00
|
|
|
Messages: chatCompletionMessages,
|
2023-11-05 11:19:30 -07:00
|
|
|
MaxTokens: maxTokens,
|
2023-11-25 10:17:33 -07:00
|
|
|
N: 1, // limit responses to 1 "choice". we use choices[0] to reference it
|
2023-10-30 16:23:27 -06:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-11-05 11:19:30 -07:00
|
|
|
// CreateChatCompletion submits a Chat Completion API request and returns the
|
|
|
|
// response.
|
2023-11-08 23:07:52 -07:00
|
|
|
func CreateChatCompletion(model string, messages []Message, maxTokens int) (string, error) {
|
2023-11-18 18:14:00 -07:00
|
|
|
client := openai.NewClient(*config.OpenAI.APIKey)
|
2023-11-08 23:07:52 -07:00
|
|
|
req := CreateChatCompletionRequest(model, messages, maxTokens)
|
2023-11-05 11:19:30 -07:00
|
|
|
resp, err := client.CreateChatCompletion(context.Background(), req)
|
2023-10-30 15:23:07 -06:00
|
|
|
if err != nil {
|
2023-11-04 16:07:06 -06:00
|
|
|
return "", err
|
2023-10-30 15:23:07 -06:00
|
|
|
}
|
|
|
|
|
|
|
|
return resp.Choices[0].Message.Content, nil
|
|
|
|
}
|
2023-10-30 15:45:21 -06:00
|
|
|
|
2023-11-05 11:19:30 -07:00
|
|
|
// CreateChatCompletionStream submits a streaming Chat Completion API request
|
2023-11-23 20:45:43 -07:00
|
|
|
// and both returns and streams the response to the provided output channel.
|
|
|
|
// May return a partial response if an error occurs mid-stream.
|
|
|
|
func CreateChatCompletionStream(model string, messages []Message, maxTokens int, output chan<- string) (string, error) {
|
2023-11-18 18:14:00 -07:00
|
|
|
client := openai.NewClient(*config.OpenAI.APIKey)
|
2023-11-08 23:07:52 -07:00
|
|
|
req := CreateChatCompletionRequest(model, messages, maxTokens)
|
2023-10-30 15:45:21 -06:00
|
|
|
|
2023-11-05 11:19:30 -07:00
|
|
|
stream, err := client.CreateChatCompletionStream(context.Background(), req)
|
2023-10-30 15:45:21 -06:00
|
|
|
if err != nil {
|
2023-11-23 20:45:43 -07:00
|
|
|
return "", err
|
2023-10-30 15:45:21 -06:00
|
|
|
}
|
|
|
|
defer stream.Close()
|
|
|
|
|
2023-11-23 20:45:43 -07:00
|
|
|
sb := strings.Builder{}
|
2023-10-30 15:45:21 -06:00
|
|
|
for {
|
2023-11-23 20:45:43 -07:00
|
|
|
response, e := stream.Recv()
|
|
|
|
if errors.Is(e, io.EOF) {
|
|
|
|
break
|
2023-10-30 15:45:21 -06:00
|
|
|
}
|
|
|
|
|
2023-11-23 20:45:43 -07:00
|
|
|
if e != nil {
|
|
|
|
err = e
|
|
|
|
break
|
2023-10-30 15:45:21 -06:00
|
|
|
}
|
2023-11-23 20:45:43 -07:00
|
|
|
chunk := response.Choices[0].Delta.Content
|
|
|
|
output <- chunk
|
|
|
|
sb.WriteString(chunk)
|
2023-10-30 15:45:21 -06:00
|
|
|
}
|
2023-11-23 20:45:43 -07:00
|
|
|
return sb.String(), err
|
2023-10-30 15:45:21 -06:00
|
|
|
}
|