add health / refresh dependency (#34)

* add health / refresh dependency * health check --------- Co-authored-by: hobo.l <[email protected]>
langgenius · Aug 19, 2024 · f1e5df2 · f1e5df2
1 parent 072abe7
commit f1e5df2
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 19 deletions.
diff --git a/internal/controller/router.go b/internal/controller/router.go
@@ -4,24 +4,43 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/langgenius/dify-sandbox/internal/middleware"
 	"github.com/langgenius/dify-sandbox/internal/static"
+	"net/http"
 )
 
-func Setup(eng *gin.Engine) {
-	eng.Use(middleware.Auth())
+func Setup(Router *gin.Engine) {
+	PublicGroup := Router.Group("")
+	PrivateGroup := Router.Group("/v1/sandbox/")
 
-	eng.POST(
-		"/v1/sandbox/run",
-		middleware.MaxRequest(static.GetDifySandboxGlobalConfigurations().MaxRequests),
-		middleware.MaxWorker(static.GetDifySandboxGlobalConfigurations().MaxWorkers),
-		RunSandboxController,
-	)
-	eng.GET(
-		"/v1/sandbox/dependencies",
-		GetDependencies,
-	)
+	PrivateGroup.Use(middleware.Auth())
 
-	eng.POST(
-		"/v1/sandbox/dependencies/update",
-		UpdateDependencies,
-	)
+	{
+		// health check
+		PublicGroup.GET("/health", func(c *gin.Context) {
+			c.JSON(http.StatusOK, "ok")
+		})
+	}
+
+	InitRunRouter(PrivateGroup)
+	InitDependencyRouter(PrivateGroup)
+}
+
+func InitDependencyRouter(Router *gin.RouterGroup) {
+	dependencyRouter := Router.Group("dependencies")
+	{
+		dependencyRouter.GET("", GetDependencies)
+		dependencyRouter.POST("update", UpdateDependencies)
+		dependencyRouter.GET("refresh", RefreshDependencies)
+	}
+}
+
+func InitRunRouter(Router *gin.RouterGroup) {
+	runRouter := Router.Group("")
+	{
+		runRouter.POST(
+			"run",
+			middleware.MaxRequest(static.GetDifySandboxGlobalConfigurations().MaxRequests),
+			middleware.MaxWorker(static.GetDifySandboxGlobalConfigurations().MaxWorkers),
+			RunSandboxController,
+		)
+	}
 }
diff --git a/internal/controller/run.go b/internal/controller/run.go
@@ -54,3 +54,16 @@ func UpdateDependencies(c *gin.Context) {
 		}
 	})
 }
+
+func RefreshDependencies(c *gin.Context) {
+	BindRequest(c, func(req struct {
+		Language string `json:"language" form:"language" binding:"required"`
+	}) {
+		switch req.Language {
+		case "python3":
+			c.JSON(200, service.RefreshPython3Dependencies())
+		default:
+			c.JSON(400, types.ErrorResponse(-400, "unsupported language"))
+		}
+	})
+}
diff --git a/internal/core/runner/python/setup.go b/internal/core/runner/python/setup.go
@@ -3,6 +3,7 @@ package python
 import (
 	_ "embed"
 	"fmt"
+	"github.com/langgenius/dify-sandbox/internal/static"
 	"os"
 	"os/exec"
 	"path"
@@ -142,3 +143,14 @@ func InstallDependencies(requirements string) error {
 func ListDependencies() []types.Dependency {
 	return python_dependencies.ListDependencies()
 }
+
+func RefreshDependencies() []types.Dependency {
+	log.Info("updating python dependencies...")
+	dependencies := static.GetRunnerDependencies()
+	err := InstallDependencies(dependencies.PythonRequirements)
+	if err != nil {
+		log.Error("failed to update python dependencies: %v", err)
+	}
+	log.Info("python dependencies updated")
+	return python_dependencies.ListDependencies()
+}
diff --git a/internal/server/server.go b/internal/server/server.go
@@ -33,15 +33,20 @@ func initServer() {
 	}
 
 	r := gin.Default()
+	r.Use(gin.Recovery())
+	if gin.Mode() == gin.DebugMode {
+		r.Use(gin.Logger())
+	}
+
 	controller.Setup(r)
 
 	r.Run(fmt.Sprintf(":%d", config.App.Port))
 }
 
 func initDependencies() {
 	log.Info("installing python dependencies...")
-	dependenices := static.GetRunnerDependencies()
-	err := python.InstallDependencies(dependenices.PythonRequirements)
+	dependencies := static.GetRunnerDependencies()
+	err := python.InstallDependencies(dependencies.PythonRequirements)
 	if err != nil {
 		log.Panic("failed to install python dependencies: %v", err)
 	}
@@ -59,7 +64,7 @@ func initDependencies() {
 		ticker := time.NewTicker(30 * time.Minute)
 		for range ticker.C {
 			log.Info("updating python dependencies...")
-			err := python.InstallDependencies(dependenices.PythonRequirements)
+			err := python.InstallDependencies(dependencies.PythonRequirements)
 			if err != nil {
 				log.Error("failed to update python dependencies: %v", err)
 			}

diff --git a/internal/service/python.go b/internal/service/python.go
@@ -63,6 +63,16 @@ func ListPython3Dependencies() *types.DifySandboxResponse {
 	})
 }
 
+type RefreshDependenciesResponse struct {
+	Dependencies []runner_types.Dependency `json:"dependencies"`
+}
+
+func RefreshPython3Dependencies() *types.DifySandboxResponse {
+	return types.SuccessResponse(&RefreshDependenciesResponse{
+		Dependencies: python.RefreshDependencies(),
+	})
+}
+
 type UpdateDependenciesResponse struct{}
 
 func UpdateDependencies() *types.DifySandboxResponse {