mlow/lmcli
1
0
Fork 0
Code Issues 3 Pull Requests Packages Projects Releases Wiki Activity

Parameterize maxTokens

Browse Source 
Minor formatting/commet changes
This commit is contained in:
Matt Low 2023-11-05 18:19:30 +00:00
parent 9c9b8fa412
commit 168e0cf5d3
3 changed files with 22 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
pkg/cli/cmd.go
View File

@ -7,6 +7,9 @@ import (
"github.com/spf13/cobra" "github.com/spf13/cobra"
) )
// TODO: allow setting with flag
const MAX_TOKENS = 256
var rootCmd = &cobra.Command{ var rootCmd = &cobra.Command{
Use: "lmcli", Use: "lmcli",
Short: "Interact with Large Language Models", Short: "Interact with Large Language Models",
@ -112,7 +115,6 @@ var newCmd = &cobra.Command{
ConversationID: conversation.ID, ConversationID: conversation.ID,
Role: "assistant", Role: "assistant",
} }
reply.RenderTTY(false) reply.RenderTTY(false)
receiver := make(chan string) receiver := make(chan string)
@ -120,7 +122,8 @@ var newCmd = &cobra.Command{
go func() { go func() {
response <- HandleDelayedResponse(receiver) response <- HandleDelayedResponse(receiver)
}() }()
err = CreateChatCompletionStream(messages, receiver)
err = CreateChatCompletionStream(messages, MAX_TOKENS, receiver)
if err != nil { if err != nil {
Fatal("%v\n", err) Fatal("%v\n", err)
} }
@ -160,7 +163,7 @@ var promptCmd = &cobra.Command{
receiver := make(chan string) receiver := make(chan string)
go HandleDelayedResponse(receiver) go HandleDelayedResponse(receiver)
err := CreateChatCompletionStream(messages, receiver) err := CreateChatCompletionStream(messages, MAX_TOKENS, receiver)
if err != nil { if err != nil {
Fatal("%v\n", err) Fatal("%v\n", err)
} }

36
pkg/cli/openai.go
View File

@ -8,7 +8,7 @@ import (
openai "github.com/sashabaranov/go-openai" openai "github.com/sashabaranov/go-openai"
) )
func CreateChatCompletionRequest(messages []Message) *openai.ChatCompletionRequest { func CreateChatCompletionRequest(messages []Message, maxTokens int) openai.ChatCompletionRequest {
chatCompletionMessages := []openai.ChatCompletionMessage{} chatCompletionMessages := []openai.ChatCompletionMessage{}
for _, m := range messages { for _, m := range messages {
chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{ chatCompletionMessages = append(chatCompletionMessages, openai.ChatCompletionMessage{
@ -17,23 +17,19 @@ func CreateChatCompletionRequest(messages []Message) *openai.ChatCompletionReque
}) })
} }
return &openai.ChatCompletionRequest{ return openai.ChatCompletionRequest{
Model: openai.GPT4, Model: openai.GPT3Dot5Turbo,
MaxTokens: 256,
Messages: chatCompletionMessages, Messages: chatCompletionMessages,
Stream: true, MaxTokens: maxTokens,
} }
} }
// CreateChatCompletion accepts a slice of Message and returns the response // CreateChatCompletion submits a Chat Completion API request and returns the
// of the Large Language Model. // response.
func CreateChatCompletion(messages []Message) (string, error) { func CreateChatCompletion(messages []Message, maxTokens int) (string, error) {
client := openai.NewClient(config.OpenAI.APIKey) client := openai.NewClient(config.OpenAI.APIKey)
resp, err := client.CreateChatCompletion( req := CreateChatCompletionRequest(messages, maxTokens)
context.Background(), resp, err := client.CreateChatCompletion(context.Background(), req)
*CreateChatCompletionRequest(messages),
)
if err != nil { if err != nil {
return "", err return "", err
} }
@ -41,18 +37,15 @@ func CreateChatCompletion(messages []Message) (string, error) {
return resp.Choices[0].Message.Content, nil return resp.Choices[0].Message.Content, nil
} }
// CreateChatCompletionStream submits an streaming Chat Completion API request // CreateChatCompletionStream submits a streaming Chat Completion API request
// and sends the received data to the output channel. // and streams the response to the provided output channel.
func CreateChatCompletionStream(messages []Message, output chan string) error { func CreateChatCompletionStream(messages []Message, maxTokens int, output chan string) error {
client := openai.NewClient(config.OpenAI.APIKey) client := openai.NewClient(config.OpenAI.APIKey)
ctx := context.Background() req := CreateChatCompletionRequest(messages, maxTokens)
req := CreateChatCompletionRequest(messages)
req.Stream = true
defer close(output) defer close(output)
stream, err := client.CreateChatCompletionStream(ctx, *req) stream, err := client.CreateChatCompletionStream(context.Background(), req)
if err != nil { if err != nil {
return err return err
} }
@ -68,7 +61,6 @@ func CreateChatCompletionStream(messages []Message, output chan string) error {
if err != nil { if err != nil {
return err return err
} }
output <- response.Choices[0].Delta.Content output <- response.Choices[0].Delta.Content
} }
} }

2
pkg/cli/tty.go
View File

@ -45,7 +45,9 @@ func HandleDelayedResponse(response chan string) string {
firstChunk := true firstChunk := true
for chunk := range response { for chunk := range response {
if firstChunk { if firstChunk {
// notify wait animation that we've received data
waitSignal <- "" waitSignal <- ""
// wait for signal that wait animation has completed
<-waitSignal <-waitSignal
firstChunk = false firstChunk = false
} }