From 86fa6da6b449e374ace5b4a0e5b36fd6fe41e622 Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Tue, 31 Oct 2023 23:35:12 +0800
Subject: [PATCH] add gpt-4-vision model

---
 adapter/chatgpt/struct.go |  1 +
 auth/rule.go              |  2 +-
 globals/variables.go      |  3 +++
 utils/tokenizer.go        | 19 +++++++++----------
 4 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/adapter/chatgpt/struct.go b/adapter/chatgpt/struct.go
index 64e06b5..0d6f17d 100644
--- a/adapter/chatgpt/struct.go
+++ b/adapter/chatgpt/struct.go
@@ -49,6 +49,7 @@ func NewChatInstanceFromConfig(v string) *ChatInstance {
 func NewChatInstanceFromModel(props *InstanceProps) *ChatInstance {
 	switch props.Model {
 	case globals.GPT4,
+		globals.GPT4Vision,
 		globals.GPT40314,
 		globals.GPT40613:
 		return NewChatInstanceFromConfig("gpt4")
diff --git a/auth/rule.go b/auth/rule.go
index c4521f4..236506f 100644
--- a/auth/rule.go
+++ b/auth/rule.go
@@ -12,7 +12,7 @@ func CanEnableModel(db *sql.DB, user *User, model string) bool {
 	case globals.GPT3Turbo, globals.GPT3Turbo0301, globals.GPT3Turbo0613,
 		globals.Claude2:
 		return true
-	case globals.GPT4, globals.GPT40613, globals.GPT40314:
+	case globals.GPT4, globals.GPT4Vision, globals.GPT40613, globals.GPT40314:
 		return user != nil && user.GetQuota(db) >= 5
 	case globals.GPT432k, globals.GPT432k0613, globals.GPT432k0314:
 		return user != nil && user.GetQuota(db) >= 50
diff --git a/globals/variables.go b/globals/variables.go
index 076a3e4..c27eda3 100644
--- a/globals/variables.go
+++ b/globals/variables.go
@@ -45,6 +45,7 @@ const (
 	GPT3Turbo16k0613 = "gpt-3.5-turbo-16k-0613"
 	GPT3Turbo16k0301 = "gpt-3.5-turbo-16k-0301"
 	GPT4             = "gpt-4"
+	GPT4Vision       = "gpt-4v"
 	GPT40314         = "gpt-4-0314"
 	GPT40613         = "gpt-4-0613"
 	GPT432k          = "gpt-4-32k"
@@ -80,6 +81,7 @@ var GPT3Turbo16kArray = []string{
 
 var GPT4Array = []string{
 	GPT4,
+	GPT4Vision,
 	GPT40314,
 	GPT40613,
 }
@@ -149,6 +151,7 @@ var AllModels = []string{
 	GPT4,
 	GPT40314,
 	GPT40613,
+	GPT4Vision,
 	GPT432k,
 	GPT432k0314,
 	GPT432k0613,
diff --git a/utils/tokenizer.go b/utils/tokenizer.go
index fd157e3..c6dc932 100644
--- a/utils/tokenizer.go
+++ b/utils/tokenizer.go
@@ -48,6 +48,7 @@ func GetWeightByModel(model string) int {
 		globals.GPT3Turbo16k0613,
 
 		globals.GPT4,
+		globals.GPT4Vision,
 		globals.GPT40314,
 		globals.GPT40613,
 		globals.SparkDesk,
@@ -102,13 +103,12 @@ func CountTokenPrice(messages []globals.Message, model string) int {
 
 func CountInputToken(model string, v []globals.Message) float32 {
 	switch model {
-	case globals.GPT3Turbo:
+	case globals.GPT3Turbo, globals.GPT3Turbo0613, globals.GPT3Turbo0301,
+		globals.GPT3Turbo16k, globals.GPT3Turbo16k0613, globals.GPT3Turbo16k0301:
 		return 0
-	case globals.GPT3Turbo16k:
-		return 0
-	case globals.GPT4:
+	case globals.GPT4, globals.GPT4Vision, globals.GPT40314, globals.GPT40613:
 		return float32(CountTokenPrice(v, model)) / 1000 * 2.1
-	case globals.GPT432k:
+	case globals.GPT432k, globals.GPT432k0613, globals.GPT432k0314:
 		return float32(CountTokenPrice(v, model)) / 1000 * 4.2
 	case globals.SparkDesk:
 		return float32(CountTokenPrice(v, model)) / 1000 * 0.15
@@ -129,13 +129,12 @@ func CountInputToken(model string, v []globals.Message) float32 {
 
 func CountOutputToken(model string, t int) float32 {
 	switch model {
-	case globals.GPT3Turbo:
+	case globals.GPT3Turbo, globals.GPT3Turbo0613, globals.GPT3Turbo0301,
+		globals.GPT3Turbo16k, globals.GPT3Turbo16k0613, globals.GPT3Turbo16k0301:
 		return 0
-	case globals.GPT3Turbo16k:
-		return 0
-	case globals.GPT4:
+	case globals.GPT4, globals.GPT4Vision, globals.GPT40314, globals.GPT40613:
 		return float32(t*GetWeightByModel(model)) / 1000 * 4.3
-	case globals.GPT432k:
+	case globals.GPT432k, globals.GPT432k0613, globals.GPT432k0314:
 		return float32(t*GetWeightByModel(model)) / 1000 * 8.6
 	case globals.SparkDesk:
 		return float32(t*GetWeightByModel(model)) / 1000 * 0.15