feat: support plugins and llm shield

sjy3 · sjy3 · commit a01efcb1d9d5 · 2026-01-29T10:51:48.000+08:00
diff --git a/.gitignore b/.gitignore
@@ -22,5 +22,6 @@
 go.work
 
 .idea/
+.trae
 agent/.env
 examples/quickstart/config.yaml
diff --git a/apps/a2a_app/app.go b/apps/a2a_app/app.go
@@ -96,6 +96,7 @@ func (a *agentkitA2AServerApp) SetupRouters(router *mux.Router, config *apps.Run
 			SessionService:  config.SessionService,
 			ArtifactService: config.ArtifactService,
 			MemoryService:   config.MemoryService,
+			PluginConfig:    config.PluginConfig,
 		},
 	})
 	reqHandler := a2asrv.NewHandler(executor, config.A2AOptions...)
diff --git a/apps/agentkit_server_app/app.go b/apps/agentkit_server_app/app.go
@@ -93,6 +93,7 @@ func (a *agentkitServerApp) SetupRouters(router *mux.Router, config *apps.RunCon
 		MemoryService:   config.MemoryService,
 		AgentLoader:     config.AgentLoader,
 		A2AOptions:      config.A2AOptions,
+		PluginConfig:    config.PluginConfig,
 	}
 
 	// setup webui routers
diff --git a/apps/basic_app.go b/apps/basic_app.go
@@ -24,6 +24,7 @@ import (
 	"google.golang.org/adk/agent"
 	"google.golang.org/adk/artifact"
 	"google.golang.org/adk/memory"
+	"google.golang.org/adk/runner"
 	"google.golang.org/adk/session"
 )
 
