From a3b3e6cc386045faacb425f99040be83f59c4aa3 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Wed, 15 May 2024 16:32:00 +0800
Subject: [PATCH] chore: update InitTokenEncoders (#255)

---
 service/token_counter.go | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/service/token_counter.go b/service/token_counter.go
index 59138d8..697ee15 100644
--- a/service/token_counter.go
+++ b/service/token_counter.go
@@ -26,6 +26,7 @@ func InitTokenEncoders() {
 	}
 	defaultTokenEncoder = gpt35TokenEncoder
 	gpt4TokenEncoder, err := tiktoken.EncodingForModel("gpt-4")
+	gpt4oTokenEncoder, err := tiktoken.EncodingForModel("gpt-4o")
 	if err != nil {
 		common.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
 	}
@@ -33,7 +34,11 @@ func InitTokenEncoders() {
 		if strings.HasPrefix(model, "gpt-3.5") {
 			tokenEncoderMap[model] = gpt35TokenEncoder
 		} else if strings.HasPrefix(model, "gpt-4") {
-			tokenEncoderMap[model] = gpt4TokenEncoder
+			if strings.HasPrefix(model, "gpt-4o") {
+				tokenEncoderMap[model] = gpt4oTokenEncoder
+			} else {
+				tokenEncoderMap[model] = gpt4TokenEncoder
+			}
 		} else {
 			tokenEncoderMap[model] = nil
 		}