feat: support agent management

langgenius · Dec 9, 2024 · 1aecaf9 · 1aecaf9
1 parent 7bb4705
commit 1aecaf9
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 3 deletions.
diff --git a/internal/server/controllers/agent.go b/internal/server/controllers/agent.go
@@ -1,6 +1,8 @@
 package controllers
 
 import (
+	"net/http"
+
 	"github.com/gin-gonic/gin"
 	"github.com/langgenius/dify-plugin-daemon/internal/service"
 	"github.com/langgenius/dify-plugin-daemon/internal/types/app"
@@ -20,3 +22,23 @@ func InvokeAgent(config *app.Config) gin.HandlerFunc {
 		)
 	}
 }
+
+func ListAgents(c *gin.Context) {
+	BindRequest(c, func(request struct {
+		TenantID string `uri:"tenant_id" validate:"required"`
+		Page     int    `form:"page" validate:"required,min=1"`
+		PageSize int    `form:"page_size" validate:"required,min=1,max=256"`
+	}) {
+		c.JSON(http.StatusOK, service.ListAgents(request.TenantID, request.Page, request.PageSize))
+	})
+}
+
+func GetAgent(c *gin.Context) {
+	BindRequest(c, func(request struct {
+		TenantID string `uri:"tenant_id" validate:"required"`
+		PluginID string `form:"plugin_id" validate:"required"`
+		Provider string `form:"provider" validate:"required"`
+	}) {
+		c.JSON(http.StatusOK, service.GetAgent(request.TenantID, request.PluginID, request.Provider))
+	})
+}
diff --git a/internal/server/http_server.go b/internal/server/http_server.go
@@ -115,16 +115,16 @@ func (app *App) endpointManagementGroup(group *gin.RouterGroup) {
 }
 
 func (app *App) pluginManagementGroup(group *gin.RouterGroup, config *app.Config) {
-	group.POST("/install/upload/package", controllers.UploadPlugin(config))
-	group.POST("/install/upload/bundle", controllers.UploadBundle(config))
+	group.POST("/install/upload/package", gzip.Gzip(gzip.DefaultCompression), controllers.UploadPlugin(config))
+	group.POST("/install/upload/bundle", gzip.Gzip(gzip.DefaultCompression), controllers.UploadBundle(config))
 	group.POST("/install/identifiers", controllers.InstallPluginFromIdentifiers(config))
 	group.POST("/install/upgrade", controllers.UpgradePlugin(config))
 	group.GET("/install/tasks/:id", controllers.FetchPluginInstallationTask)
 	group.POST("/install/tasks/delete_all", controllers.DeleteAllPluginInstallationTasks)
 	group.POST("/install/tasks/:id/delete", controllers.DeletePluginInstallationTask)
 	group.POST("/install/tasks/:id/delete/*identifier", controllers.DeletePluginInstallationItemFromTask)
 	group.GET("/install/tasks", controllers.FetchPluginInstallationTasks)
-	group.GET("/fetch/manifest", controllers.FetchPluginManifest)
+	group.GET("/fetch/manifest", gzip.Gzip(gzip.DefaultCompression), controllers.FetchPluginManifest)
 	group.GET("/fetch/identifier", controllers.FetchPluginFromIdentifier)
 	group.POST("/uninstall", controllers.UninstallPlugin)
 	group.GET("/list", gzip.Gzip(gzip.DefaultCompression), controllers.ListPlugins)
@@ -134,6 +134,8 @@ func (app *App) pluginManagementGroup(group *gin.RouterGroup, config *app.Config
 	group.GET("/tools", gzip.Gzip(gzip.DefaultCompression), controllers.ListTools)
 	group.GET("/tool", gzip.Gzip(gzip.DefaultCompression), controllers.GetTool)
 	group.POST("/tools/check_existence", controllers.CheckToolExistence)
+	group.GET("/agents", gzip.Gzip(gzip.DefaultCompression), controllers.ListAgents)
+	group.GET("/agent", gzip.Gzip(gzip.DefaultCompression), controllers.GetAgent)
 }
 
 func (app *App) pluginAssetGroup(group *gin.RouterGroup) {

diff --git a/internal/service/manage_plugin.go b/internal/service/manage_plugin.go
@@ -218,6 +218,10 @@ func GetTool(tenant_id string, plugin_id string, provider string) *entities.Resp
 	)
 
 	if err != nil {
+		if err == db.ErrDatabaseNotFound {
+			return exception.ErrPluginNotFound().ToResponse()
+		}
+
 		return exception.InternalServerError(err).ToResponse()
 	}
 
@@ -262,3 +266,37 @@ func CheckToolExistence(tenantId string, providerIds []RequestCheckToolExistence
 
 	return entities.NewSuccessResponse(existence)
 }
+
+func ListAgents(tenant_id string, page int, page_size int) *entities.Response {
+	providers, err := db.GetAll[models.AgentInstallation](
+		db.Equal("tenant_id", tenant_id),
+		db.Page(page, page_size),
+	)
+
+	if err != nil {
+		return exception.InternalServerError(err).ToResponse()
+	}
+
+	return entities.NewSuccessResponse(providers)
+}
+
+func GetAgent(tenant_id string, plugin_id string, provider string) *entities.Response {
+	agent, err := db.GetOne[models.AgentInstallation](
+		db.Equal("tenant_id", tenant_id),
+		db.Equal("plugin_id", plugin_id),
+	)
+
+	if err != nil {
+		if err == db.ErrDatabaseNotFound {
+			return exception.ErrPluginNotFound().ToResponse()
+		}
+
+		return exception.InternalServerError(err).ToResponse()
+	}
+
+	if agent.Provider != provider {
+		return exception.ErrPluginNotFound().ToResponse()
+	}
+
+	return entities.NewSuccessResponse(agent)
+}
diff --git a/internal/types/models/agent.go b/internal/types/models/agent.go
@@ -0,0 +1,12 @@
+package models
+
+import "github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
+
+type AgentInstallation struct {
+	Model
+	TenantID               string                                   `json:"tenant_id" gorm:"column:tenant_id;type:uuid;index;not null"`
+	Provider               string                                   `json:"provider" gorm:"column:provider;size:127;index;not null"`
+	PluginUniqueIdentifier string                                   `json:"plugin_unique_identifier" gorm:"index;size:255"`
+	PluginID               string                                   `json:"plugin_id" gorm:"index;size:255"`
+	Declaration            plugin_entities.AgentProviderDeclaration `json:"declaration" gorm:"serializer:json;type:text;size:65535;not null"`
+}
diff --git a/internal/types/models/curd/atomic.go b/internal/types/models/curd/atomic.go
@@ -116,6 +116,22 @@ func InstallPlugin(
 			}
 		}
 
