fix: add check to model name (#83)

Author: j178

README.md

@@ -290,7 +290,7 @@ Notes:
 
 - `api_type` should be "AZURE" or "AZURE_AD".
 - `api_version` defaults to "2023-05-15" if not specified.
-- Configure `model_mapping` to map model names to your deployment names. If not specified, the model name will be used as the deployment name with `.` or `:` removed (e.g. "gpt-3.5-turbo" -> "gpt-35-turbo").
+- Configure `model_mapping` to map model names to your deployment names. The key must be a valid OpenAI model name. If not specified, the model name will be used as the deployment name with `.` or `:` removed (e.g. "gpt-3.5-turbo" -> "gpt-35-turbo").
 
 Find more details about Azure OpenAI service here: https://learn.microsoft.com/en-US/azure/ai-services/openai/reference.
 

cmd/chatgpt/main.go

@@ -93,7 +93,10 @@ func main() {
 		defer func() { _ = lockFile.Unlock() }()
 	}
 
-	conversations := chatgpt.NewConversationManager(conf, chatgpt.ConversationHistoryFile())
+	conversations, err := chatgpt.NewConversationManager(conf, chatgpt.ConversationHistoryFile())
+	if err != nil {
+		exit(err)
+	}
 
 	if *startNewConversation {
 		conversations.New(conf.Conversation)

config.go

@@ -10,6 +10,8 @@ import (
 
 	"github.com/mitchellh/go-homedir"
 	"github.com/sashabaranov/go-openai"
+
+	"github.com/j178/chatgpt/tokenizer"
 )
 
 type ConversationConfig struct {
@@ -140,7 +142,7 @@ func InitConfig() (GlobalConfig, error) {
 			Prompt:        "default",
 			ContextLength: 6,
 			Stream:        true,
-			Temperature:   0,
+			Temperature:   1.0,
 			MaxTokens:     1024,
 		},
 		KeyMap: defaultKeyMapConfig(),
@@ -157,14 +159,21 @@ func InitConfig() (GlobalConfig, error) {
 	if endpoint != "" {
 		conf.Endpoint = endpoint
 	}
+
 	if conf.APIKey == "" {
 		return GlobalConfig{}, errors.New("Missing API key. Set it in `~/.config/chatgpt/config.json` or by setting the `OPENAI_API_KEY` environment variable. You can find or create your API key at https://platform.openai.com/account/api-keys.")
 	}
+
 	conf.APIType = openai.APIType(strings.ToUpper(string(conf.APIType)))
 	switch conf.APIType {
+	case openai.APITypeOpenAI, openai.APITypeAzure, openai.APITypeAzureAD:
 	default:
 		return GlobalConfig{}, fmt.Errorf("unknown API type: %s", conf.APIType)
-	case openai.APITypeOpenAI, openai.APITypeAzure, openai.APITypeAzureAD:
+	}
+
+	err = tokenizer.CheckModel(conf.Conversation.Model)
+	if err != nil {
+		return GlobalConfig{}, fmt.Errorf("invalid model %s", conf.Conversation.Model)
 	}
 	return conf, nil
 }

conversation.go

@@ -3,7 +3,7 @@ package chatgpt
 import (
 	"encoding/json"
 	"errors"
-	"log"
+	"fmt"
 	"os"
 
 	"github.com/sashabaranov/go-openai"
@@ -18,7 +18,7 @@ type ConversationManager struct {
 	Idx           int             `json:"last_idx"`
 }
 
-func NewConversationManager(conf GlobalConfig, historyFile string) *ConversationManager {
+func NewConversationManager(conf GlobalConfig, historyFile string) (*ConversationManager, error) {
 	h := &ConversationManager{
 		file:       historyFile,
 		globalConf: conf,
@@ -27,9 +27,9 @@ func NewConversationManager(conf GlobalConfig, historyFile string) *Conversation
 
 	err := h.Load()
 	if err != nil {
-		log.Println("Failed to load history:", err)
+		return nil, fmt.Errorf("Failed to load conversation history: %w", err)
 	}
-	return h
+	return h, nil
 }
 
 func (m *ConversationManager) Dump() error {
@@ -65,7 +65,11 @@ func (m *ConversationManager) Load() error {
 	if err != nil {
 		return err
 	}
-	for _, c := range m.Conversations {
+	for i, c := range m.Conversations {
+		err = tokenizer.CheckModel(c.Config.Model)
+		if err != nil {
+			return fmt.Errorf("invalid model %s in conversation %d", c.Config.Model, i+1)
+		}
 		c.manager = m
 	}
 	return nil

tokenizer/tokenize.go

@@ -1,32 +1,52 @@
 package tokenizer
 
 import (
+	"strings"
+
 	"github.com/pkoukk/tiktoken-go"
 	"github.com/sashabaranov/go-openai"
 )
 
 var encoders = map[string]*tiktoken.Tiktoken{}
 
-func CountTokens(model, text string) int {
+func getEncoding(model string) (*tiktoken.Tiktoken, error) {
 	enc, ok := encoders[model]
+	var err error
 	if !ok {
-		enc, _ = tiktoken.EncodingForModel(model)
+		enc, err = tiktoken.EncodingForModel(model)
+		if err != nil {
+			return nil, err
+		}
 		encoders[model] = enc
 	}
+	return enc, nil
+}
+
+func CheckModel(model string) error {
+	_, err := getEncoding(model)
+	return err
+}
+
+func CountTokens(model, text string) int {
+	enc, err := getEncoding(model)
+	if err != nil {
+		panic(err)
+	}
 	return len(enc.Encode(text, nil, nil))
 }
 
 // CountMessagesTokens based on https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
 func CountMessagesTokens(model string, messages []openai.ChatCompletionMessage) int {
-	var tokens int
-	var tokensPerMessage int
-	var tokensPerName int
+	var (
+		tokens           int
+		tokensPerMessage int
+		tokensPerName    int
+	)
 
-	switch model {
-	case openai.GPT3Dot5Turbo, openai.GPT3Dot5Turbo0301:
+	if strings.HasPrefix(model, "gpt-3.5") {
 		tokensPerMessage = 4 // every message follows <|start|>{role/name}\n{content}<|end|>\n
 		tokensPerName = -1   // if there's a name, the role is omitted
-	case openai.GPT4, openai.GPT40314, openai.GPT432K, openai.GPT432K0314:
+	} else if strings.HasPrefix(model, "gpt-4") {
 		tokensPerMessage = 3
 		tokensPerName = 1
 	}