@@ -33,6 +34,7 @@ type RunConfig struct {
 	MemoryService   memory.Service
 	AgentLoader     agent.Loader
 	A2AOptions      []a2asrv.RequestHandlerOption
+	PluginConfig    runner.PluginConfig
 }
 
 type ApiConfig struct {
diff --git a/apps/simple_app/app.go b/apps/simple_app/app.go
@@ -78,6 +78,7 @@ func (app *agentkitSimpleApp) SetupRouters(router *mux.Router, config *apps.RunC
 		SessionService:  sessionService,
 		ArtifactService: config.ArtifactService,
 		MemoryService:   config.MemoryService,
+		PluginConfig:    config.PluginConfig,
 	})
 	if err != nil {
 		return fmt.Errorf("new runner error: %w", err)
diff --git a/auth/veauth/utils.go b/auth/veauth/utils.go
@@ -16,10 +16,13 @@ package veauth
 
 import (
 	"encoding/json"
+	"log"
 	"os"
 	"strings"
 
 	"github.com/volcengine/veadk-go/common"
+	"github.com/volcengine/veadk-go/configs"
+	"github.com/volcengine/veadk-go/utils"
 )
 
 type VeIAMCredential struct {
@@ -50,3 +53,20 @@ func RefreshAKSK(accessKey string, secretKey string) (VeIAMCredential, error) {
 	}
 	return GetCredentialFromVeFaaSIAM()
 }
+
+func GetAuthInfo() (ak, sk, sessionToken string) {
+	ak = utils.GetEnvWithDefault(common.VOLCENGINE_ACCESS_KEY, configs.GetGlobalConfig().Volcengine.AK)
+	sk = utils.GetEnvWithDefault(common.VOLCENGINE_SECRET_KEY, configs.GetGlobalConfig().Volcengine.SK)
+
+	if strings.TrimSpace(ak) == "" || strings.TrimSpace(sk) == "" {
+		iam, err := GetCredentialFromVeFaaSIAM()
+		if err != nil {
+			log.Printf("GetAuthInfo error: %s\n", err.Error())
+		} else {
+			ak = iam.AccessKeyID
+			sk = iam.SecretAccessKey
+			sessionToken = iam.SessionToken
+		}
+	}
+	return
+}
diff --git a/common/consts.go b/common/consts.go
@@ -89,3 +89,11 @@ const (
 	AGENTPILOT_API_KEY      = "AGENTPILOT_API_KEY"
 	AGENTPILOT_WORKSPACE_ID = "AGENTPILOT_WORKSPACE_ID"
 )
+
+// LLM Shield
+const (
+	TOOL_LLM_SHIELD_URL     = "TOOL_LLM_SHIELD_URL"
+	TOOL_LLM_SHIELD_REGION  = "TOOL_LLM_SHIELD_REGION"
+	TOOL_LLM_SHIELD_APP_ID  = "TOOL_LLM_SHIELD_APP_ID"
+	TOOL_LLM_SHIELD_API_KEY = "TOOL_LLM_SHIELD_API_KEY"
+)
diff --git a/common/defaults.go b/common/defaults.go
@@ -69,6 +69,12 @@ const (
 	DEFAULT_AGENTKIT_TOOL_SERVICE_CODE = "agentkit"
 )
 
+// prompt pilot
 const (
 	DEFAULT_AGENTPILOT_API_URL = "https://prompt-pilot.cn-beijing.volces.com"
 )
+
+// LLM Shield
+const (
+	DEFAULT_LLM_SHIELD_REGION = "cn-beijing"
+)
diff --git a/configs/configs.go b/configs/configs.go
@@ -69,6 +69,7 @@ func SetupVeADKConfig() error {
 		Tool: &BuiltinToolConfigs{
 			MCPRouter: &MCPRouter{},
 			RunCode:   &RunCode{},
+			LLMShield: &LLMShield{},
 		},
 		PromptPilot: &PromptPilotConfig{},
 		TlsConfig:   &TLSConfig{},
diff --git a/configs/tool.go b/configs/tool.go
@@ -22,11 +22,13 @@ import (
 type BuiltinToolConfigs struct {
 	MCPRouter *MCPRouter `yaml:"mcp_router"`
 	RunCode   *RunCode   `yaml:"run_code"`
+	LLMShield *LLMShield `yaml:"llm_shield"`
 }
 
 func (b *BuiltinToolConfigs) MapEnvToConfig() {
 	b.MCPRouter.MapEnvToConfig()
 	b.RunCode.MapEnvToConfig()
+	b.LLMShield.MapEnvToConfig()
 }
 
 type MCPRouter struct {
@@ -52,3 +54,17 @@ func (r *RunCode) MapEnvToConfig() {
 	r.ServiceCode = utils.GetEnvWithDefault(common.AGENTKIT_TOOL_SERVICE_CODE)
 	r.Region = utils.GetEnvWithDefault(common.AGENTKIT_TOOL_REGION)
 }
+
+type LLMShield struct {
+	Url    string `yaml:"url"`
+	Region string `yaml:"region"`
+	AppId  string `yaml:"app_id"`
+	ApiKey string `yaml:"api_key"`
+}
+
+func (r *LLMShield) MapEnvToConfig() {
+	r.Url = utils.GetEnvWithDefault(common.TOOL_LLM_SHIELD_URL)
+	r.Region = utils.GetEnvWithDefault(common.TOOL_LLM_SHIELD_REGION)
+	r.ApiKey = utils.GetEnvWithDefault(common.TOOL_LLM_SHIELD_API_KEY)
+	r.AppId = utils.GetEnvWithDefault(common.TOOL_LLM_SHIELD_APP_ID)
+}
diff --git a/examples/plugins/agent.go b/examples/plugins/agent.go
@@ -0,0 +1,146 @@
+// Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"fmt"
+
+	veagent "github.com/volcengine/veadk-go/agent/llmagent"
+	"github.com/volcengine/veadk-go/agent/workflowagents/sequentialagent"
+	"github.com/volcengine/veadk-go/apps"
+	"github.com/volcengine/veadk-go/apps/agentkit_server_app"
+	"github.com/volcengine/veadk-go/tool/builtin_tools"
+	"github.com/volcengine/veadk-go/tool/builtin_tools/web_search"
+	"github.com/volcengine/veadk-go/utils"
+	"google.golang.org/adk/agent"
+	"google.golang.org/adk/agent/llmagent"
+	"google.golang.org/adk/model"
+	"google.golang.org/adk/plugin"
+	"google.golang.org/adk/runner"
+	"google.golang.org/adk/session"
+	"google.golang.org/adk/tool"
+)
+
+func main() {
+	ctx := context.Background()
+
+	webSearch, err := web_search.NewWebSearchTool(&web_search.Config{})
+	if err != nil {
+		fmt.Printf("NewWebSearchTool failed: %v", err)
+		return
+	}
+
+	greetingAgent, err := veagent.New(&veagent.Config{
+		Config: llmagent.Config{
+			Name:        "greeting_agent",
+			Description: "A friendly agent that greets the user.",
+			Instruction: "Greet the user warmly.",
+			Tools: []tool.Tool{
+				webSearch,
+			},
+		},
+		ModelExtraConfig: map[string]any{
+			"extra_body": map[string]any{
+				"thinking": map[string]string{
+					"type": "disabled",
+				},
+			},
+		},
+	})
+	if err != nil {
+		fmt.Printf("NewLLMAgent greetingAgent failed: %v", err)
+		return
+	}
+
+	goodbyeAgent, err := veagent.New(&veagent.Config{
+		Config: llmagent.Config{
+			Name:        "goodbye_agent",
+			Description: "A polite agent that says goodbye to the user.",
+			Instruction: "Directly return goodbye",
+		},
+		ModelExtraConfig: map[string]any{
+			"extra_body": map[string]any{
+				"thinking": map[string]string{
+					"type": "disabled",
+				},
+			},
+		},
+	})
+	if err != nil {
+		fmt.Printf("NewLLMAgent goodbyeAgent failed: %v", err)
+		return
+	}
+
+	rootAgent, err := sequentialagent.New(sequentialagent.Config{
+		AgentConfig: agent.Config{
+			Name:        "veAgent",
+			SubAgents:   []agent.Agent{greetingAgent, goodbyeAgent},
+			Description: "Executes a sequence of greeting and goodbye.",
+		},
+	})
+
+	if err != nil {
+		fmt.Printf("NewSequentialAgent failed: %v", err)
+		return
+	}
+
+	app := agentkit_server_app.NewAgentkitServerApp(apps.DefaultApiConfig())
+
+	err = app.Run(ctx, &apps.RunConfig{
+		AgentLoader:    agent.NewSingleLoader(rootAgent),
+		SessionService: session.InMemoryService(),
+		PluginConfig: runner.PluginConfig{
+			Plugins: []*plugin.Plugin{
+				//NewTestPlugins(),
+				utils.Must(builtin_tools.NewLLMShieldPlugins()),
+			},
+		},
+	})
+	if err != nil {
+		fmt.Printf("Run failed: %v", err)
+	}
+}
+
+func beforeModelCallBack(ctx agent.CallbackContext, llmRequest *model.LLMRequest) (*model.LLMResponse, error) {
+	fmt.Printf("%s BeforeModelCallBack called\n", ctx.AgentName())
+	return nil, nil
+}
+
+func afterModelCallBack(ctx agent.CallbackContext, llmResponse *model.LLMResponse, llmResponseError error) (*model.LLMResponse, error) {
+	fmt.Printf("%s afterModelCallback called\n", ctx.AgentName())
+	return nil, nil
+}
+
+func beforeToolCallback(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) {
+	fmt.Printf("%s beforeToolCallBack called\n", tool.Name())
+	return nil, nil
+}
+
+func afterToolCallback(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) {
+	fmt.Printf("%s afterToolCallback called\n", tool.Name())
+	return nil, nil
+}
+
+func NewTestPlugins() *plugin.Plugin {
+	plugins, _ := plugin.New(plugin.Config{
+		Name:                "llm_shield_test",
+		BeforeModelCallback: beforeModelCallBack,
+		AfterModelCallback:  afterModelCallBack,
+		BeforeToolCallback:  beforeToolCallback,
+		AfterToolCallback:   afterToolCallback,
+	})
+	return plugins
+}
diff --git a/go.mod b/go.mod
@@ -15,7 +15,7 @@ require (
 	github.com/volcengine/volcengine-go-sdk v1.1.53
 	go.uber.org/zap v1.27.1
 	golang.org/x/oauth2 v0.32.0
-	google.golang.org/adk v0.3.1-0.20260113130926-012d380e4056
+	google.golang.org/adk v0.3.1-0.20260123125504-aec89487da29
 	google.golang.org/genai v1.40.0
 	gopkg.in/go-playground/validator.v8 v8.18.2
 	gopkg.in/yaml.v3 v3.0.1
@@ -81,4 +81,4 @@ require (
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	rsc.io/omap v1.2.0 // indirect
 	rsc.io/ordered v1.1.1 // indirect
-)
+)
diff --git a/go.sum b/go.sum
@@ -220,8 +220,8 @@ golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
 gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
-google.golang.org/adk v0.3.1-0.20260113130926-012d380e4056 h1:SfDR+nm6VdVjSsvKbHGAV/B6csmMge24nqQBSbJi5Go=
-google.golang.org/adk v0.3.1-0.20260113130926-012d380e4056/go.mod h1:iE1Kgc8JtYHiNxfdLa9dxcV4DqTn0D8q4eqhBi012Ak=
+google.golang.org/adk v0.3.1-0.20260123125504-aec89487da29 h1:xYzqSSpn/WhNVyiieXnvYLGKPRJk4cJFgImrhsA9Pi4=
+google.golang.org/adk v0.3.1-0.20260123125504-aec89487da29/go.mod h1:jVeb7Ir53+3XKTncdY7k3pVdPneKcm5+60sXpxHQnao=
 google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
 google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
 google.golang.org/genai v1.40.0 h1:kYxyQSH+vsib8dvsgyLJzsVEIv5k3ZmHJyVqdvGncmc=
diff --git a/tool/builtin_tools/llm_shield.go b/tool/builtin_tools/llm_shield.go
diff --git a/tool/builtin_tools/llm_shield_test.go b/tool/builtin_tools/llm_shield_test.go

Original file line number	Diff line number	Diff line change
`@@ -93,6 +93,7 @@ func (a agentkitServerApp) SetupRouters(router mux.Router, config *apps.RunCon`
`93`	`93`	`MemoryService: config.MemoryService,`
`94`	`94`	`AgentLoader: config.AgentLoader,`
`95`	`95`	`A2AOptions: config.A2AOptions,`
	`96`	`+ PluginConfig: config.PluginConfig,`
`96`	`97`	`}`
`97`	`98`
`98`	`99`	`// setup webui routers`
Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@ import (`
`24`	`24`	`"google.golang.org/adk/agent"`
`25`	`25`	`"google.golang.org/adk/artifact"`
`26`	`26`	`"google.golang.org/adk/memory"`
	`27`	`+ "google.golang.org/adk/runner"`
`27`	`28`	`"google.golang.org/adk/session"`
`28`	`29`	`)`
`29`	`30`
`@@ -33,6 +34,7 @@ type RunConfig struct {`
`33`	`34`	`MemoryService memory.Service`
`34`	`35`	`AgentLoader agent.Loader`
`35`	`36`	`A2AOptions []a2asrv.RequestHandlerOption`
	`37`	`+ PluginConfig runner.PluginConfig`
`36`	`38`	`}`
`37`	`39`
`38`	`40`	`type ApiConfig struct {`
Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,12 @@ const (`
`69`	`69`	`DEFAULT_AGENTKIT_TOOL_SERVICE_CODE = "agentkit"`
`70`	`70`	`)`
`71`	`71`
	`72`	`+// prompt pilot`
`72`	`73`	`const (`
`73`	`74`	`DEFAULT_AGENTPILOT_API_URL = "https://prompt-pilot.cn-beijing.volces.com"`
`74`	`75`	`)`
	`76`	`+`
	`77`	`+// LLM Shield`
	`78`	`+const (`
	`79`	`+ DEFAULT_LLM_SHIELD_REGION = "cn-beijing"`
	`80`	`+)`