feat: aiproxy probe (#5229)

* feat: aiproxy probe * fix: aiproxy use startup probe
labring · Nov 22, 2024 · f9dc424 · f9dc424
1 parent 26bb340
commit f9dc424
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 8 deletions.
diff --git a/service/aiproxy/controller/misc.go b/service/aiproxy/controller/misc.go
@@ -8,12 +8,16 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
+type StatusData struct {
+	StartTime int64 `json:"startTime"`
+}
+
 func GetStatus(c *gin.Context) {
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
-		"data": gin.H{
-			"start_time": common.StartTime,
+		"data": &StatusData{
+			StartTime: common.StartTime,
 		},
 	})
 }
diff --git a/service/aiproxy/controller/model.go b/service/aiproxy/controller/model.go
@@ -306,6 +306,7 @@ func RetrieveModel(c *gin.Context) {
 				Code:    "model_not_found",
 			},
 		})
+		return
 	}
 	c.JSON(200, model)
 }
diff --git a/service/aiproxy/deploy/manifests/deploy.yaml.tmpl b/service/aiproxy/deploy/manifests/deploy.yaml.tmpl
@@ -56,5 +56,14 @@ spec:
           ports:
             - containerPort: 3000
           imagePullPolicy: Always
+          startupProbe:
+            httpGet:
+              port: 3000
+              path: /api/status
+            initialDelaySeconds: 5
+            periodSeconds: 3
+            failureThreshold: 30
+            successThreshold: 1
+            timeoutSeconds: 1
       serviceAccountName: default
       automountServiceAccountToken: false
diff --git a/service/aiproxy/relay/adaptor/ali/main.go b/service/aiproxy/relay/adaptor/ali/main.go
@@ -20,10 +20,12 @@ func ConvertRequest(request *model.GeneralOpenAIRequest) *model.GeneralOpenAIReq
 	if request.TopP != nil && *request.TopP >= 1 {
 		*request.TopP = 0.9999
 	}
-	if request.StreamOptions == nil {
-		request.StreamOptions = &model.StreamOptions{}
+	if request.Stream {
+		if request.StreamOptions == nil {
+			request.StreamOptions = &model.StreamOptions{}
+		}
+		request.StreamOptions.IncludeUsage = true
 	}
-	request.StreamOptions.IncludeUsage = true
 	return request
 }
 

diff --git a/service/aiproxy/router/api.go b/service/aiproxy/router/api.go
@@ -10,13 +10,17 @@ import (
 )
 
 func SetAPIRouter(router *gin.Engine) {
-	apiRouter := router.Group("/api")
+	api := router.Group("/api")
 	if env.Bool("GZIP_ENABLED", false) {
-		apiRouter.Use(gzip.Gzip(gzip.DefaultCompression))
+		api.Use(gzip.Gzip(gzip.DefaultCompression))
 	}
+
+	healthRouter := api.Group("")
+	healthRouter.GET("/status", controller.GetStatus)
+
+	apiRouter := api.Group("")
 	apiRouter.Use(middleware.AdminAuth)
 	{
-		apiRouter.GET("/status", controller.GetStatus)
 		apiRouter.GET("/models", controller.BuiltinModels)
 		apiRouter.GET("/models/price", controller.ModelPrice)
 		apiRouter.GET("/models/enabled", controller.EnabledModels)