pkoukk / tiktoken-go

go version of tiktoken
MIT License
649 stars 71 forks source link

fix: field of name also need to be count #13

Closed Doublemine closed 1 year ago

Doublemine commented 1 year ago

Below is the unit test code based on openai-cookbook examples :


func TestNumTokensFromMessages(t *testing.T) {
    type args struct {
        messages []openai.ChatCompletionMessage
    }
    tests := []struct {
        args    args
        wantErr bool
    }{
        {
            args: args{
                messages: []openai.ChatCompletionMessage{
                    {
                        Role:    "system",
                        Content: "You are a helpful, pattern-following assistant that translates corporate jargon into plain English.",
                    },
                    {
                        Role:    "system",
                        Name:    "example_user",
                        Content: "New synergies will help drive top-line growth.",
                    },
                    {
                        Role:    "system",
                        Name:    "example_assistant",
                        Content: "Things working well together will increase revenue.",
                    },
                    {
                        Role:    "system",
                        Name:    "example_user",
                        Content: "Let's circle back when we have more bandwidth to touch base on opportunities for increased leverage.",
                    },
                    {
                        Role:    "system",
                        Name:    "example_assistant",
                        Content: "Let's talk later when we're less busy about how to do better.",
                    },
                    {
                        Role:    "user",
                        Content: "This late pivot means we don't have time to boil the ocean for the client deliverable.",
                    },
                },
            },
            wantErr: false,
        },
    }
    {
        wantTokenNum := 127
        for _, tt := range tests {
            t.Run("test gpt3.5 trubo", func(t *testing.T) {
                got, err := NumTokensFromMessages(tt.args.messages,"gpt-3.5-turbo-0301")
                if (err != nil) != tt.wantErr {
                    t.Errorf("NumTokensFromMessages() error = %v, wantErr %v", err, tt.wantErr)
                    return
                }
                if got != wantTokenNum {
                    t.Errorf("NumTokensFromMessages() = %v, want %v", got, wantTokenNum)
                }
            })
        }
    }
    {

        wantTokenNum := 129
        for _, tt := range tests {
            t.Run("test gpt4", func(t *testing.T) {
                got, err := NumTokensFromMessages(tt.args.messages,"gpt-4-0314")
                if (err != nil) != tt.wantErr {
                    t.Errorf("NumTokensFromMessages() error = %v, wantErr %v", err, tt.wantErr)
                    return
                }
                if got != wantTokenNum {
                    t.Errorf("NumTokensFromMessages() = %v, want %v", got, wantTokenNum)
                }
            })
        }
    }
}