Skip to content

Commit

Permalink
get around the token limit
Browse files Browse the repository at this point in the history
  • Loading branch information
@quantonganh
quantonganh committed May 9, 2023
1 parent 507b0c1 commit 23c36f8
Show file tree
Hide file tree
Showing 3 changed files with 126 additions and 34 deletions.
3 changes: 3 additions & 0 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,15 @@ require (
github.com/gdamore/tcell/v2 v2.6.0
github.com/kljensen/snowball v0.8.0
github.com/mitchellh/go-homedir v1.1.0
github.com/pkoukk/tiktoken-go v0.1.1
github.com/rivo/tview v0.0.0-20230320095235-84f9c0ff9de8
github.com/tidwall/buntdb v1.2.10
)

require (
github.com/dlclark/regexp2 v1.8.1 // indirect
github.com/gdamore/encoding v1.0.0 // indirect
github.com/google/uuid v1.3.0 // indirect
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
github.com/mattn/go-runewidth v0.0.14 // indirect
github.com/rivo/uniseg v0.4.4 // indirect
Expand Down
10 changes: 10 additions & 0 deletions go.sum
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0=
github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/gdamore/encoding v1.0.0 h1:+7OoQ1Bc6eTm5niUzBa0Ctsh6JbMW6Ra+YNuAtDBdko=
github.com/gdamore/encoding v1.0.0/go.mod h1:alR0ol34c49FCSBLjhosxzcPHQbf2trDkoo5dl+VrEg=
github.com/gdamore/tcell/v2 v2.6.0 h1:OKbluoP9VYmJwZwq/iLb4BxwKcwGthaa1YNBJIyCySg=
github.com/gdamore/tcell/v2 v2.6.0/go.mod h1:be9omFATkdr0D9qewWW3d+MEvl5dha+Etb5y65J2H8Y=
github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/kljensen/snowball v0.8.0 h1:WU4cExxK6sNW33AiGdbn4e8RvloHrhkAssu2mVJ11kg=
github.com/kljensen/snowball v0.8.0/go.mod h1:OGo5gFWjaeXqCu4iIrMl5OYip9XUJHGOU5eSkPjVg2A=
github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
Expand All @@ -10,12 +15,16 @@ github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWV
github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
github.com/pkoukk/tiktoken-go v0.1.1 h1:jtkYlIECjyM9OW1w4rjPmTohK4arORP9V25y6TM6nXo=
github.com/pkoukk/tiktoken-go v0.1.1/go.mod h1:boMWvk9pQCOTx11pgu0DrIdrAKgQzzJKUP6vLXaz7Rw=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/rivo/tview v0.0.0-20230320095235-84f9c0ff9de8 h1:wthS/rREJ6WlALtQ3Ysp4Cty/qiY2LNslV90U71bNg0=
github.com/rivo/tview v0.0.0-20230320095235-84f9c0ff9de8/go.mod h1:nVwGv4MP47T0jvlk7KuTTjjuSmrGO4JF0iaiNt4bufE=
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
github.com/rivo/uniseg v0.4.3/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis=
github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
github.com/tidwall/assert v0.1.0 h1:aWcKyRBUAdLoVebxo95N7+YZVTFF/ASTr7BN4sLP6XI=
github.com/tidwall/btree v1.6.0 h1:LDZfKfQIBHGHWSwckhXI0RPSXzlo+KYdjK7FWSqOzzg=
github.com/tidwall/btree v1.6.0/go.mod h1:twD9XRA5jj9VUQGELzDO4HPQTNJsoWWfYEL+EUQ2cKY=
Expand Down Expand Up @@ -66,3 +75,4 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
147 changes: 113 additions & 34 deletions main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,14 @@ import (
"net/http"
"os"
"path/filepath"
"regexp"
"strconv"
"strings"
"time"

"github.com/gdamore/tcell/v2"
"github.com/mitchellh/go-homedir"
"github.com/pkoukk/tiktoken-go"
"github.com/rivo/tview"
"github.com/tidwall/buntdb"
)
Expand All @@ -25,6 +28,8 @@ const (
roleUser = "user"
roleAssistant = "assistant"

systemMessage = "You are ChatGPT, a large language model trained by OpenAI. Answer as concisely as possible."

prefixSuggestTitle = "suggest me a short title for "

pageMain = "main"
Expand All @@ -33,6 +38,8 @@ const (

buttonCancel = "Cancel"
buttonDelete = "Delete"

maxTokens = 4097
)

type Conversation struct {
Expand Down Expand Up @@ -314,7 +321,6 @@ func main() {
case tcell.KeyESC:
if textView.GetText(false) != "" || !isNewChat {
app.SetFocus(textView)
textView.ScrollToBeginning()
}
case tcell.KeyEnter:
content := textArea.GetText()
Expand All @@ -324,34 +330,14 @@ func main() {
textArea.SetText("", false)
textArea.SetDisabled(true)

textView.ScrollToEnd()
if textView.GetText(false) != "" {
isNewChat = false
fmt.Fprintf(textView, "\n\n")
}
fmt.Fprintln(textView, "[red::]You:[-]")
fmt.Fprintf(textView, "%s\n\n", content)

titleCh := make(chan string)
messages := make([]Message, 0)
if isNewChat {
if textView.GetText(false) == "" {
messages = append(messages, Message{
Role: roleSystem,
Content: "You are ChatGPT, a large language model trained by OpenAI. Answer as concisely as possible.",
Content: systemMessage,
})
} else if list.GetItemCount() > 0 {
title, _ := list.GetItemText(list.GetCurrentItem())
if c, ok := m[title]; ok {
messages = c.Messages
}
}

messages = append(messages, Message{
Role: roleUser,
Content: content,
})

titleCh := make(chan string)
if list.GetItemCount() == 0 || isNewChat {
go func() {
resp, err := createChatCompletion([]Message{
{
Expand All @@ -374,30 +360,87 @@ func main() {
titleCh <- titleResp.Choices[0].Message.Content
}
}()
} else {
isNewChat = false

title, _ := list.GetItemText(list.GetCurrentItem())
if c, ok := m[title]; ok {
messages = c.Messages
}

textView.ScrollToEnd()
fmt.Fprintf(textView, "\n\n")
}

messages = append(messages, Message{
Role: roleUser,
Content: content,
})

numTokens, err := NumTokensFromMessages(messages, gpt3Dot5Turbo)
if err != nil {
log.Println(err)
return nil
}

if numTokens > maxTokens {
isNewChat = true
title, _ := list.GetItemText(list.GetCurrentItem())
go func() {
titleCh <- addSuffixNumber(title)
}()

messages = []Message{
{
Role: roleSystem,
Content: systemMessage,
},
{
Role: roleUser,
Content: fmt.Sprintf("%s: %s", title, content),
},
}

textView.Clear()
}

fmt.Fprintln(textView, "[red::]You:[-]")
fmt.Fprintf(textView, "%s\n\n", content)

respCh := make(chan string)
errCh := make(chan error, 1)
go func() {
resp, err := createChatCompletion(messages, true)
if err != nil {
log.Panic(err)
errCh <- err
}

reader := bufio.NewReader(resp.Body)
for {
line, err := reader.ReadBytes('\n')
if err == nil {
var streamingResp *StreamingResponse
if err := json.Unmarshal(bytes.TrimPrefix(line, []byte("data: ")), &streamingResp); err == nil {
respCh <- streamingResp.Choices[0].Delta.Content
if err != nil {
if errors.Is(err, io.EOF) {
close(respCh)
return
} else {
errCh <- err
}
} else if errors.Is(err, io.EOF) {
close(respCh)
return
}

var streamingResp *StreamingResponse
if err := json.Unmarshal(bytes.TrimPrefix(line, []byte("data: ")), &streamingResp); err == nil {
respCh <- streamingResp.Choices[0].Delta.Content
}
}
}()

select {
case err := <-errCh:
log.Println("received error:", err)
return nil
default:
}

fmt.Fprintln(textView, "[green::]ChatGPT:[-]")
go func() {
var fullContent strings.Builder
Expand Down Expand Up @@ -502,11 +545,47 @@ func main() {
}
}

const completionsURL = "https://api.openai.com/v1/chat/completions"
func NumTokensFromMessages(messages []Message, model string) (int, error) {
t, err := tiktoken.EncodingForModel(model)
if err != nil {
return 0, err
}

var tokensPerMessage int
if model == gpt3Dot5Turbo {
tokensPerMessage = 4
} else {
tokensPerMessage = 3
}

numTokens := 0
for _, message := range messages {
numTokens += tokensPerMessage
numTokens += len(t.Encode(message.Content, nil, nil))
numTokens += len(t.Encode(message.Role, nil, nil))
}
numTokens += 3
return numTokens, nil
}

func addSuffixNumber(title string) string {
re := regexp.MustCompile(`(.*)\s-\s(\d+)$`)
match := re.FindStringSubmatch(title)
if match == nil {
return fmt.Sprintf("%s - %d", title, 2)
}
suffixNumber, _ := strconv.Atoi(match[2])
return fmt.Sprintf("%s - %d", match[1], suffixNumber+1)
}

const (
completionsURL = "https://api.openai.com/v1/chat/completions"
gpt3Dot5Turbo = "gpt-3.5-turbo"
)

func createChatCompletion(messages []Message, stream bool) (*http.Response, error) {
reqBody, err := json.Marshal(&Request{
Model: "gpt-3.5-turbo",
Model: gpt3Dot5Turbo,
Messages: messages,
Stream: stream,
})
Expand Down

0 comments on commit 23c36f8

Please sign in to comment.