+		// create agent installation
+		if declaration.Agent != nil {
+			agentInstallation := &models.AgentInstallation{
+				PluginID:               pluginToBeReturns.PluginID,
+				PluginUniqueIdentifier: pluginToBeReturns.PluginUniqueIdentifier,
+				TenantID:               tenant_id,
+				Provider:               declaration.Agent.Identity.Name,
+				Declaration:            *declaration.Agent,
+			}
+
+			err := db.Create(agentInstallation, tx)
+			if err != nil {
+				return err
+			}
+		}
+
 		// create model installation
 		if declaration.Model != nil {
 			modelInstallation := &models.AIModelInstallation{
@@ -229,6 +245,20 @@ func UninstallPlugin(
 			}
 		}
 
+		// delete agent installation
+		if declaration.Agent != nil {
+			agentInstallation := &models.AgentInstallation{
+				PluginID:               pluginToBeReturns.PluginID,
+				PluginUniqueIdentifier: pluginToBeReturns.PluginUniqueIdentifier,
+				TenantID:               tenant_id,
+			}
+
+			err := db.DeleteByCondition(&agentInstallation, tx)
+			if err != nil {
+				return err
+			}
+		}
+
 		// delete model installation
 		if declaration.Model != nil {
 			modelInstallation := &models.AIModelInstallation{