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.
This commit is contained in:
Max
parent
a38d115465
commit
a28d61ed78
2 changed files with 144 additions and 0 deletions
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue