add gpt-4-vision model

coaidev · Oct 31, 2023 · 86fa6da · 86fa6da
1 parent 4f817f6
commit 86fa6da
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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:
-		return 0
-	case globals.GPT3Turbo16k:
+	case globals.GPT3Turbo, globals.GPT3Turbo0613, globals.GPT3Turbo0301,
+		globals.GPT3Turbo16k, globals.GPT3Turbo16k0613, globals.GPT3Turbo16k0301:
 		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