diff --git a/.env.example b/.env.example index 02dfe80b5..4263fd893 100644 --- a/.env.example +++ b/.env.example @@ -235,6 +235,9 @@ PLUGIN_RUNTIME_MAX_BUFFER_SIZE=5242880 # plugin install/upgrade timeout in minutes PLUGIN_INSTALL_TIMEOUT=15 +# how long a cached model schema is kept, in minutes +PLUGIN_MODEL_SCHEMA_CACHE_TTL=60 + # dify backwards invocation write timeout in milliseconds DIFY_BACKWARDS_INVOCATION_WRITE_TIMEOUT=5000 # dify backwards invocation read timeout in milliseconds diff --git a/internal/core/io_tunnel/datasource.gen.go b/internal/core/io_tunnel/datasource.gen.go index a92fdf951..0cffe3d50 100644 --- a/internal/core/io_tunnel/datasource.gen.go +++ b/internal/core/io_tunnel/datasource.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/datasource_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func DatasourceValidateCredentials( diff --git a/internal/core/io_tunnel/dynamic_parameter.gen.go b/internal/core/io_tunnel/dynamic_parameter.gen.go index 28968a0dd..03b51dcdc 100644 --- a/internal/core/io_tunnel/dynamic_parameter.gen.go +++ b/internal/core/io_tunnel/dynamic_parameter.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/dynamic_select_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func FetchDynamicParameterOptions( diff --git a/internal/core/io_tunnel/model.gen.go b/internal/core/io_tunnel/model.gen.go index ba4d2ea8a..3607ebe24 100644 --- a/internal/core/io_tunnel/model.gen.go +++ b/internal/core/io_tunnel/model.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func InvokeLLM( @@ -177,16 +177,3 @@ func ValidateModelCredentials( 1, ) } - -func GetAIModelSchema( - session *session_manager.Session, - request *requests.RequestGetAIModelSchema, -) ( - *stream.Stream[model_entities.GetModelSchemasResponse], error, -) { - return GenericInvokePlugin[requests.RequestGetAIModelSchema, model_entities.GetModelSchemasResponse]( - session, - request, - 1, - ) -} diff --git a/internal/core/io_tunnel/oauth.gen.go b/internal/core/io_tunnel/oauth.gen.go index 9b1621998..b652fd29c 100644 --- a/internal/core/io_tunnel/oauth.gen.go +++ b/internal/core/io_tunnel/oauth.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/oauth_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func GetAuthorizationURL( diff --git a/internal/core/io_tunnel/tool.gen.go b/internal/core/io_tunnel/tool.gen.go index 080009814..965777985 100644 --- a/internal/core/io_tunnel/tool.gen.go +++ b/internal/core/io_tunnel/tool.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" "github.com/langgenius/dify-plugin-daemon/pkg/entities/tool_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func InvokeTool( diff --git a/internal/core/io_tunnel/trigger.gen.go b/internal/core/io_tunnel/trigger.gen.go index 4941bd7e9..48ec208a6 100644 --- a/internal/core/io_tunnel/trigger.gen.go +++ b/internal/core/io_tunnel/trigger.gen.go @@ -4,8 +4,8 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func TriggerInvokeEvent( diff --git a/internal/core/plugin_manager/manager.go b/internal/core/plugin_manager/manager.go index c3c76eac8..c27377a3f 100644 --- a/internal/core/plugin_manager/manager.go +++ b/internal/core/plugin_manager/manager.go @@ -2,6 +2,7 @@ package plugin_manager import ( "fmt" + "time" lru "github.com/hashicorp/golang-lru/v2" "github.com/langgenius/dify-cloud-kit/oss" @@ -16,6 +17,7 @@ import ( "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/plugin_packager/decoder" "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache/helper" "github.com/langgenius/dify-plugin-daemon/pkg/utils/log" "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" ) @@ -138,6 +140,9 @@ func (p *PluginManager) Launch(configuration *app.Config) { } cache.SetKeyPrefix(configuration.RedisKeyPrefix) + helper.SetModelSchemaCacheTTL( + time.Duration(configuration.PluginModelSchemaCacheTTL) * time.Minute, + ) // init redis client if configuration.RedisUseSentinel { diff --git a/internal/server/controllers/definitions/definitions.go b/internal/server/controllers/definitions/definitions.go index 9f3b3f134..8efe2281b 100644 --- a/internal/server/controllers/definitions/definitions.go +++ b/internal/server/controllers/definitions/definitions.go @@ -203,17 +203,17 @@ var PluginDispatchers = []PluginDispatcher{ BufferSize: 1, Path: "/model/validate_model_credentials", }, - { - Name: "GetAIModelSchema", - RequestType: requests.RequestGetAIModelSchema{}, - ResponseType: model_entities.GetModelSchemasResponse{}, - AccessType: access_types.PLUGIN_ACCESS_TYPE_MODEL, - AccessAction: access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS, - AccessTypeString: "access_types.PLUGIN_ACCESS_TYPE_MODEL", - AccessActionString: "access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS", - BufferSize: 1, - Path: "/model/schema", - }, + // { // No need to implement this for now, it is cached in the model schema service + // Name: "GetAIModelSchema", + // RequestType: requests.RequestGetAIModelSchema{}, + // ResponseType: model_entities.GetModelSchemasResponse{}, + // AccessType: access_types.PLUGIN_ACCESS_TYPE_MODEL, + // AccessAction: access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS, + // AccessTypeString: "access_types.PLUGIN_ACCESS_TYPE_MODEL", + // AccessActionString: "access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS", + // BufferSize: 1, + // Path: "/model/schema", + // }, // { // No need to implement this for now, it has its special implementation in the agent service // Name: "InvokeAgentStrategy", // RequestType: requests.RequestInvokeAgentStrategy{}, diff --git a/internal/server/controllers/generator/generator.go b/internal/server/controllers/generator/generator.go index 54756a64e..0449dd744 100644 --- a/internal/server/controllers/generator/generator.go +++ b/internal/server/controllers/generator/generator.go @@ -114,7 +114,7 @@ func GeneratePluginDaemon(accessType access_types.PluginAccessType, dispatchers tmpl := template.Must(template.New("pluginDaemon").Parse(pluginDaemonTemplate)) // Create output file - outputPath := filepath.Join("internal", "core", "plugin_daemon", strings.ToLower(string(accessType))+".gen.go") + outputPath := filepath.Join("internal", "core", "io_tunnel", strings.ToLower(string(accessType))+".gen.go") // Execute template var buf strings.Builder diff --git a/internal/server/controllers/model.gen.go b/internal/server/controllers/model.gen.go index 7fb488937..e7e5d3c30 100644 --- a/internal/server/controllers/model.gen.go +++ b/internal/server/controllers/model.gen.go @@ -178,16 +178,3 @@ func ValidateModelCredentials(config *app.Config) gin.HandlerFunc { ) } } - -func GetAIModelSchema(config *app.Config) gin.HandlerFunc { - type request = plugin_entities.InvokePluginRequest[requests.RequestGetAIModelSchema] - - return func(c *gin.Context) { - BindPluginDispatchRequest( - c, - func(itr request) { - service.GetAIModelSchema(&itr, c, config.PluginMaxExecutionTimeout) - }, - ) - } -} diff --git a/internal/server/controllers/model_schema.go b/internal/server/controllers/model_schema.go new file mode 100644 index 000000000..3d828d043 --- /dev/null +++ b/internal/server/controllers/model_schema.go @@ -0,0 +1,22 @@ +package controllers + +import ( + "github.com/gin-gonic/gin" + "github.com/langgenius/dify-plugin-daemon/internal/service" + "github.com/langgenius/dify-plugin-daemon/internal/types/app" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" +) + +func GetAIModelSchema(config *app.Config) gin.HandlerFunc { + type request = plugin_entities.InvokePluginRequest[requests.RequestGetAIModelSchema] + + return func(c *gin.Context) { + BindPluginDispatchRequest( + c, + func(itr request) { + service.GetAIModelSchema(&itr, c, config.PluginMaxExecutionTimeout) + }, + ) + } +} diff --git a/internal/server/http_server.gen.go b/internal/server/http_server.gen.go index 4392b0121..9e7d4ed81 100644 --- a/internal/server/http_server.gen.go +++ b/internal/server/http_server.gen.go @@ -24,7 +24,6 @@ func (app *App) setupGeneratedRoutes(group *gin.RouterGroup, config *app.Config) group.POST("/moderation/invoke", controllers.InvokeModeration(config)) group.POST("/model/validate_provider_credentials", controllers.ValidateProviderCredentials(config)) group.POST("/model/validate_model_credentials", controllers.ValidateModelCredentials(config)) - group.POST("/model/schema", controllers.GetAIModelSchema(config)) group.POST("/oauth/get_authorization_url", controllers.GetAuthorizationURL(config)) group.POST("/oauth/get_credentials", controllers.GetCredentials(config)) group.POST("/oauth/refresh_credentials", controllers.RefreshCredentials(config)) diff --git a/internal/server/http_server.go b/internal/server/http_server.go index 925cbdce3..cae40de09 100644 --- a/internal/server/http_server.go +++ b/internal/server/http_server.go @@ -121,6 +121,7 @@ func (app *App) pluginDispatchGroup(group *gin.RouterGroup, config *app.Config) group.Use(app.InitClusterID()) group.POST("/agent_strategy/invoke", controllers.InvokeAgentStrategy(config)) + group.POST("/model/schema", controllers.GetAIModelSchema(config)) group.POST("/model/polling/start", controllers.StartPolling()) group.POST("/model/polling/check", controllers.CheckPolling()) @@ -245,6 +246,8 @@ func (app *App) invokeGroup(group *gin.RouterGroup, config *app.Config) { dispatchGroup.POST("/agent_strategy/invoke", controllers.InvokeAgentStrategy(config)) + dispatchGroup.POST("/model/schema", + controllers.GetAIModelSchema(config)) app.setupGeneratedRoutes(dispatchGroup, config) } diff --git a/internal/service/datasource.gen.go b/internal/service/datasource.gen.go index deee6ee1d..b9a5c9b7e 100644 --- a/internal/service/datasource.gen.go +++ b/internal/service/datasource.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/datasource_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func DatasourceValidateCredentials( diff --git a/internal/service/dynamic_parameter.gen.go b/internal/service/dynamic_parameter.gen.go index 7f28cb59b..35b8e0c9b 100644 --- a/internal/service/dynamic_parameter.gen.go +++ b/internal/service/dynamic_parameter.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/dynamic_select_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func FetchDynamicParameterOptions( diff --git a/internal/service/model.gen.go b/internal/service/model.gen.go index 219580f24..d77758619 100644 --- a/internal/service/model.gen.go +++ b/internal/service/model.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func InvokeLLM( @@ -74,7 +74,7 @@ func InvokeMultimodalEmbedding( return io_tunnel.InvokeMultimodalEmbedding(session, &r.Data) }, access_types.PLUGIN_ACCESS_TYPE_MODEL, - access_types.PLUGIN_ACCESS_ACTION_INVOKE_MULTIMODAL_EMBEDDING, + access_types.PLUGIN_ACCESS_ACTION_INVOKE_TEXT_EMBEDDING, r, ctx, max_timeout_seconds, @@ -125,7 +125,7 @@ func InvokeMultimodalRerank( return io_tunnel.InvokeMultimodalRerank(session, &r.Data) }, access_types.PLUGIN_ACCESS_TYPE_MODEL, - access_types.PLUGIN_ACCESS_ACTION_INVOKE_MULTIMODAL_RERANK, + access_types.PLUGIN_ACCESS_ACTION_INVOKE_RERANK, r, ctx, max_timeout_seconds, @@ -233,20 +233,3 @@ func ValidateModelCredentials( max_timeout_seconds, ) } - -func GetAIModelSchema( - r *plugin_entities.InvokePluginRequest[requests.RequestGetAIModelSchema], - ctx *gin.Context, - max_timeout_seconds int, -) { - baseSSEWithSession( - func(session *session_manager.Session) (*stream.Stream[model_entities.GetModelSchemasResponse], error) { - return io_tunnel.GetAIModelSchema(session, &r.Data) - }, - access_types.PLUGIN_ACCESS_TYPE_MODEL, - access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS, - r, - ctx, - max_timeout_seconds, - ) -} diff --git a/internal/service/model_schema.go b/internal/service/model_schema.go new file mode 100644 index 000000000..6efdeac91 --- /dev/null +++ b/internal/service/model_schema.go @@ -0,0 +1,88 @@ +package service + +import ( + "github.com/gin-gonic/gin" + "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" + "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" + "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" + "github.com/langgenius/dify-plugin-daemon/pkg/entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache/helper" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" +) + +func GetAIModelSchema( + r *plugin_entities.InvokePluginRequest[requests.RequestGetAIModelSchema], + ctx *gin.Context, + max_timeout_seconds int, +) { + identity := helper.ModelSchemaCacheIdentity{ + UniqueIdentifier: r.UniqueIdentifier, + ModelType: string(r.Data.ModelType), + Model: r.Data.Model, + CredentialType: r.Data.Credentials.CredentialType, + Credentials: r.Data.Credentials.Credentials, + } + + if cached := helper.GetCachedModelSchema(identity); cached != nil { + writeSingleSSEChunk(ctx, *cached) + return + } + + baseSSEWithSession( + func(session *session_manager.Session) (*stream.Stream[model_entities.GetModelSchemasResponse], error) { + response, err := io_tunnel.GenericInvokePlugin[ + requests.RequestGetAIModelSchema, + model_entities.GetModelSchemasResponse, + ](session, &r.Data, 1) + if err != nil { + return nil, err + } + + cacheModelSchemaStream(response, identity) + return response, nil + }, + access_types.PLUGIN_ACCESS_TYPE_MODEL, + access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS, + r, + ctx, + max_timeout_seconds, + ) +} + +// The plugin answers with a single chunk. Anything else, including an error mid-stream, +// leaves the cache untouched rather than storing a partial answer. +func cacheModelSchemaStream( + response *stream.Stream[model_entities.GetModelSchemasResponse], + identity helper.ModelSchemaCacheIdentity, +) { + var chunks []model_entities.GetModelSchemasResponse + + response.Filter(func(chunk model_entities.GetModelSchemasResponse) error { + chunks = append(chunks, chunk) + return nil + }) + + response.OnClose(func() { + if len(chunks) != 1 || chunks[0].ModelSchema == nil { + return + } + helper.StoreModelSchema(identity, chunks[0]) + }) +} + +// Mirrors the wire format baseSSEService writes, so a cache hit is indistinguishable to the +// caller from a plugin invocation. +func writeSingleSSEChunk(ctx *gin.Context, chunk model_entities.GetModelSchemasResponse) { + writer := ctx.Writer + writer.WriteHeader(200) + writer.Header().Set("Content-Type", "text/event-stream") + + writer.Write([]byte("data: ")) + writer.Write(parser.MarshalJsonBytes(entities.NewSuccessResponse(chunk))) + writer.Write([]byte("\n\n")) + writer.Flush() +} diff --git a/internal/service/oauth.gen.go b/internal/service/oauth.gen.go index 19d426582..ddddcc09f 100644 --- a/internal/service/oauth.gen.go +++ b/internal/service/oauth.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/oauth_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func GetAuthorizationURL( diff --git a/internal/service/tool.gen.go b/internal/service/tool.gen.go index 33dd4f23c..200198cef 100644 --- a/internal/service/tool.gen.go +++ b/internal/service/tool.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" "github.com/langgenius/dify-plugin-daemon/pkg/entities/tool_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func InvokeTool( diff --git a/internal/service/trigger.gen.go b/internal/service/trigger.gen.go index 8b73133d3..d7798efa5 100644 --- a/internal/service/trigger.gen.go +++ b/internal/service/trigger.gen.go @@ -7,9 +7,9 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func TriggerInvokeEvent( diff --git a/internal/types/app/config.go b/internal/types/app/config.go index 04c33de8f..5582af55e 100644 --- a/internal/types/app/config.go +++ b/internal/types/app/config.go @@ -115,6 +115,9 @@ type Config struct { PluginInstalledPath string `envconfig:"PLUGIN_INSTALLED_PATH" validate:"required"` // where the plugin finally installed PluginPackageCachePath string `envconfig:"PLUGIN_PACKAGE_CACHE_PATH"` // where plugin packages stored + // how long a cached model schema is kept, in minutes; the key pins the plugin version, so this only bounds memory + PluginModelSchemaCacheTTL int `envconfig:"PLUGIN_MODEL_SCHEMA_CACHE_TTL"` + // request timeout PluginMaxExecutionTimeout int `envconfig:"PLUGIN_MAX_EXECUTION_TIMEOUT" validate:"required"` diff --git a/internal/types/app/default.go b/internal/types/app/default.go index 34f8b1241..1def38e12 100644 --- a/internal/types/app/default.go +++ b/internal/types/app/default.go @@ -40,6 +40,7 @@ func (config *Config) SetDefault() { setDefaultInt(&config.PluginLocalLaunchingConcurrent, 2) setDefaultInt(&config.PersistenceStorageMaxSize, 100*1024*1024) setDefaultString(&config.PluginPackageCachePath, "plugin_packages") + setDefaultInt(&config.PluginModelSchemaCacheTTL, 60) setDefaultString(&config.PythonInterpreterPath, "/usr/bin/python3") setDefaultInt(&config.PythonEnvInitTimeout, 120) setDefaultInt(&config.DifyInvocationWriteTimeout, 5000) diff --git a/pkg/utils/cache/helper/model_schema.go b/pkg/utils/cache/helper/model_schema.go new file mode 100644 index 000000000..95e20acd4 --- /dev/null +++ b/pkg/utils/cache/helper/model_schema.go @@ -0,0 +1,93 @@ +package helper + +import ( + "crypto/sha256" + "encoding/hex" + "strings" + "sync/atomic" + "time" + + "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/log" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" +) + +const DEFAULT_MODEL_SCHEMA_CACHE_TTL = 60 * time.Minute + +var modelSchemaCacheTTL atomic.Int64 + +func SetModelSchemaCacheTTL(ttl time.Duration) { + if ttl <= 0 { + ttl = DEFAULT_MODEL_SCHEMA_CACHE_TTL + } + modelSchemaCacheTTL.Store(int64(ttl)) +} + +func currentModelSchemaCacheTTL() time.Duration { + if ttl := modelSchemaCacheTTL.Load(); ttl > 0 { + return time.Duration(ttl) + } + return DEFAULT_MODEL_SCHEMA_CACHE_TTL +} + +type ModelSchemaCacheIdentity struct { + UniqueIdentifier plugin_entities.PluginUniqueIdentifier + ModelType string + Model string + CredentialType string + Credentials map[string]any +} + +// Remote and debugging plugins keep one identifier across code changes, so their schemas +// are never cached. +func (i ModelSchemaCacheIdentity) cacheable() bool { + return !i.UniqueIdentifier.RemoteLike() +} + +func (i ModelSchemaCacheIdentity) cacheKey() string { + digest := sha256.Sum256(parser.MarshalJsonBytes(map[string]any{ + "credential_type": i.CredentialType, + "credentials": i.Credentials, + })) + + return strings.Join( + []string{ + "model_schema", + i.UniqueIdentifier.String(), + i.ModelType, + i.Model, + hex.EncodeToString(digest[:]), + }, + ":", + ) +} + +// The key pins the plugin version and checksum, so an upgrade lands on a different key and +// a stale entry is unreachable. Nothing invalidates this cache. +func GetCachedModelSchema(identity ModelSchemaCacheIdentity) *model_entities.GetModelSchemasResponse { + if !identity.cacheable() { + return nil + } + + schema, err := cache.Get[model_entities.GetModelSchemasResponse](identity.cacheKey()) + if err != nil { + if err != cache.ErrNotFound { + log.Warn("failed to read model schema cache", "model", identity.Model, "error", err) + } + return nil + } + + return schema +} + +func StoreModelSchema(identity ModelSchemaCacheIdentity, schema model_entities.GetModelSchemasResponse) { + if !identity.cacheable() { + return + } + + if err := cache.Store(identity.cacheKey(), schema, currentModelSchemaCacheTTL()); err != nil { + log.Warn("failed to store model schema cache", "model", identity.Model, "error", err) + } +} diff --git a/pkg/utils/cache/helper/model_schema_test.go b/pkg/utils/cache/helper/model_schema_test.go new file mode 100644 index 000000000..4174f565a --- /dev/null +++ b/pkg/utils/cache/helper/model_schema_test.go @@ -0,0 +1,98 @@ +package helper + +import ( + "testing" + + miniredis "github.com/alicebob/miniredis/v2" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + SCHEMA_TEST_IDENTIFIER_V1 = "acme/vertex-connector:0.4.1@0123456789abcdef" + SCHEMA_TEST_IDENTIFIER_V2 = "acme/vertex-connector:0.5.0@fedcba9876543210" + SCHEMA_TEST_REMOTE = "5f8d1c2e-9a3b-4c6d-8e7f-0a1b2c3d4e5f/debug-plugin:0.0.1@aaaabbbbccccdddd" + SCHEMA_TEST_MODEL_TYPE = "llm" + SCHEMA_TEST_MODEL = "gemini-2.5-pro" +) + +func setupModelSchemaTest(t *testing.T) { + t.Helper() + + redisServer := miniredis.RunT(t) + require.NoError(t, cache.InitRedisClient(redisServer.Addr(), cache.RedisCredentials{}, false, 0, nil)) + t.Cleanup(func() { + _ = cache.Close() + }) +} + +func schemaTestIdentity(identifier string, credentials map[string]any) ModelSchemaCacheIdentity { + return ModelSchemaCacheIdentity{ + UniqueIdentifier: plugin_entities.PluginUniqueIdentifier(identifier), + ModelType: SCHEMA_TEST_MODEL_TYPE, + Model: SCHEMA_TEST_MODEL, + CredentialType: "api-key", + Credentials: credentials, + } +} + +func schemaWithLabel(label string) model_entities.GetModelSchemasResponse { + return model_entities.GetModelSchemasResponse{ + ModelSchema: &plugin_entities.ModelDeclaration{ + Model: label, + }, + } +} + +func TestModelSchemaCacheRoundTrip(t *testing.T) { + setupModelSchemaTest(t) + + identity := schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V1, map[string]any{"key": "secret-one"}) + assert.Nil(t, GetCachedModelSchema(identity), "nothing is cached before the first store") + + StoreModelSchema(identity, schemaWithLabel("from-cache")) + + cached := GetCachedModelSchema(identity) + require.NotNil(t, cached) + require.NotNil(t, cached.ModelSchema) + assert.Equal(t, "from-cache", cached.ModelSchema.Model) +} + +// An upgrade changes the unique identifier, so the old entry is unreachable rather than stale. +func TestModelSchemaCacheIsKeyedOnPluginVersion(t *testing.T) { + setupModelSchemaTest(t) + + credentials := map[string]any{"key": "secret-one"} + StoreModelSchema(schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V1, credentials), schemaWithLabel("old-version")) + + assert.Nil(t, + GetCachedModelSchema(schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V2, credentials)), + "an upgraded plugin must not read the previous version's schema", + ) +} + +func TestModelSchemaCacheIsKeyedOnCredentials(t *testing.T) { + setupModelSchemaTest(t) + + StoreModelSchema( + schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V1, map[string]any{"key": "secret-one"}), + schemaWithLabel("first-credentials"), + ) + + assert.Nil(t, + GetCachedModelSchema(schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V1, map[string]any{"key": "secret-two"})), + "a different credential set must not read the first one's schema", + ) +} + +func TestModelSchemaCacheSkipsRemotePlugins(t *testing.T) { + setupModelSchemaTest(t) + + identity := schemaTestIdentity(SCHEMA_TEST_REMOTE, map[string]any{"key": "secret-one"}) + StoreModelSchema(identity, schemaWithLabel("debug-build")) + + assert.Nil(t, GetCachedModelSchema(identity), "a debugging plugin reuses one identifier across builds") +}