feat: dependencies

internal/controller/router.go

@@ -7,9 +7,16 @@ import (
 )
 
 func Setup(eng *gin.Engine) {
-	eng.Use(middleware.MaxRequest(static.GetDifySandboxGlobalConfigurations().MaxRequests))
 	eng.Use(middleware.Auth())
-	eng.Use(middleware.MaxWorker(static.GetDifySandboxGlobalConfigurations().MaxWorkers))
 
-	eng.POST("/v1/sandbox/run", RunSandboxController)
+	eng.POST(
+		"/v1/sandbox/run",
+		middleware.MaxRequest(static.GetDifySandboxGlobalConfigurations().MaxRequests),
+		middleware.MaxWorker(static.GetDifySandboxGlobalConfigurations().MaxWorkers),
+		RunSandboxController,
+	)
+	eng.GET(
+		"/v1/sandbox/dependencies",
+		GetDependencies,
+	)
 }

internal/controller/run.go

@@ -13,7 +13,7 @@ func RunSandboxController(c *gin.Context) {
 		Code          string                    `json:"code" form:"code" binding:"required"`
 		Preload       string                    `json:"preload" form:"preload"`
 		EnableNetwork bool                      `json:"enable_network" form:"enable_network"`
-		Dependencies  []runner_types.Dependency `json:"denpendencies" form:"denpendencies"`
+		Dependencies  []runner_types.Dependency `json:"dependencies" form:"dependencies"`
 	}) {
 		switch req.Language {
 		case "python3":
@@ -30,3 +30,16 @@ func RunSandboxController(c *gin.Context) {
 		}
 	})
 }
+
+func GetDependencies(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.ListPython3Dependencies())
+		default:
+			c.JSON(400, types.ErrorResponse(-400, "unsupported language"))
+		}
+	})
+}

internal/core/runner/python/preload/generic.go

