Browse Source

feat: support plugin assets

Yeuoly 10 months ago
parent
commit
1c2d5e6bca

+ 1 - 1
.env.example

@@ -24,7 +24,7 @@ DB_PORT=5432
 DB_DATABASE=dify_plugin
 
 PERSISTENCE_STORAGE_TYPE=local
-PERSISTENCE_STORAGE_LOCAL_PATH=./storage
+PERSISTENCE_STORAGE_LOCAL_PATH=./storage/persistence
 PERSISTENCE_STORAGE_AWS_S3_REGION=us-east-1
 PERSISTENCE_STORAGE_AWS_S3_ACCESS_KEY=
 PERSISTENCE_STORAGE_AWS_S3_SECRET_KEY=

+ 4 - 1
.gitignore

@@ -5,4 +5,7 @@ logs/
 cmd/**/__debug_bin*
 *.zip
 .DS_Store
-storage/
+storage/
+__pycache__
+media-cache
+subprocesses

+ 3 - 2
cmd/server/main.go

@@ -47,8 +47,9 @@ func setDefault(config *app.Config) {
 	setDefaultBool(&config.PluginRemoteInstallingEnabled, true)
 	setDefaultBool(&config.PluginEndpointEnabled, true)
 	setDefaultString(&config.DBSslMode, "disable")
-	setDefaultString(&config.PluginMediaCachePath, "/var/dify-plugin-daemon/media-cache")
-	setDefaultString(&config.ProcessCachingPath, "/var/dify-plugin-daemon/subprocesses")
+	setDefaultString(&config.PluginMediaCachePath, "./storage/assets")
+	setDefaultString(&config.PersistenceStorageLocalPath, "./storage/persistence")
+	setDefaultString(&config.ProcessCachingPath, "./storage/subprocesses")
 }
 
 func setDefaultInt[T constraints.Integer](value *T, defaultValue T) {

+ 113 - 2
internal/core/plugin_manager/basic_manager/remap_assets.go

@@ -1,12 +1,123 @@
 package basic_manager
 
-import "github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
+import (
+	"fmt"
+
+	"github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
+)
 
 // RemapAssets will take the assets and remap them to a media id
 func (r *BasicPluginRuntime) RemapAssets(
 	declaration *plugin_entities.PluginDeclaration,
 	assets map[string][]byte,
 ) error {
-	// TODO: implement
+	remapped_asset_ids := make(map[string]string)
+	remap := func(filename string) (string, error) {
+		if id, ok := remapped_asset_ids[filename]; ok {
+			return id, nil
+		}
+
+		file, ok := assets[filename]
+		if !ok {
+			return "", fmt.Errorf("file not found: %s", filename)
+		}
+
+		id, err := r.mediaManager.Upload(file)
+		if err != nil {
+			return "", err
+		}
+
+		r.assets_ids = append(r.assets_ids, id)
+
+		remapped_asset_ids[filename] = id
+		return id, nil
+	}
+
+	var err error
+
+	if declaration.Model != nil {
+		if declaration.Model.IconSmall != nil {
+			if declaration.Model.IconSmall.EnUS != "" {
+				declaration.Model.IconSmall.EnUS, err = remap(declaration.Model.IconSmall.EnUS)
+				if err != nil {
+					return err
+				}
+			}
+
+			if declaration.Model.IconSmall.ZhHans != "" {
+				declaration.Model.IconSmall.ZhHans, err = remap(declaration.Model.IconSmall.ZhHans)
+				if err != nil {
+					return err
+				}
+			}
+
+			if declaration.Model.IconSmall.JaJp != "" {
+				declaration.Model.IconSmall.JaJp, err = remap(declaration.Model.IconSmall.JaJp)
+				if err != nil {
+					return err
+				}
+			}
+
+			if declaration.Model.IconSmall.PtBr != "" {
+				declaration.Model.IconSmall.PtBr, err = remap(declaration.Model.IconSmall.PtBr)
+				if err != nil {
+					return err
+				}
+			}
+		}
+
+		if declaration.Model.IconLarge != nil {
+			if declaration.Model.IconLarge.EnUS != "" {
+				declaration.Model.IconLarge.EnUS, err = remap(declaration.Model.IconLarge.EnUS)
+				if err != nil {
+					return err
+				}
+			}
+
+			if declaration.Model.IconLarge.ZhHans != "" {
+				declaration.Model.IconLarge.ZhHans, err = remap(declaration.Model.IconLarge.ZhHans)
+				if err != nil {
+					return err
+				}
+			}
+
+			if declaration.Model.IconLarge.JaJp != "" {
+				declaration.Model.IconLarge.JaJp, err = remap(declaration.Model.IconLarge.JaJp)
+				if err != nil {
+					return err
+				}
+			}
+
+			if declaration.Model.IconLarge.PtBr != "" {
+				declaration.Model.IconLarge.PtBr, err = remap(declaration.Model.IconLarge.PtBr)
+				if err != nil {
+					return err
+				}
+			}
+		}
+	}
+
+	if declaration.Tool != nil {
+		if declaration.Tool.Identity.Icon != "" {
+			declaration.Tool.Identity.Icon, err = remap(declaration.Tool.Identity.Icon)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	if declaration.Icon != "" {
+		declaration.Icon, err = remap(declaration.Icon)
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
+
+func (r *BasicPluginRuntime) ClearAssets() {
+	for _, id := range r.assets_ids {
+		r.mediaManager.Delete(id)
+	}
+}

+ 2 - 0
internal/core/plugin_manager/basic_manager/type.go

@@ -4,6 +4,8 @@ import "github.com/langgenius/dify-plugin-daemon/internal/core/plugin_manager/me
 
 type BasicPluginRuntime struct {
 	mediaManager *media_manager.MediaManager
+
+	assets_ids []string
 }
 
 func NewBasicPluginRuntime(mediaManager *media_manager.MediaManager) BasicPluginRuntime {

+ 4 - 0
internal/core/plugin_manager/manager.go

@@ -83,6 +83,10 @@ func (p *PluginManager) Get(identity plugin_entities.PluginUniqueIdentifier) plu
 	return nil
 }
 
+func (p *PluginManager) GetAsset(id string) ([]byte, error) {
+	return p.mediaManager.Get(id)
+}
+
 func (p *PluginManager) Init(configuration *app.Config) {
 	// TODO: init plugin manager
 	log.Info("start plugin manager daemon...")

+ 9 - 0
internal/core/plugin_manager/media_manager/type.go

@@ -69,3 +69,12 @@ func (m *MediaManager) Get(id string) ([]byte, error) {
 
 	return file, nil
 }
+
+func (m *MediaManager) Delete(id string) error {
+	// delete from cache
+	m.cache.Remove(id)
+
+	// delete from storage
+	filepath := path.Join(m.storagePath, id)
+	return os.Remove(filepath)
+}

+ 10 - 0
internal/core/plugin_manager/remote_manager/hooks.go

@@ -5,6 +5,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/langgenius/dify-plugin-daemon/internal/core/plugin_manager/basic_manager"
 	"github.com/langgenius/dify-plugin-daemon/internal/core/plugin_manager/media_manager"
 	"github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/cache"
@@ -49,6 +50,10 @@ func (s *DifyServer) OnOpen(c gnet.Conn) (out []byte, action gnet.Action) {
 	// new plugin connected
 	c.SetContext(&codec{})
 	runtime := &RemotePluginRuntime{
+		BasicPluginRuntime: basic_manager.NewBasicPluginRuntime(
+			s.mediaManager,
+		),
+
 		conn:           c,
 		response:       stream.NewStreamResponse[[]byte](512),
 		callbacks:      make(map[string][]func([]byte)),
@@ -91,6 +96,9 @@ func (s *DifyServer) OnClose(c gnet.Conn, err error) (action gnet.Action) {
 	// close plugin
 	plugin.onDisconnected()
 
+	// clear assets
+	plugin.ClearAssets()
+
 	// uninstall plugin
 	if plugin.handshake && plugin.registration_transferred &&
 		plugin.endpoints_registration_transferred &&
@@ -253,6 +261,8 @@ func (s *DifyServer) onMessage(runtime *RemotePluginRuntime, message []byte) {
 			return
 		}
 
+		runtime.assets_transferred = true
+
 		runtime.checksum = runtime.calculateChecksum()
 		runtime.InitState()
 		runtime.SetActiveAt(time.Now())

+ 4 - 3
internal/core/plugin_manager/remote_manager/server_test.go

@@ -168,9 +168,10 @@ func TestAcceptConnection(t *testing.T) {
 	conn.Write([]byte("\n"))
 	conn.Write(handle_shake_message)
 	conn.Write([]byte("\n"))
-	conn.Write([]byte("[]\n"))
-	conn.Write([]byte("[]\n"))
-	conn.Write([]byte("[]\n"))
+	conn.Write([]byte("[]\n")) // transfer tool
+	conn.Write([]byte("[]\n")) // transfer model
+	conn.Write([]byte("[]\n")) // transfer endpoint
+	conn.Write([]byte("[]\n")) // transfer file
 	closed_chan := make(chan bool)
 
 	msg := ""

+ 0 - 1
internal/core/plugin_manager/watcher.go

@@ -176,7 +176,6 @@ func (p *PluginManager) loadPlugin(plugin_path string) (*pluginRuntimeWithDecode
 		return nil, errors.Join(fmt.Errorf("plugin already exists: %s", manifest.Identity()), err)
 	}
 
-	// TODO: use plugin unique id as the working directory
 	checksum, err := checksum.CalculateChecksum(decoder)
 	if err != nil {
 		return nil, errors.Join(err, fmt.Errorf("calculate checksum error"))

+ 19 - 0
internal/server/controllers/plugins.go

@@ -1 +1,20 @@
 package controllers
+
+import (
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"github.com/langgenius/dify-plugin-daemon/internal/core/plugin_manager"
+)
+
+func GetAsset(c *gin.Context) {
+	plugin_manager := plugin_manager.GetGlobalPluginManager()
+	asset, err := plugin_manager.GetAsset(c.Param("id"))
+
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.Data(http.StatusOK, "application/octet-stream", asset)
+}

+ 7 - 0
internal/server/http_server.go

@@ -23,6 +23,7 @@ func (app *App) server(config *app.Config) func() {
 	app.endpointGroup(engine.Group("/e"), config)
 	app.awsLambdaTransactionGroup(engine.Group("/backwards-invocation"), config)
 	app.endpointManagementGroup(engine.Group("/endpoint"))
+	app.pluginGroup(engine.Group("/plugin"), config)
 
 	srv := &http.Server{
 		Addr:    fmt.Sprintf(":%d", config.ServerPort),
@@ -93,3 +94,9 @@ func (app *App) endpointManagementGroup(group *gin.RouterGroup) {
 	group.POST("/remove", controllers.RemoveEndpoint)
 	group.GET("/list", controllers.ListEndpoints)
 }
+
+func (app *App) pluginGroup(group *gin.RouterGroup, config *app.Config) {
+	group.Use(CheckingKey(config.PluginInnerApiKey))
+
+	group.GET("/asset/:id", controllers.GetAsset)
+}

+ 1 - 0
internal/types/entities/plugin_entities/plugin_declaration.go

@@ -131,6 +131,7 @@ type PluginDeclarationWithoutAdvancedFields struct {
 	Type      DifyManifestType          `json:"type" yaml:"type,omitempty" validate:"required,eq=plugin"`
 	Author    string                    `json:"author" yaml:"author,omitempty" validate:"required,max=128"`
 	Name      string                    `json:"name" yaml:"name,omitempty" validate:"required,max=128"`
+	Icon      string                    `json:"icon" yaml:"icon,omitempty" validate:"required,max=128"`
 	Label     I18nObject                `json:"label" yaml:"label" validate:"required"`
 	CreatedAt time.Time                 `json:"created_at" yaml:"created_at,omitempty" validate:"required"`
 	Resource  PluginResourceRequirement `json:"resource" yaml:"resource,omitempty" validate:"required"`

+ 1 - 1
internal/types/entities/plugin_entities/tool_declaration.go

@@ -160,7 +160,7 @@ type ToolProviderIdentity struct {
 	Author      string      `json:"author" validate:"required"`
 	Name        string      `json:"name" validate:"required"`
 	Description I18nObject  `json:"description" validate:"required"`
-	Icon        []byte      `json:"icon" validate:"required"`
+	Icon        string      `json:"icon" validate:"required"`
 	Label       I18nObject  `json:"label" validate:"required"`
 	Tags        []ToolLabel `json:"tags" validate:"required,dive,tool_label"`
 }