Closed Doublemine closed 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) } }) } } }
Below is the unit test code based on openai-cookbook examples :