@@ -1,4 +1,4 @@
-package preload
+package dependencies
 
 /*
 import json
@@ -24,7 +24,7 @@ import itertools
 
 func init() {
 	SetupDependency("json", "", "import json")
-	SetupDependency("datetime", "", "import datetime")
+	SetupDependency("datetime", "", "import datetime\nfrom datetime import datetime\ndatetime.strptime('2021-01-01', '%Y-%m-%d')")
 	SetupDependency("math", "", "import math")
 	SetupDependency("random", "", "import random")
 	SetupDependency("re", "", "import re")
@@ -42,5 +42,4 @@ func init() {
 	SetupDependency("functools", "", "import functools")
 	SetupDependency("operator", "", "import operator")
 	SetupDependency("itertools", "", "import itertools")
-	SetupDependency("datetime.datetime", "", "import datetime.datetime\ndatetime.datetime.strptime('2024-05-08', '%Y-%m-%d')")
 }

internal/core/runner/python/dependencies/init.go

@@ -0,0 +1,52 @@
+package dependencies
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/langgenius/dify-sandbox/internal/core/runner/types"
+)
+
+var preload_script_map = map[string]string{}
+var preload_script_map_lock = &sync.RWMutex{}
+
+func SetupDependency(package_name string, version string, script string) {
+	preload_script_map_lock.Lock()
+	defer preload_script_map_lock.Unlock()
+	preload_script_map[fmt.Sprintf("%s==%s", package_name, version)] = script
+}
+
+func GetDependencies(package_name string, version string) string {
+	preload_script_map_lock.RLock()
+	defer preload_script_map_lock.RUnlock()
+	if script, ok := preload_script_map[fmt.Sprintf("%s==%s", package_name, version)]; ok {
+		return script
+	}
+
+	return ""
+}
+
+func ListDependencies() []types.Dependency {
+	dependencies := []types.Dependency{}
+	preload_script_map_lock.RLock()
+	for k := range preload_script_map {
+		parts := strings.Split(k, "==")
+		package_name := ""
+		version := ""
+		if len(parts) == 0 {
+			continue
+		} else if len(parts) == 1 {
+			package_name = parts[0]
+		} else if len(parts) == 2 {
+			package_name = parts[0]
+			version = parts[1]
+		}
+		dependencies = append(dependencies, types.Dependency{
+			Name:    package_name,
+			Version: version,
+		})
+	}
+
+	return dependencies
+}

internal/core/runner/python/preload/jinja2.go

@@ -1,4 +1,4 @@
-package preload
+package dependencies
 
 import "fmt"
 

internal/core/runner/python/preload/init.go

@@ -1,25 +0,0 @@
-package preload
-
-import (
-	"fmt"
-	"sync"
-)
-
-var preload_script_map = map[string]string{}
-var preload_script_map_lock = &sync.RWMutex{}
-
-func SetupDependency(package_name string, version string, script string) {
-	preload_script_map_lock.Lock()
-	defer preload_script_map_lock.Unlock()
-	preload_script_map[package_name] = script
-}
-
-func GetDependencies(package_name string, version string) string {
-	preload_script_map_lock.RLock()
-	defer preload_script_map_lock.RUnlock()
-	if script, ok := preload_script_map[package_name]; ok {
-		return script
-	}
-
-	return fmt.Sprintf("import %s", package_name)
-}

internal/core/runner/python/prescript.py

@@ -3,26 +3,28 @@ if __name__ == "__main__":
     import os
     import sys
     import json
-    import typing
     import time
     import traceback
-    import jinja2
-    import re
 
-    if len(sys.argv) != 4:
+    # setup sys.excepthook
+    def excepthook(type, value, tb):
+        sys.stderr.write("".join(traceback.format_exception(type, value, tb)))
+        sys.stderr.flush()
+        sys.exit(-1)
+    
+    sys.excepthook = excepthook
+
+    if len(sys.argv) != 5:
         sys.exit(-1)
 
     lib = ctypes.CDLL("/tmp/sandbox-python/python.so")
     module = sys.argv[1]
     code = open(module).read()
 
-    with open("/tmp/1.txt", "w") as f:
-        f.write(code)
-
-    def sandbox(uid, gid):
-        lib.DifySeccomp.argtypes = [ctypes.c_uint32, ctypes.c_uint32]
+    def sandbox(uid, gid, enable_network):
+        lib.DifySeccomp.argtypes = [ctypes.c_uint32, ctypes.c_uint32, ctypes.c_bool]
         lib.DifySeccomp.restype = None
-        lib.DifySeccomp(uid, gid)
+        lib.DifySeccomp(uid, gid, enable_network)
     
     uid = int(sys.argv[2])
     gid = int(sys.argv[3])
@@ -34,12 +36,4 @@ if __name__ == "__main__":
 
     sandbox(uid, gid, options.get("enable_network", False))
 
-    # setup sys.excepthook
-    def excepthook(type, value, tb):
-        sys.stderr.write("".join(traceback.format_exception(type, value, tb)))
-        sys.stderr.flush()
-        sys.exit(-1)
-    
-    sys.excepthook = excepthook
-
     exec(code)

internal/core/runner/python/python.go

@@ -11,7 +11,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/langgenius/dify-sandbox/internal/core/runner"
-	python_preload "github.com/langgenius/dify-sandbox/internal/core/runner/python/preload"
+	python_dependencies "github.com/langgenius/dify-sandbox/internal/core/runner/python/dependencies"
 	"github.com/langgenius/dify-sandbox/internal/core/runner/types"
 	"github.com/langgenius/dify-sandbox/internal/static"
 )
@@ -103,7 +103,7 @@ func (p *PythonRunner) InitializeEnvironment(code string, preload string, option
 
 	packages_preload := make([]string, len(options.Dependencies))
 	for i, dependency := range options.Dependencies {
-		packages_preload[i] = python_preload.GetDependencies(dependency.Name, dependency.Version)
+		packages_preload[i] = python_dependencies.GetDependencies(dependency.Name, dependency.Version)
 	}
 	if len(packages_preload) != 0 {
 		preload_script = fmt.Sprintf("%s\n%s", strings.Join(packages_preload, "\n"), preload_script)

internal/core/runner/python/setup.go

@@ -7,6 +7,8 @@ import (
 	"path"
 
 	"github.com/langgenius/dify-sandbox/internal/core/runner"
+	python_dependencies "github.com/langgenius/dify-sandbox/internal/core/runner/python/dependencies"
+	"github.com/langgenius/dify-sandbox/internal/core/runner/types"
 	"github.com/langgenius/dify-sandbox/internal/utils/log"
 )
 
@@ -70,3 +72,7 @@ func InstallDependencies(requirements string) error {
 		return nil
 	})
 }
+
+func ListDependencies() []types.Dependency {
+	return python_dependencies.ListDependencies()
+}

internal/service/python.go

@@ -52,3 +52,13 @@ func RunPython3Code(code string, preload string, options *runner_types.RunnerOpt
 		}
 	}
 }
+
+type ListDependenciesResponse struct {
+	Dependencies []runner_types.Dependency `json:"dependencies"`
+}
+
+func ListPython3Dependencies() *types.DifySandboxResponse {
+	return types.SuccessResponse(&ListDependenciesResponse{
+		Dependencies: python.ListDependencies(),
+	})
+}