From 8f69e09de11e1f10e69442707ea1f61f2019b76d Mon Sep 17 00:00:00 2001
From: Kelvin <kidult.x@gmail.com>
Date: Wed, 14 Jun 2023 09:08:10 +0800
Subject: [PATCH] new model support

---
 main.go | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/main.go b/main.go
index 649c0c2..34b3ff7 100644
--- a/main.go
+++ b/main.go
@@ -166,7 +166,7 @@ func InitRedis() *redis.Client {
 
 // 计算Messages中的token数量
 func numTokensFromMessages(messages []Message, model string) int {
-	if model == "gpt-3.5-turbo-0301" {
+	if strings.Contains(model, "gpt-3.5") {
 		model = "gpt-3.5-turbo"
 	}
 	tkm, err := tiktoken.EncodingForModel(model)
@@ -185,7 +185,7 @@ func numTokensFromMessages(messages []Message, model string) int {
 
 // 计算String中的token数量
 func numTokensFromString(msg string, model string) int {
-	if model == "gpt-3.5-turbo-0301" {
+	if strings.Contains(model, "gpt-3.5") {
 		model = "gpt-3.5-turbo"
 	}
 	tkm, err := tiktoken.EncodingForModel(model)
@@ -603,7 +603,7 @@ func completions(c *gin.Context) {
 func handleGetModels(c *gin.Context) {
 	// BUGFIX: fix options request, see https://github.com/diemus/azure-openai-proxy/issues/3
 	//models := []string{"gpt-4", "gpt-4-0314", "gpt-4-32k", "gpt-4-32k-0314", "gpt-3.5-turbo", "gpt-3.5-turbo-0301", "text-davinci-003", "text-embedding-ada-002"}
-	models := []string{"gpt-3.5-turbo", "gpt-3.5-turbo-0301", "text-davinci-003", "text-embedding-ada-002", "text-davinci-edit-001"}
+	models := []string{"gpt-3.5-turbo", "gpt-3.5-turbo-0301", "gpt-3.5-turbo-0613", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-16k-0613", "text-davinci-003", "text-embedding-ada-002", "text-davinci-edit-001", "code-davinci-edit-001", "images-generations"}
 	result := ListModelResponse{
 		Object: "list",
 	}
@@ -731,11 +731,26 @@ func main() {
 	// Redis.Set(context.Background(), "server:2", serverInfoStr2, 0)
 
 	// var modelInfo ModelInfo = ModelInfo{
-	// 	ModelName:       "gpt-3.5-turbo",
+	// 	ModelName:       "gpt-3.5-turbo-0613",
 	// 	ModelPrice:      0.0001,
 	// 	ModelPrepayment: 4000,
 	// }
 	// modelInfoStr, _ := json.Marshal(&modelInfo)
+	// Redis.Set(context.Background(), "model:gpt-3.5-turbo-0613", modelInfoStr, 0)
+	// var modelInfo1 ModelInfo = ModelInfo{
+	// 	ModelName:       "gpt-3.5-turbo-16k",
+	// 	ModelPrice:      0.0001,
+	// 	ModelPrepayment: 4000,
+	// }
+	// modelInfoStr1, _ := json.Marshal(&modelInfo1)
+	// Redis.Set(context.Background(), "model:gpt-3.5-turbo-16k", modelInfoStr1, 0)
+	// var modelInfo2 ModelInfo = ModelInfo{
+	// 	ModelName:       "gpt-3.5-turbo-16k-061",
+	// 	ModelPrice:      0.0001,
+	// 	ModelPrepayment: 4000,
+	// }
+	// modelInfoStr2, _ := json.Marshal(&modelInfo2)
+	// Redis.Set(context.Background(), "model:gpt-3.5-turbo-16k-0613", modelInfoStr2, 0)
 	// var modelInfo2 ModelInfo = ModelInfo{
 	// 	ModelName:       "text-davinci-003",
 	// 	ModelPrice:      0.001,