Skip to content

Commit

Permalink
feat: support plugin assets
Browse files Browse the repository at this point in the history
  • Loading branch information
Yeuoly committed Sep 4, 2024
1 parent d02ef5e commit 1c2d5e6
Show file tree
Hide file tree
Showing 14 changed files with 178 additions and 11 deletions.
2 changes: 1 addition & 1 deletion .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
5 changes: 4 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -5,4 +5,7 @@ logs/
cmd/**/__debug_bin*
*.zip
.DS_Store
storage/
storage/
__pycache__
media-cache
subprocesses
5 changes: 3 additions & 2 deletions cmd/server/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
115 changes: 113 additions & 2 deletions internal/core/plugin_manager/basic_manager/remap_assets.go
Original file line number Diff line number Diff line change
@@ -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 changes: 2 additions & 0 deletions internal/core/plugin_manager/basic_manager/type.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
4 changes: 4 additions & 0 deletions internal/core/plugin_manager/manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -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...")
Expand Down
9 changes: 9 additions & 0 deletions internal/core/plugin_manager/media_manager/type.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 changes: 10 additions & 0 deletions internal/core/plugin_manager/remote_manager/hooks.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)),
Expand Down Expand Up @@ -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 &&
Expand Down Expand Up @@ -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())
Expand Down
7 changes: 4 additions & 3 deletions internal/core/plugin_manager/remote_manager/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 := ""
Expand Down
1 change: 0 additions & 1 deletion internal/core/plugin_manager/watcher.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
Expand Down
19 changes: 19 additions & 0 deletions internal/server/controllers/plugins.go
Original file line number Diff line number Diff line change
@@ -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 changes: 7 additions & 0 deletions internal/server/http_server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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)
}
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
}
Expand Down

0 comments on commit 1c2d5e6

Please sign in to comment.