From 3af4649b526b115cc4ae94905736d1134309dc9d Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Sat, 6 Apr 2024 20:42:35 +0800
Subject: [PATCH] fix: only check model when request path in whitelist

---
 middleware/auth.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/middleware/auth.go b/middleware/auth.go
index 01b2cce3..64ce6608 100644
--- a/middleware/auth.go
+++ b/middleware/auth.go
@@ -116,7 +116,7 @@ func TokenAuth() func(c *gin.Context) {
 			return
 		}
 		requestModel, err := getRequestModel(c)
-		if err != nil && !strings.HasPrefix(c.Request.URL.Path, "/v1/models") {
+		if err != nil && shouldCheckModel(c) {
 			abortWithMessage(c, http.StatusBadRequest, err.Error())
 			return
 		}
@@ -142,3 +142,19 @@ func TokenAuth() func(c *gin.Context) {
 		c.Next()
 	}
 }
+
+func shouldCheckModel(c *gin.Context) bool {
+	if strings.HasPrefix(c.Request.URL.Path, "/v1/completions") {
+		return true
+	}
+	if strings.HasPrefix(c.Request.URL.Path, "/v1/chat/completions") {
+		return true
+	}
+	if strings.HasPrefix(c.Request.URL.Path, "/v1/images") {
+		return true
+	}
+	if strings.HasPrefix(c.Request.URL.Path, "/v1/audio") {
+		return true
+	}
+	return false
+}