From a28d61ed78ed9bb89da0f7d152a5c0a8781b3508 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 24 Feb 2026 09:50:23 +0800 Subject: [PATCH] Add environment variable resolution for agent settings - Introduce `resolveEnvStrings` function to handle `$ENV.XXX` references in the agent's DSL settings, ensuring that environment variables are correctly substituted in system and uses fields. - Enhance `Load` function to call `resolveEnvStrings` during the loading process, improving the configuration handling for agents. - Add comprehensive unit tests for `resolveEnvStrings` to validate the correct resolution of environment variables across various fields and scenarios, including handling of undefined variables and plain strings. - Ensure that the implementation maintains existing functionality while enhancing flexibility for environment-based configurations. --- agent/load.go | 37 ++++++++++++++++ agent/load_test.go | 107 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+) diff --git a/agent/load.go b/agent/load.go index 37208fa8..a5714cd1 100644 --- a/agent/load.go +++ b/agent/load.go @@ -6,6 +6,7 @@ import ( "github.com/yaoapp/gou/application" "github.com/yaoapp/gou/connector" + "github.com/yaoapp/gou/helper" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" @@ -46,6 +47,9 @@ func Load(cfg config.Config) error { setting.StoreSetting.MaxSize = 20 // default is 20 } + // Resolve $ENV.XXX references in system and uses fields + resolveEnvStrings(&setting) + // Default Assistant, Agent is the developer name, Mohe is the brand name of the assistant if setting.Uses == nil { setting.Uses = &types.Uses{Default: "mohe"} // Agent is the developer name, Mohe is the brand name of the assistant @@ -444,3 +448,36 @@ func defaultAssistant() (*assistant.Assistant, error) { } return assistant.Get(agentDSL.Uses.Default) } + +// resolveEnvStrings resolves $ENV.XXX references in agent.yml string fields. +// agent.yml is parsed via yaml.Unmarshal which does not handle $ENV substitution, +// unlike connector files which call helper.EnvString explicitly during Register. +func resolveEnvStrings(setting *types.DSL) { + if setting.System != nil { + setting.System.Default = helper.EnvString(setting.System.Default) + setting.System.Keyword = helper.EnvString(setting.System.Keyword) + setting.System.QueryDSL = helper.EnvString(setting.System.QueryDSL) + setting.System.Title = helper.EnvString(setting.System.Title) + setting.System.Prompt = helper.EnvString(setting.System.Prompt) + setting.System.RobotPrompt = helper.EnvString(setting.System.RobotPrompt) + setting.System.NeedSearch = helper.EnvString(setting.System.NeedSearch) + setting.System.Entity = helper.EnvString(setting.System.Entity) + } + + if setting.Uses != nil { + setting.Uses.Default = helper.EnvString(setting.Uses.Default) + setting.Uses.Title = helper.EnvString(setting.Uses.Title) + setting.Uses.Prompt = helper.EnvString(setting.Uses.Prompt) + setting.Uses.RobotPrompt = helper.EnvString(setting.Uses.RobotPrompt) + setting.Uses.Vision = helper.EnvString(setting.Uses.Vision) + setting.Uses.Audio = helper.EnvString(setting.Uses.Audio) + setting.Uses.Search = helper.EnvString(setting.Uses.Search) + setting.Uses.Fetch = helper.EnvString(setting.Uses.Fetch) + setting.Uses.Web = helper.EnvString(setting.Uses.Web) + setting.Uses.Keyword = helper.EnvString(setting.Uses.Keyword) + setting.Uses.QueryDSL = helper.EnvString(setting.Uses.QueryDSL) + setting.Uses.Rerank = helper.EnvString(setting.Uses.Rerank) + } + + setting.Cache = helper.EnvString(setting.Cache) +} diff --git a/agent/load_test.go b/agent/load_test.go index 3b219c00..653d10d8 100644 --- a/agent/load_test.go +++ b/agent/load_test.go @@ -8,6 +8,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/yaoapp/yao/agent/assistant" + "github.com/yaoapp/yao/agent/types" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/test" ) @@ -211,6 +212,112 @@ func TestGetGlobalPromptsWithDisableFlag(t *testing.T) { }) } +func TestResolveEnvStrings(t *testing.T) { + t.Setenv("TEST_CONNECTOR", "openai.gpt-5") + t.Setenv("TEST_ASSISTANT", "my-assistant") + t.Setenv("TEST_CACHE", "my-cache") + + t.Run("SystemFields", func(t *testing.T) { + setting := &types.DSL{ + System: &types.System{ + Default: "$ENV.TEST_CONNECTOR", + Keyword: "$ENV.TEST_CONNECTOR", + QueryDSL: "$ENV.TEST_CONNECTOR", + Title: "$ENV.TEST_CONNECTOR", + Prompt: "$ENV.TEST_CONNECTOR", + RobotPrompt: "$ENV.TEST_CONNECTOR", + NeedSearch: "$ENV.TEST_CONNECTOR", + Entity: "$ENV.TEST_CONNECTOR", + }, + } + resolveEnvStrings(setting) + + assert.Equal(t, "openai.gpt-5", setting.System.Default) + assert.Equal(t, "openai.gpt-5", setting.System.Keyword) + assert.Equal(t, "openai.gpt-5", setting.System.QueryDSL) + assert.Equal(t, "openai.gpt-5", setting.System.Title) + assert.Equal(t, "openai.gpt-5", setting.System.Prompt) + assert.Equal(t, "openai.gpt-5", setting.System.RobotPrompt) + assert.Equal(t, "openai.gpt-5", setting.System.NeedSearch) + assert.Equal(t, "openai.gpt-5", setting.System.Entity) + }) + + t.Run("UsesFields", func(t *testing.T) { + setting := &types.DSL{ + Uses: &types.Uses{ + Default: "$ENV.TEST_ASSISTANT", + Title: "$ENV.TEST_ASSISTANT", + Prompt: "$ENV.TEST_ASSISTANT", + RobotPrompt: "$ENV.TEST_ASSISTANT", + Vision: "$ENV.TEST_ASSISTANT", + Audio: "$ENV.TEST_ASSISTANT", + Search: "$ENV.TEST_ASSISTANT", + Fetch: "$ENV.TEST_ASSISTANT", + Web: "$ENV.TEST_ASSISTANT", + Keyword: "$ENV.TEST_ASSISTANT", + QueryDSL: "$ENV.TEST_ASSISTANT", + Rerank: "$ENV.TEST_ASSISTANT", + }, + } + resolveEnvStrings(setting) + + assert.Equal(t, "my-assistant", setting.Uses.Default) + assert.Equal(t, "my-assistant", setting.Uses.Title) + assert.Equal(t, "my-assistant", setting.Uses.Prompt) + assert.Equal(t, "my-assistant", setting.Uses.RobotPrompt) + assert.Equal(t, "my-assistant", setting.Uses.Vision) + assert.Equal(t, "my-assistant", setting.Uses.Audio) + assert.Equal(t, "my-assistant", setting.Uses.Search) + assert.Equal(t, "my-assistant", setting.Uses.Fetch) + assert.Equal(t, "my-assistant", setting.Uses.Web) + assert.Equal(t, "my-assistant", setting.Uses.Keyword) + assert.Equal(t, "my-assistant", setting.Uses.QueryDSL) + assert.Equal(t, "my-assistant", setting.Uses.Rerank) + }) + + t.Run("CacheField", func(t *testing.T) { + setting := &types.DSL{Cache: "$ENV.TEST_CACHE"} + resolveEnvStrings(setting) + assert.Equal(t, "my-cache", setting.Cache) + }) + + t.Run("PlainStringsUnchanged", func(t *testing.T) { + setting := &types.DSL{ + Cache: "plain-cache", + System: &types.System{ + Default: "openai.gpt-5", + }, + Uses: &types.Uses{ + Default: "mohe", + Title: "__yao.title", + }, + } + resolveEnvStrings(setting) + + assert.Equal(t, "plain-cache", setting.Cache) + assert.Equal(t, "openai.gpt-5", setting.System.Default) + assert.Equal(t, "mohe", setting.Uses.Default) + assert.Equal(t, "__yao.title", setting.Uses.Title) + }) + + t.Run("NilSystemAndUses", func(t *testing.T) { + setting := &types.DSL{Cache: "test"} + assert.NotPanics(t, func() { + resolveEnvStrings(setting) + }) + }) + + t.Run("UndefinedEnvReturnsEmpty", func(t *testing.T) { + setting := &types.DSL{ + System: &types.System{ + Default: "$ENV.UNDEFINED_VAR_12345", + }, + } + resolveEnvStrings(setting) + assert.Equal(t, "", setting.System.Default) + }) +} + func TestGlobalPromptsContent(t *testing.T) { prepare(t) defer test.Clean()