feat(llm): refactor audio capabilities and enhance role management

- Replaced 'Voice' with 'Audio' in the system configuration and related tests to better reflect functionality.
- Introduced new methods for role management in the llmprovider, allowing for dynamic retrieval of roles based on user and team context.
- Updated the OpenAPI settings to support new role management endpoints and capabilities.
- Enhanced the handling of API keys in provider management, allowing for optional plain-text retrieval.
This commit is contained in:
Max 2026-04-30 12:33:05 +08:00
parent a58cac6d5c
commit 6f78f6066b
24 changed files with 2348 additions and 94 deletions

1
.gitignore vendored
View file

@ -84,3 +84,4 @@ sandbox/v2/PID-KILL-UPGRADE.md
sandbox/v2/*.md
POSTGRESQL_COMPAT.md
openapi/setting/*.md
agent/docs/design/*.md

View file

@ -15,6 +15,7 @@ import (
"github.com/yaoapp/yao/agent/output/message"
agentsandbox "github.com/yaoapp/yao/agent/sandbox"
sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types"
"github.com/yaoapp/yao/llmprovider"
infraV2 "github.com/yaoapp/yao/sandbox/v2"
)
@ -628,7 +629,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
}
// GetConnector get the connector object, capabilities, and error with priority:
// opts.Connector > ast.Connector > defaultConnector (fallback)
// opts.Connector > ast.Connector > GetRoleBy("default", identity) > GetRole("default") > error
// Note: opts.Connector may be set by Create hook's applyOptionsAdjustments
// Returns: (connector, capabilities, error)
func (ast *Assistant) GetConnector(ctx *context.Context, opts ...*context.Options) (connector.Connector, *goullm.Capabilities, error) {
@ -637,6 +638,21 @@ func (ast *Assistant) GetConnector(ctx *context.Context, opts ...*context.Option
connectorID = opts[0].Connector
}
// Fallback to unified role resolution via llmprovider (team > user > system)
if connectorID == "" && llmprovider.Global != nil {
if ctx != nil && ctx.Authorized != nil {
if cid, err := llmprovider.Global.GetRoleBy("default", ctx.Authorized); err == nil {
connectorID = cid
}
}
if connectorID == "" {
if cid, err := llmprovider.Global.GetRole("default"); err == nil {
connectorID = cid
}
}
}
// Legacy fallback
if connectorID == "" {
connectorID = defaultConnector
}

View file

@ -161,9 +161,6 @@ func (ast *Assistant) Validate() error {
if ast.Name == "" {
return fmt.Errorf("name is required")
}
if ast.Connector == "" {
return fmt.Errorf("connector is required")
}
return nil
}

View file

@ -881,20 +881,11 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
// Init init the assistant
// Choose the connector and initialize the assistant
func (ast *Assistant) initialize() error {
conn := defaultConnector
if ast.Connector != "" {
conn = ast.Connector
}
ast.Connector = conn
// Register scripts as process handlers
if len(ast.Scripts) > 0 {
if err := ast.RegisterScripts(); err != nil {
return fmt.Errorf("failed to register scripts: %w", err)
}
}
return nil
}

View file

@ -12,6 +12,7 @@ import (
"github.com/yaoapp/yao/agent/i18n"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/data"
"github.com/yaoapp/yao/llmprovider"
"gopkg.in/yaml.v3"
)
@ -43,7 +44,7 @@ type SystemConfig struct {
NeedSearch string // Connector for __yao.needsearch agent
Entity string // Connector for __yao.entity agent
Vision string // Connector for vision capabilities
Voice string // Connector for voice/STT capabilities
Audio string // Connector for audio/STT capabilities
}
// systemConfig holds the system agents configuration (global variable like others in load.go)
@ -208,7 +209,7 @@ func loadSystemAgent(id, pathPrefix string) (*Assistant, error) {
}
// resolveSystemConnector resolves the connector for a system agent
// Priority: specific agent config > system.default > defaultConnector > fallback to first capable connector
// Priority: specific agent config > system.default > llmprovider role > defaultConnector > fallback
func resolveSystemConnector(agentID string) string {
// Try specific agent config first
if systemConfig != nil {
@ -245,9 +246,9 @@ func resolveSystemConnector(agentID string) string {
if systemConfig.Vision != "" {
return systemConfig.Vision
}
case "__yao.voice":
if systemConfig.Voice != "" {
return systemConfig.Voice
case "__yao.audio":
if systemConfig.Audio != "" {
return systemConfig.Audio
}
}
@ -257,6 +258,17 @@ func resolveSystemConnector(agentID string) string {
}
}
// Try unified role resolution: strip __yao. prefix as role name
if llmprovider.Global != nil {
role := strings.TrimPrefix(agentID, "__yao.")
if cid, err := llmprovider.Global.GetRole(role); err == nil && cid != "" {
return cid
}
if cid, err := llmprovider.Global.GetRole("default"); err == nil && cid != "" {
return cid
}
}
// Try global default connector
if defaultConnector != "" {
return defaultConnector
@ -267,7 +279,7 @@ func resolveSystemConnector(agentID string) string {
}
// GetVisionConnector returns the connector for vision capabilities.
// Priority: system.vision > system.default > defaultConnector > findCapableConnector
// Priority: system.vision > system.default > llmprovider GetRole("vision") > defaultConnector > findCapableConnector
func GetVisionConnector() string {
if systemConfig != nil {
if systemConfig.Vision != "" {
@ -277,23 +289,33 @@ func GetVisionConnector() string {
return systemConfig.Default
}
}
if llmprovider.Global != nil {
if cid, err := llmprovider.Global.GetRole("vision"); err == nil && cid != "" {
return cid
}
}
if defaultConnector != "" {
return defaultConnector
}
return findCapableConnector()
}
// GetVoiceConnector returns the connector for voice/STT capabilities.
// Priority: system.voice > system.default > defaultConnector > findCapableConnector
func GetVoiceConnector() string {
// GetAudioConnector returns the connector for audio/STT capabilities.
// Priority: system.audio > system.default > llmprovider GetRole("audio") > defaultConnector > findCapableConnector
func GetAudioConnector() string {
if systemConfig != nil {
if systemConfig.Voice != "" {
return systemConfig.Voice
if systemConfig.Audio != "" {
return systemConfig.Audio
}
if systemConfig.Default != "" {
return systemConfig.Default
}
}
if llmprovider.Global != nil {
if cid, err := llmprovider.Global.GetRole("audio"); err == nil && cid != "" {
return cid
}
}
if defaultConnector != "" {
return defaultConnector
}

View file

@ -35,12 +35,48 @@ func GetCapabilitiesFromConn(conn connector.Connector) *goullm.Capabilities {
if capabilities, ok := caps.(goullm.Capabilities); ok {
return &capabilities
}
if capsMap, ok := caps.(map[string]interface{}); ok {
return capabilitiesFromMap(capsMap)
}
}
}
return getDefaultCapabilities()
}
// capabilitiesFromMap converts a JSON-deserialized map into goullm.Capabilities.
func capabilitiesFromMap(m map[string]interface{}) *goullm.Capabilities {
caps := getDefaultCapabilities()
if v, ok := m["streaming"].(bool); ok {
caps.Streaming = v
}
if v, ok := m["tool_calls"].(bool); ok {
caps.ToolCalls = v
}
if v, ok := m["vision"]; ok {
caps.Vision = v
}
if v, ok := m["audio"].(bool); ok {
caps.Audio = v
}
if v, ok := m["stt"].(bool); ok {
caps.STT = v
}
if v, ok := m["reasoning"].(bool); ok {
caps.Reasoning = v
}
if v, ok := m["json"].(bool); ok {
caps.JSON = v
}
if v, ok := m["multimodal"].(bool); ok {
caps.Multimodal = v
}
if v, ok := m["temperature_adjustable"].(bool); ok {
caps.TemperatureAdjustable = v
}
return caps
}
// getDefaultCapabilities returns minimal default capabilities
func getDefaultCapabilities() *goullm.Capabilities {
return &goullm.Capabilities{

View file

@ -19,6 +19,7 @@ import (
"github.com/yaoapp/yao/agent/store/xun"
"github.com/yaoapp/yao/agent/types"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/llmprovider"
)
var agentDSL *types.DSL
@ -234,7 +235,7 @@ func initAssistant() error {
NeedSearch: agentDSL.System.NeedSearch,
Entity: agentDSL.System.Entity,
Vision: agentDSL.System.Vision,
Voice: agentDSL.System.Voice,
Audio: agentDSL.System.Audio,
})
}
@ -453,6 +454,22 @@ func GetSearchConfig() *searchTypes.Config {
return agentDSL.Search
}
// SyncLLMDefaults writes the agent.yml system role defaults into setting.Global.
// Must be called after both llmprovider.Init() and setting.Init() have completed.
func SyncLLMDefaults() error {
if agentDSL == nil || agentDSL.System == nil {
return nil
}
if llmprovider.Global == nil {
return fmt.Errorf("llmprovider.Global not initialized")
}
roles := buildSystemRoles(agentDSL.System)
if len(roles) == 0 {
return nil
}
return llmprovider.Global.SetDefaults(roles)
}
// defaultAssistant get the default assistant
func defaultAssistant() (*assistant.Assistant, error) {
if agentDSL.Uses == nil || agentDSL.Uses.Default == "" {
@ -461,6 +478,28 @@ func defaultAssistant() (*assistant.Assistant, error) {
return assistant.Get(agentDSL.Uses.Default)
}
// buildSystemRoles converts the System config block into a role→connectorID map
// for llmprovider.SetDefaults.
func buildSystemRoles(sys *types.System) map[string]string {
roles := make(map[string]string)
add := func(role, cid string) {
if cid != "" {
roles[role] = cid
}
}
add("default", sys.Default)
add("keyword", sys.Keyword)
add("querydsl", sys.QueryDSL)
add("title", sys.Title)
add("prompt", sys.Prompt)
add("robot_prompt", sys.RobotPrompt)
add("needsearch", sys.NeedSearch)
add("entity", sys.Entity)
add("vision", sys.Vision)
add("audio", sys.Audio)
return roles
}
// 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.
@ -475,7 +514,7 @@ func resolveEnvStrings(setting *types.DSL) {
setting.System.NeedSearch = helper.EnvString(setting.System.NeedSearch)
setting.System.Entity = helper.EnvString(setting.System.Entity)
setting.System.Vision = helper.EnvString(setting.System.Vision)
setting.System.Voice = helper.EnvString(setting.System.Voice)
setting.System.Audio = helper.EnvString(setting.System.Audio)
}
if setting.Uses != nil {

View file

@ -229,7 +229,7 @@ func TestResolveEnvStrings(t *testing.T) {
NeedSearch: "$ENV.TEST_CONNECTOR",
Entity: "$ENV.TEST_CONNECTOR",
Vision: "$ENV.TEST_CONNECTOR",
Voice: "$ENV.TEST_CONNECTOR",
Audio: "$ENV.TEST_CONNECTOR",
},
}
resolveEnvStrings(setting)
@ -243,24 +243,24 @@ func TestResolveEnvStrings(t *testing.T) {
assert.Equal(t, "openai.gpt-5", setting.System.NeedSearch)
assert.Equal(t, "openai.gpt-5", setting.System.Entity)
assert.Equal(t, "openai.gpt-5", setting.System.Vision)
assert.Equal(t, "openai.gpt-5", setting.System.Voice)
assert.Equal(t, "openai.gpt-5", setting.System.Audio)
})
t.Run("SystemVisionVoiceSeparateEnv", func(t *testing.T) {
t.Run("SystemVisionAudioSeparateEnv", func(t *testing.T) {
t.Setenv("TEST_VISION_CONN", "openai.gpt-4o")
t.Setenv("TEST_VOICE_CONN", "whisper-1")
t.Setenv("TEST_AUDIO_CONN", "whisper-1")
setting := &types.DSL{
System: &types.System{
Default: "$ENV.TEST_CONNECTOR",
Vision: "$ENV.TEST_VISION_CONN",
Voice: "$ENV.TEST_VOICE_CONN",
Audio: "$ENV.TEST_AUDIO_CONN",
},
}
resolveEnvStrings(setting)
assert.Equal(t, "openai.gpt-5", setting.System.Default)
assert.Equal(t, "openai.gpt-4o", setting.System.Vision)
assert.Equal(t, "whisper-1", setting.System.Voice)
assert.Equal(t, "whisper-1", setting.System.Audio)
})
t.Run("UsesFields", func(t *testing.T) {

View file

@ -102,7 +102,7 @@ type System struct {
NeedSearch string `json:"needsearch,omitempty" yaml:"needsearch,omitempty"` // Connector for __yao.needsearch agent
Entity string `json:"entity,omitempty" yaml:"entity,omitempty"` // Connector for __yao.entity agent
Vision string `json:"vision,omitempty" yaml:"vision,omitempty"` // Connector for vision capabilities
Voice string `json:"voice,omitempty" yaml:"voice,omitempty"` // Connector for voice/STT capabilities
Audio string `json:"audio,omitempty" yaml:"audio,omitempty"` // Connector for audio/STT capabilities
}
// Mention Structure

View file

@ -448,6 +448,11 @@ func Load(cfg config.Config, options LoadOption, progressCallback ...func(string
warnings = append(warnings, Warning{Widget: "Setting Registry", Error: err})
}
// Sync agent.yml system defaults into setting.Global (must run after llmprovider + setting init)
if err := agent.SyncLLMDefaults(); err != nil {
warnings = append(warnings, Warning{Widget: "LLM Defaults Sync", Error: err})
}
for name, hook := range LoadHooks {
err = hook(cfg)
if err != nil {
@ -721,6 +726,11 @@ func Reload(cfg config.Config, options LoadOption) (err error) {
}
}
// Sync agent.yml system defaults into setting.Global (must run after llmprovider + setting init)
if err := agent.SyncLLMDefaults(); err != nil {
printErr(cfg.Mode, "LLM Defaults Sync", err)
}
// Load OpenAPI
_, err = openapi.Load(cfg)
if err != nil {

364
llmprovider/models.go Normal file
View file

@ -0,0 +1,364 @@
package llmprovider
import (
"fmt"
"strings"
"github.com/yaoapp/gou/connector"
goullm "github.com/yaoapp/gou/llm"
)
// ---------------------------------------------------------------------------
// GetModel — by connectorID
// ---------------------------------------------------------------------------
// GetModel returns the runtime connector for a given connectorID.
// Lookup order:
// 1. connector.Select (already registered in runtime)
// 2. Model-level ID with ":" separator (e.g. "t123.openai:gpt-4o")
// 3. r.Get by store Key (works when connectorID == Key, e.g. builtin)
// 4. r.GetByConnectorID (linear scan by ConnectorID field, for dynamic providers)
func (r *Registry) GetModel(connectorID string) (connector.Connector, error) {
if conn, err := connector.Select(connectorID); err == nil {
return conn, nil
}
// Path 2: model-level ID "providerCID:modelID"
if parts := strings.SplitN(connectorID, ":", 2); len(parts) == 2 {
return r.getModelConnector(parts[0], parts[1])
}
// Path 3: try by Key (fast, works for builtin where Key == ConnectorID)
if p, err := r.Get(connectorID, true); err == nil {
if eerr := ensureConnector(p); eerr != nil {
return nil, fmt.Errorf("model %q ensure connector: %w", connectorID, eerr)
}
cid := p.ConnectorID
if cid == "" {
cid = connectorID
}
return connector.Select(cid)
}
// Path 4: reverse lookup by ConnectorID field (dynamic providers where Key != ConnectorID)
p, err := r.GetByConnectorID(connectorID, true)
if err != nil {
return nil, fmt.Errorf("model %q not found", connectorID)
}
cid := p.ConnectorID
if cid == "" {
cid = connectorID
}
return connector.Select(cid)
}
// getModelConnector finds a provider by connectorID, locates the model, and
// ensures a per-model connector is registered in the runtime.
func (r *Registry) getModelConnector(providerCID, modelID string) (connector.Connector, error) {
p, err := r.GetByConnectorID(providerCID, true)
if err != nil {
if p2, err2 := r.Get(providerCID, true); err2 == nil {
p = p2
} else {
return nil, fmt.Errorf("provider %q not found for model %q", providerCID, modelID)
}
}
var model *ModelInfo
for i, m := range p.Models {
if m.ID == modelID {
model = &p.Models[i]
break
}
}
if model == nil {
return nil, fmt.Errorf("model %q not found in provider %q", modelID, providerCID)
}
if err := ensureModelConnector(p, model); err != nil {
return nil, err
}
cid := providerCID + ":" + modelID
return connector.Select(cid)
}
// ---------------------------------------------------------------------------
// GetRoleModel — role → connector
// ---------------------------------------------------------------------------
// GetRoleModel returns the connector for a role at system scope.
func (r *Registry) GetRoleModel(role string) (connector.Connector, error) {
cid, err := r.GetRole(role)
if err != nil {
return nil, err
}
return r.GetModel(cid)
}
// GetRoleModelByUser returns the connector for a role, merged user > system.
func (r *Registry) GetRoleModelByUser(role, userID string) (connector.Connector, error) {
cid, err := r.GetRoleByUser(role, userID)
if err != nil {
return nil, err
}
return r.GetModel(cid)
}
// GetRoleModelByTeam returns the connector for a role, merged team > system.
func (r *Registry) GetRoleModelByTeam(role, teamID string) (connector.Connector, error) {
cid, err := r.GetRoleByTeam(role, teamID)
if err != nil {
return nil, err
}
return r.GetModel(cid)
}
// ---------------------------------------------------------------------------
// Built-in role shortcuts
// ---------------------------------------------------------------------------
func (r *Registry) GetDefaultModel() (connector.Connector, error) { return r.GetRoleModel("default") }
func (r *Registry) GetDefaultModelByUser(userID string) (connector.Connector, error) {
return r.GetRoleModelByUser("default", userID)
}
func (r *Registry) GetDefaultModelByTeam(teamID string) (connector.Connector, error) {
return r.GetRoleModelByTeam("default", teamID)
}
func (r *Registry) GetVisionModel() (connector.Connector, error) { return r.GetRoleModel("vision") }
func (r *Registry) GetVisionModelByUser(userID string) (connector.Connector, error) {
return r.GetRoleModelByUser("vision", userID)
}
func (r *Registry) GetVisionModelByTeam(teamID string) (connector.Connector, error) {
return r.GetRoleModelByTeam("vision", teamID)
}
func (r *Registry) GetAudioModel() (connector.Connector, error) { return r.GetRoleModel("audio") }
func (r *Registry) GetAudioModelByUser(userID string) (connector.Connector, error) {
return r.GetRoleModelByUser("audio", userID)
}
func (r *Registry) GetAudioModelByTeam(teamID string) (connector.Connector, error) {
return r.GetRoleModelByTeam("audio", teamID)
}
func (r *Registry) GetEmbeddingModel() (connector.Connector, error) {
return r.GetRoleModel("embedding")
}
func (r *Registry) GetEmbeddingModelByUser(userID string) (connector.Connector, error) {
return r.GetRoleModelByUser("embedding", userID)
}
func (r *Registry) GetEmbeddingModelByTeam(teamID string) (connector.Connector, error) {
return r.GetRoleModelByTeam("embedding", teamID)
}
// ---------------------------------------------------------------------------
// Capabilities
// ---------------------------------------------------------------------------
// GetCapabilities returns capabilities for a connector by connectorID.
func (r *Registry) GetCapabilities(connectorID string) (*goullm.Capabilities, error) {
conn, err := r.GetModel(connectorID)
if err != nil {
return nil, err
}
return capabilitiesFromConn(conn), nil
}
// GetRoleCapabilities returns capabilities for a role at system scope.
func (r *Registry) GetRoleCapabilities(role string) (*goullm.Capabilities, error) {
conn, err := r.GetRoleModel(role)
if err != nil {
return nil, err
}
return capabilitiesFromConn(conn), nil
}
// GetRoleCapabilitiesByUser returns capabilities for a role, merged user > system.
func (r *Registry) GetRoleCapabilitiesByUser(role, userID string) (*goullm.Capabilities, error) {
conn, err := r.GetRoleModelByUser(role, userID)
if err != nil {
return nil, err
}
return capabilitiesFromConn(conn), nil
}
// GetRoleCapabilitiesByTeam returns capabilities for a role, merged team > system.
func (r *Registry) GetRoleCapabilitiesByTeam(role, teamID string) (*goullm.Capabilities, error) {
conn, err := r.GetRoleModelByTeam(role, teamID)
if err != nil {
return nil, err
}
return capabilitiesFromConn(conn), nil
}
// ---------------------------------------------------------------------------
// ListModels
// ---------------------------------------------------------------------------
// ListModels returns all enabled models as []connector.Option (system scope, no owner filter).
func (r *Registry) ListModels() []connector.Option {
return r.listModels(nil)
}
// ListModelsByUser returns builtin + user-owned dynamic models.
func (r *Registry) ListModelsByUser(userID string) []connector.Option {
return r.listModels(&ProviderOwner{Type: "user", UserID: userID})
}
// ListModelsByTeam returns builtin + team-owned dynamic models.
func (r *Registry) ListModelsByTeam(teamID string) []connector.Option {
return r.listModels(&ProviderOwner{Type: "team", TeamID: teamID})
}
// ListModelsBy returns models scoped to the caller's identity (team > user).
func (r *Registry) ListModelsBy(id Identity) []connector.Option {
if id.GetTeamID() != "" {
return r.ListModelsByTeam(id.GetTeamID())
}
return r.ListModelsByUser(id.GetUserID())
}
// ---------------------------------------------------------------------------
// By — Identity-scoped convenience methods
// ---------------------------------------------------------------------------
// GetRoleBy returns the connectorID for a role, scoped by identity.
func (r *Registry) GetRoleModelBy(role string, id Identity) (connector.Connector, error) {
if id.GetTeamID() != "" {
return r.GetRoleModelByTeam(role, id.GetTeamID())
}
return r.GetRoleModelByUser(role, id.GetUserID())
}
func (r *Registry) GetDefaultModelBy(id Identity) (connector.Connector, error) {
return r.GetRoleModelBy("default", id)
}
func (r *Registry) GetVisionModelBy(id Identity) (connector.Connector, error) {
return r.GetRoleModelBy("vision", id)
}
func (r *Registry) GetAudioModelBy(id Identity) (connector.Connector, error) {
return r.GetRoleModelBy("audio", id)
}
func (r *Registry) GetEmbeddingModelBy(id Identity) (connector.Connector, error) {
return r.GetRoleModelBy("embedding", id)
}
func (r *Registry) GetRoleCapabilitiesBy(role string, id Identity) (*goullm.Capabilities, error) {
conn, err := r.GetRoleModelBy(role, id)
if err != nil {
return nil, err
}
return capabilitiesFromConn(conn), nil
}
// ---------------------------------------------------------------------------
// internal
// ---------------------------------------------------------------------------
// listModels returns enabled models. When owner is non-nil, returns builtin
// providers plus dynamic providers belonging to that owner.
// Builtin providers are one-connector-per-model; dynamic providers are expanded
// to per-model options here.
func (r *Registry) listModels(owner *ProviderOwner) []connector.Option {
enabled := true
providers, err := r.List(&ProviderFilter{
Source: ProviderSourceAll,
Enabled: &enabled,
}, true)
if err != nil {
return nil
}
var result []connector.Option
for _, p := range providers {
if owner != nil && p.Source == ProviderSourceDynamic {
if !ownerMatch(&p.Owner, owner) {
continue
}
}
if p.Source == ProviderSourceDynamic && len(p.Models) > 0 {
for _, m := range p.Models {
if !m.Enabled {
continue
}
_ = ensureModelConnector(&p, &m)
cid := p.ConnectorID + ":" + m.ID
label := p.Name + " / " + m.Name
if m.Name == "" {
label = p.Name + " / " + m.ID
}
result = append(result, connector.Option{
Label: label,
Value: cid,
})
}
} else {
result = append(result, connector.Option{
Label: p.Name,
Value: p.ConnectorID,
})
}
}
return result
}
// ownerMatch returns true if the provider owner matches the requested scope.
func ownerMatch(po, want *ProviderOwner) bool {
if want.Type == "team" {
return po.Type == "team" && po.TeamID == want.TeamID
}
return po.Type == "user" && po.UserID == want.UserID
}
// capabilitiesFromConn extracts *llm.Capabilities from a connector's settings.
func capabilitiesFromConn(conn connector.Connector) *goullm.Capabilities {
if conn == nil {
return defaultCaps()
}
settings := conn.Setting()
if settings != nil {
if caps, ok := settings["capabilities"]; ok {
if c, ok := caps.(*goullm.Capabilities); ok {
return c
}
if c, ok := caps.(goullm.Capabilities); ok {
return &c
}
}
}
return defaultCaps()
}
// capsToMap converts Capabilities to map[string]interface{} for process handlers.
func capsToMap(caps *goullm.Capabilities) map[string]interface{} {
if caps == nil {
return nil
}
result := make(map[string]interface{})
if caps.Vision != nil {
result["vision"] = caps.Vision
}
result["audio"] = caps.Audio
result["stt"] = caps.STT
result["tool_calls"] = caps.ToolCalls
result["reasoning"] = caps.Reasoning
result["streaming"] = caps.Streaming
result["json"] = caps.JSON
result["multimodal"] = caps.Multimodal
result["temperature_adjustable"] = caps.TemperatureAdjustable
return result
}
func defaultCaps() *goullm.Capabilities {
return &goullm.Capabilities{
Vision: false,
ToolCalls: false,
Audio: false,
Reasoning: false,
Streaming: false,
JSON: false,
Multimodal: false,
TemperatureAdjustable: true,
}
}

239
llmprovider/models_test.go Normal file
View file

@ -0,0 +1,239 @@
package llmprovider_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/yao/setting"
)
func TestGetModel(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "model-get")
conn, err := r.GetModel(p.ConnectorID)
require.NoError(t, err)
assert.NotNil(t, conn)
s := conn.Setting()
host, _ := s["host"].(string)
assert.Equal(t, "https://api.openai.com", host)
}
func TestGetModelNotFound(t *testing.T) {
r := setupRegistryWithSetting(t)
_, err := r.GetModel("nonexistent-connector")
assert.Error(t, err)
assert.Contains(t, err.Error(), "not found")
}
func TestGetModelByProviderKey(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "model-key")
conn, err := r.GetModel(p.Key)
require.NoError(t, err)
assert.NotNil(t, conn)
}
func TestGetRoleModel(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "rolemodel-prov")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
conn, err := r.GetRoleModel("default")
require.NoError(t, err)
assert.NotNil(t, conn)
s := conn.Setting()
host, _ := s["host"].(string)
assert.Equal(t, "https://api.openai.com", host)
}
func TestGetDefaultModel(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "default-model")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
conn, err := r.GetDefaultModel()
require.NoError(t, err)
assert.NotNil(t, conn)
}
func TestGetDefaultModelByUser(t *testing.T) {
r := setupRegistryWithSetting(t)
sysP := createTestProviderForRole(t, r, "dm-sys")
userP := createTestProviderForRole(t, r, "dm-user")
err := r.SetDefaults(map[string]string{"default": sysP.Key})
require.NoError(t, err)
_, err = setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeUser, UserID: "dm-u1"},
"llm.roles",
map[string]interface{}{
"default": map[string]interface{}{
"provider": userP.Key,
"model": "gpt-4o",
},
},
)
require.NoError(t, err)
conn, err := r.GetDefaultModelByUser("dm-u1")
require.NoError(t, err)
assert.NotNil(t, conn)
s := conn.Setting()
model, _ := s["model"].(string)
assert.Equal(t, "gpt-4o", model)
}
func TestGetCapabilities(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "caps-prov")
caps, err := r.GetCapabilities(p.ConnectorID)
require.NoError(t, err)
assert.NotNil(t, caps)
}
func TestGetRoleCapabilities(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "rolecaps-prov")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
caps, err := r.GetRoleCapabilities("default")
require.NoError(t, err)
assert.NotNil(t, caps)
}
func TestListModels(t *testing.T) {
r := setupRegistryWithSetting(t)
_ = createTestProviderForRole(t, r, "listm-prov")
opts := r.ListModels()
assert.NotEmpty(t, opts, "should have at least the created provider")
found := false
for _, o := range opts {
if o.Label == "Test listm-prov / GPT-4o" {
found = true
break
}
}
assert.True(t, found, "should contain the test provider's model")
}
func TestListModelsByUser(t *testing.T) {
r := setupRegistryWithSetting(t)
_ = createTestProviderForRole(t, r, "listmu-prov")
opts := r.ListModelsByUser("some-user")
assert.NotEmpty(t, opts)
}
func TestListModelsReturnsConnectorOption(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "opt-prov")
opts := r.ListModels()
modelCID := p.ConnectorID + ":gpt-4o"
found := false
for _, o := range opts {
if o.Value == modelCID {
found = true
assert.Equal(t, "Test opt-prov / GPT-4o", o.Label)
}
}
assert.True(t, found, "should contain model-level option with colon-separated CID")
}
func TestListModelsIncludesBuiltin(t *testing.T) {
r := setupRegistryWithSetting(t)
opts := r.ListModels()
builtinCount := 0
for _, o := range opts {
for _, ai := range connector.AIConnectors {
if o.Value == ai.Value {
builtinCount++
break
}
}
}
t.Logf("ListModels returned %d options, %d matching builtin AIConnectors (total: %d)",
len(opts), builtinCount, len(connector.AIConnectors))
}
func TestProcessGetModel(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "proc-model")
proc := process.New("llmprovider.getmodel", p.ConnectorID)
result, err := proc.Exec()
require.NoError(t, err)
m, ok := result.(map[string]interface{})
require.True(t, ok)
host, _ := m["host"].(string)
assert.Equal(t, "https://api.openai.com", host)
}
func TestProcessListModels(t *testing.T) {
r := setupRegistryWithSetting(t)
_ = createTestProviderForRole(t, r, "proc-listm")
proc := process.New("llmprovider.listmodels")
result, err := proc.Exec()
require.NoError(t, err)
list, ok := result.([]interface{})
require.True(t, ok)
assert.NotEmpty(t, list)
item := list[0].(map[string]interface{})
assert.Contains(t, item, "label")
assert.Contains(t, item, "value")
}
func TestProcessGetCapabilities(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "proc-caps")
proc := process.New("llmprovider.getcapabilities", p.ConnectorID)
result, err := proc.Exec()
require.NoError(t, err)
m, ok := result.(map[string]interface{})
require.True(t, ok)
assert.Contains(t, m, "streaming")
assert.Contains(t, m, "tool_calls")
}
func TestProcessGetRoleModel(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "proc-rm")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
proc := process.New("llmprovider.getrolemodel", "default")
result, err := proc.Exec()
require.NoError(t, err)
m, ok := result.(map[string]interface{})
require.True(t, ok)
host, _ := m["host"].(string)
assert.Equal(t, "https://api.openai.com", host)
}

View file

@ -3,12 +3,14 @@ package llmprovider
import (
"encoding/json"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/kun/exception"
)
func init() {
process.RegisterGroup("llmprovider", map[string]process.Handler{
// --- existing ---
"get": ProcessGet,
"getmasked": ProcessGetMasked,
"create": ProcessCreate,
@ -18,6 +20,43 @@ func init() {
"getsetting": ProcessGetSetting,
"getpresets": ProcessGetPresets,
"getpreset": ProcessGetPreset,
// --- roles ---
"getrole": ProcessGetRole,
"getrolebyuser": ProcessGetRoleByUser,
"getrolebyteam": ProcessGetRoleByTeam,
"listroles": ProcessListRoles,
"listrolesbyuser": ProcessListRolesByUser,
"listrolesbyteam": ProcessListRolesByTeam,
// --- models ---
"getmodel": ProcessGetModel,
"getrolemodel": ProcessGetRoleModel,
"getrolemodelbyuser": ProcessGetRoleModelByUser,
"getrolemodelbyteam": ProcessGetRoleModelByTeam,
"getdefaultmodel": ProcessGetDefaultModel,
"getdefaultmodelbyuser": ProcessGetDefaultModelByUser,
"getdefaultmodelbyteam": ProcessGetDefaultModelByTeam,
"getvisionmodel": ProcessGetVisionModel,
"getvisionmodelbyuser": ProcessGetVisionModelByUser,
"getvisionmodelbyteam": ProcessGetVisionModelByTeam,
"getaudiomodel": ProcessGetAudioModel,
"getaudiomodelbyuser": ProcessGetAudioModelByUser,
"getaudiomodelbyteam": ProcessGetAudioModelByTeam,
"getembeddingmodel": ProcessGetEmbeddingModel,
"getembeddingmodelbyuser": ProcessGetEmbeddingModelByUser,
"getembeddingmodelbyteam": ProcessGetEmbeddingModelByTeam,
// --- capabilities ---
"getcapabilities": ProcessGetCapabilities,
"getrolecapabilities": ProcessGetRoleCapabilities,
"getrolecapabilitiesbyuser": ProcessGetRoleCapabilitiesByUser,
"getrolecapabilitiesbyteam": ProcessGetRoleCapabilitiesByTeam,
// --- list models ---
"listmodels": ProcessListModels,
"listmodelsbyuser": ProcessListModelsByUser,
"listmodelsbyteam": ProcessListModelsByTeam,
})
}
@ -29,12 +68,14 @@ func requireGlobal() {
// ProcessGet retrieves a provider by key.
// Args[0] string: provider key
// Args[1] bool: withKey (optional, default false) — true returns plain-text APIKey
func ProcessGet(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
key := p.ArgsString(0)
provider, err := Global.Get(key)
withKey := len(p.Args) > 1 && toBool(p.Args[1])
provider, err := Global.Get(key, withKey)
if err != nil {
exception.New(err.Error(), 404).Throw()
}
@ -115,7 +156,8 @@ func ProcessDelete(p *process.Process) interface{} {
}
// ProcessList returns providers matching a filter.
// Args[0] map: ProviderFilter (optional)
// Args[0] map: ProviderFilter (optional)
// Args[1] bool: withKey (optional, default false) — true returns plain-text APIKeys
func ProcessList(p *process.Process) interface{} {
requireGlobal()
@ -130,7 +172,8 @@ func ProcessList(p *process.Process) interface{} {
}
}
result, err := Global.List(filter)
withKey := len(p.Args) > 1 && toBool(p.Args[1])
result, err := Global.List(filter, withKey)
if err != nil {
exception.New(err.Error(), 500).Throw()
}
@ -168,3 +211,384 @@ func ProcessGetPreset(p *process.Process) interface{} {
}
return preset
}
// ---------------------------------------------------------------------------
// Roles
// ---------------------------------------------------------------------------
// ProcessGetRole returns the connectorID for a role (system scope).
// Args[0] string: role name
func ProcessGetRole(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
cid, err := Global.GetRole(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return cid
}
// ProcessGetRoleByUser returns the connectorID for a role (user > system merge).
// Args[0] string: role, Args[1] string: userID
func ProcessGetRoleByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(2)
cid, err := Global.GetRoleByUser(p.ArgsString(0), p.ArgsString(1))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return cid
}
// ProcessGetRoleByTeam returns the connectorID for a role (team > system merge).
// Args[0] string: role, Args[1] string: teamID
func ProcessGetRoleByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(2)
cid, err := Global.GetRoleByTeam(p.ArgsString(0), p.ArgsString(1))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return cid
}
// ProcessListRoles returns all role assignments (system scope).
func ProcessListRoles(p *process.Process) interface{} {
requireGlobal()
roles, err := Global.ListRoles()
if err != nil {
exception.New(err.Error(), 500).Throw()
}
return rolesToMap(roles)
}
// ProcessListRolesByUser returns all role assignments (user > system merge).
// Args[0] string: userID
func ProcessListRolesByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
roles, err := Global.ListRolesByUser(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 500).Throw()
}
return rolesToMap(roles)
}
// ProcessListRolesByTeam returns all role assignments (team > system merge).
// Args[0] string: teamID
func ProcessListRolesByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
roles, err := Global.ListRolesByTeam(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 500).Throw()
}
return rolesToMap(roles)
}
// ---------------------------------------------------------------------------
// Models
// ---------------------------------------------------------------------------
// ProcessGetModel returns the connector setting map by connectorID.
// Args[0] string: connectorID
func ProcessGetModel(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetModel(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetRoleModel returns the connector setting map for a role (system scope).
// Args[0] string: role
func ProcessGetRoleModel(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetRoleModel(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetRoleModelByUser returns the connector setting map for a role (user scope).
// Args[0] string: role, Args[1] string: userID
func ProcessGetRoleModelByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(2)
conn, err := Global.GetRoleModelByUser(p.ArgsString(0), p.ArgsString(1))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetRoleModelByTeam returns the connector setting map for a role (team scope).
// Args[0] string: role, Args[1] string: teamID
func ProcessGetRoleModelByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(2)
conn, err := Global.GetRoleModelByTeam(p.ArgsString(0), p.ArgsString(1))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetDefaultModel returns the default model connector setting map.
func ProcessGetDefaultModel(p *process.Process) interface{} {
requireGlobal()
conn, err := Global.GetDefaultModel()
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetDefaultModelByUser returns the default model for a user.
// Args[0] string: userID
func ProcessGetDefaultModelByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetDefaultModelByUser(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetDefaultModelByTeam returns the default model for a team.
// Args[0] string: teamID
func ProcessGetDefaultModelByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetDefaultModelByTeam(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetVisionModel returns the vision model connector setting map.
func ProcessGetVisionModel(p *process.Process) interface{} {
requireGlobal()
conn, err := Global.GetVisionModel()
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetVisionModelByUser returns the vision model for a user.
// Args[0] string: userID
func ProcessGetVisionModelByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetVisionModelByUser(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetVisionModelByTeam returns the vision model for a team.
// Args[0] string: teamID
func ProcessGetVisionModelByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetVisionModelByTeam(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetAudioModel returns the audio model connector setting map.
func ProcessGetAudioModel(p *process.Process) interface{} {
requireGlobal()
conn, err := Global.GetAudioModel()
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetAudioModelByUser returns the audio model for a user.
// Args[0] string: userID
func ProcessGetAudioModelByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetAudioModelByUser(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetAudioModelByTeam returns the audio model for a team.
// Args[0] string: teamID
func ProcessGetAudioModelByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetAudioModelByTeam(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetEmbeddingModel returns the embedding model connector setting map.
func ProcessGetEmbeddingModel(p *process.Process) interface{} {
requireGlobal()
conn, err := Global.GetEmbeddingModel()
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetEmbeddingModelByUser returns the embedding model for a user.
// Args[0] string: userID
func ProcessGetEmbeddingModelByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetEmbeddingModelByUser(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ProcessGetEmbeddingModelByTeam returns the embedding model for a team.
// Args[0] string: teamID
func ProcessGetEmbeddingModelByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
conn, err := Global.GetEmbeddingModelByTeam(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return conn.Setting()
}
// ---------------------------------------------------------------------------
// Capabilities
// ---------------------------------------------------------------------------
// ProcessGetCapabilities returns capabilities for a connectorID.
// Args[0] string: connectorID
func ProcessGetCapabilities(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
caps, err := Global.GetCapabilities(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return capsToMap(caps)
}
// ProcessGetRoleCapabilities returns capabilities for a role (system scope).
// Args[0] string: role
func ProcessGetRoleCapabilities(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
caps, err := Global.GetRoleCapabilities(p.ArgsString(0))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return capsToMap(caps)
}
// ProcessGetRoleCapabilitiesByUser returns capabilities for a role (user scope).
// Args[0] string: role, Args[1] string: userID
func ProcessGetRoleCapabilitiesByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(2)
caps, err := Global.GetRoleCapabilitiesByUser(p.ArgsString(0), p.ArgsString(1))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return capsToMap(caps)
}
// ProcessGetRoleCapabilitiesByTeam returns capabilities for a role (team scope).
// Args[0] string: role, Args[1] string: teamID
func ProcessGetRoleCapabilitiesByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(2)
caps, err := Global.GetRoleCapabilitiesByTeam(p.ArgsString(0), p.ArgsString(1))
if err != nil {
exception.New(err.Error(), 404).Throw()
}
return capsToMap(caps)
}
// ---------------------------------------------------------------------------
// List Models
// ---------------------------------------------------------------------------
// ProcessListModels returns all enabled models as []Option (system scope).
func ProcessListModels(p *process.Process) interface{} {
requireGlobal()
return optionsToSlice(Global.ListModels())
}
// ProcessListModelsByUser returns models visible to a user.
// Args[0] string: userID
func ProcessListModelsByUser(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
return optionsToSlice(Global.ListModelsByUser(p.ArgsString(0)))
}
// ProcessListModelsByTeam returns models visible to a team.
// Args[0] string: teamID
func ProcessListModelsByTeam(p *process.Process) interface{} {
requireGlobal()
p.ValidateArgNums(1)
return optionsToSlice(Global.ListModelsByTeam(p.ArgsString(0)))
}
// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------
func rolesToMap(roles map[string]RoleTarget) map[string]interface{} {
result := make(map[string]interface{}, len(roles))
for k, v := range roles {
result[k] = map[string]interface{}{
"provider": v.Provider,
"model": v.Model,
}
}
return result
}
func optionsToSlice(opts []connector.Option) []interface{} {
result := make([]interface{}, len(opts))
for i, o := range opts {
result[i] = map[string]interface{}{
"label": o.Label,
"value": o.Value,
}
}
return result
}
func toBool(v interface{}) bool {
switch b := v.(type) {
case bool:
return b
case float64:
return b != 0
case int:
return b != 0
case string:
return b == "true" || b == "1"
default:
return false
}
}

View file

@ -37,13 +37,22 @@ func TestProcessGet(t *testing.T) {
setupRegistry(t)
createViaProcess(t, "proc-get")
// Default: masked
p := process.New("llmprovider.get", "proc-get")
result, err := p.Exec()
require.NoError(t, err)
m := toMapResult(t, result)
assert.Equal(t, "proc-get", m["key"])
assert.Equal(t, "sk-proc-test", m["api_key"])
assert.NotEqual(t, "sk-proc-test", m["api_key"], "default should be masked")
// withKey=true: plain text
p2 := process.New("llmprovider.get", "proc-get", true)
result2, err := p2.Exec()
require.NoError(t, err)
m2 := toMapResult(t, result2)
assert.Equal(t, "sk-proc-test", m2["api_key"], "withKey=true should return plain text")
}
func TestProcessGetMasked(t *testing.T) {

View file

@ -47,8 +47,15 @@ func (r *Registry) SetEncryptionKey(key string) {
r.encKey = key
}
// shouldExposeKey returns true when the caller explicitly requests plain-text APIKey.
func shouldExposeKey(withKey []bool) bool {
return len(withKey) > 0 && withKey[0]
}
// Get retrieves a provider by key. Lazily ensures its connector is registered.
func (r *Registry) Get(key string) (*Provider, error) {
// By default the APIKey is masked; pass withKey=true to get the plain-text key
// (only for internal LLM-request paths).
func (r *Registry) Get(key string, withKey ...bool) (*Provider, error) {
r.mu.RLock()
defer r.mu.RUnlock()
@ -58,18 +65,48 @@ func (r *Registry) Get(key string) (*Provider, error) {
}
_ = ensureConnector(p)
if !shouldExposeKey(withKey) {
cp := *p
cp.APIKey = maskAPIKey(cp.APIKey)
return &cp, nil
}
return p, nil
}
// GetMasked retrieves a provider with the API key masked for display.
func (r *Registry) GetMasked(key string) (*Provider, error) {
p, err := r.Get(key)
// GetByConnectorID finds a provider by its ConnectorID field (linear scan).
// Use when the caller has a ConnectorID but not the store Key.
// By default the APIKey is masked; pass withKey=true to get the plain-text key.
func (r *Registry) GetByConnectorID(cid string, withKey ...bool) (*Provider, error) {
r.mu.RLock()
defer r.mu.RUnlock()
keys, err := indexGet(r.store, r.cache)
if err != nil {
return nil, err
}
cp := *p
cp.APIKey = maskAPIKey(cp.APIKey)
return &cp, nil
for _, key := range keys {
p, err := storeGet(r.store, r.cache, key, r.encKey)
if err != nil {
continue
}
if p.ConnectorID == cid {
_ = ensureConnector(p)
if !shouldExposeKey(withKey) {
cp := *p
cp.APIKey = maskAPIKey(cp.APIKey)
return &cp, nil
}
return p, nil
}
}
return nil, fmt.Errorf("provider with connector_id %q not found", cid)
}
// Deprecated: GetMasked is equivalent to Get(key) since Get now masks by default.
func (r *Registry) GetMasked(key string) (*Provider, error) {
return r.Get(key)
}
// Create adds a new provider. Persists, caches, registers connector, and updates index.
@ -161,7 +198,8 @@ func (r *Registry) Delete(key string) error {
}
// List returns providers matching the filter.
func (r *Registry) List(filter *ProviderFilter) ([]Provider, error) {
// By default the APIKey is masked; pass withKey=true to get plain-text keys.
func (r *Registry) List(filter *ProviderFilter, withKey ...bool) ([]Provider, error) {
r.mu.RLock()
defer r.mu.RUnlock()
@ -170,6 +208,7 @@ func (r *Registry) List(filter *ProviderFilter) ([]Provider, error) {
return nil, err
}
expose := shouldExposeKey(withKey)
var result []Provider
for _, key := range keys {
p, err := storeGet(r.store, r.cache, key, r.encKey)
@ -180,7 +219,9 @@ func (r *Registry) List(filter *ProviderFilter) ([]Provider, error) {
continue
}
cp := *p
cp.APIKey = maskAPIKey(cp.APIKey)
if !expose {
cp.APIKey = maskAPIKey(cp.APIKey)
}
result = append(result, cp)
}
return result, nil
@ -217,7 +258,7 @@ func (r *Registry) Reload() error {
// GetConnector returns the runtime connector for a given provider key.
func (r *Registry) GetConnector(key string) (connector.Connector, error) {
p, err := r.Get(key)
p, err := r.Get(key, true)
if err != nil {
return nil, err
}

View file

@ -125,6 +125,143 @@ func TestGetMasked(t *testing.T) {
assert.Contains(t, got.APIKey, "xxxx")
}
// ---------------------------------------------------------------------------
// withKey behavior
// ---------------------------------------------------------------------------
func TestGet_DefaultMasked(t *testing.T) {
r := setupRegistry(t)
p := testProvider
_, err := r.Create(&p)
require.NoError(t, err)
got, err := r.Get("test-openai")
require.NoError(t, err)
assert.NotEqual(t, "sk-test-xxxxx", got.APIKey, "Get() default should mask APIKey")
assert.Contains(t, got.APIKey, "*")
}
func TestGet_WithKeyTrue(t *testing.T) {
r := setupRegistry(t)
p := testProvider
_, err := r.Create(&p)
require.NoError(t, err)
got, err := r.Get("test-openai", true)
require.NoError(t, err)
assert.Equal(t, "sk-test-xxxxx", got.APIKey, "Get(key, true) should return plain text APIKey")
}
func TestGetByConnectorID_DefaultMasked(t *testing.T) {
r := setupRegistry(t)
p := testProvider
created, err := r.Create(&p)
require.NoError(t, err)
got, err := r.GetByConnectorID(created.ConnectorID)
require.NoError(t, err)
assert.NotEqual(t, "sk-test-xxxxx", got.APIKey, "GetByConnectorID() default should mask")
assert.Contains(t, got.APIKey, "*")
}
func TestGetByConnectorID_WithKeyTrue(t *testing.T) {
r := setupRegistry(t)
p := testProvider
created, err := r.Create(&p)
require.NoError(t, err)
got, err := r.GetByConnectorID(created.ConnectorID, true)
require.NoError(t, err)
assert.Equal(t, "sk-test-xxxxx", got.APIKey, "GetByConnectorID(cid, true) should return plain text")
}
func TestList_DefaultMasked(t *testing.T) {
r := setupRegistry(t)
p := testProvider
_, err := r.Create(&p)
require.NoError(t, err)
list, err := r.List(&llmprovider.ProviderFilter{Source: llmprovider.ProviderSourceDynamic})
require.NoError(t, err)
require.True(t, len(list) > 0)
for _, item := range list {
assert.NotEqual(t, "sk-test-xxxxx", item.APIKey, "List() default should mask all APIKeys")
}
}
func TestList_WithKeyTrue(t *testing.T) {
r := setupRegistry(t)
p := testProvider
_, err := r.Create(&p)
require.NoError(t, err)
list, err := r.List(&llmprovider.ProviderFilter{Source: llmprovider.ProviderSourceDynamic}, true)
require.NoError(t, err)
found := false
for _, item := range list {
if item.Key == "test-openai" {
found = true
assert.Equal(t, "sk-test-xxxxx", item.APIKey, "List(filter, true) should return plain text")
}
}
assert.True(t, found)
}
func TestGetMasked_EqualsGetDefault(t *testing.T) {
r := setupRegistry(t)
p := testProvider
_, err := r.Create(&p)
require.NoError(t, err)
fromGet, err := r.Get("test-openai")
require.NoError(t, err)
fromGetMasked, err := r.GetMasked("test-openai")
require.NoError(t, err)
assert.Equal(t, fromGet.APIKey, fromGetMasked.APIKey, "GetMasked should equal Get (both masked by default)")
}
func TestListModels_ConnectorHasRealKey(t *testing.T) {
r := setupRegistry(t)
p := llmprovider.Provider{
Key: "realkey-prov",
Name: "RealKey Test",
Type: "openai",
APIURL: "https://api.openai.com",
APIKey: "sk-real-secret-key-12345",
Models: []llmprovider.ModelInfo{
{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true},
},
Enabled: true,
Owner: llmprovider.ProviderOwner{Type: "user", UserID: "rk-user"},
}
_, err := r.Create(&p)
require.NoError(t, err)
opts := r.ListModelsByUser("rk-user")
require.True(t, len(opts) > 0, "should have at least one model option")
var modelCID string
for _, o := range opts {
if o.Label == "RealKey Test / GPT-4o" {
modelCID = o.Value
break
}
}
require.NotEmpty(t, modelCID, "should find the model option")
conn, err := connector.Select(modelCID)
require.NoError(t, err, "model connector should be registered")
s := conn.Setting()
key, _ := s["key"].(string)
assert.Equal(t, "sk-real-secret-key-12345", key, "connector should have the real API key, not masked")
}
func TestGetLazy(t *testing.T) {
r := setupRegistry(t)
@ -369,13 +506,13 @@ func TestEncryptionRoundTrip(t *testing.T) {
_, err := r.Create(&p)
require.NoError(t, err)
got, err := r.Get("test-encrypted")
got, err := r.Get("test-encrypted", true)
require.NoError(t, err)
assert.Equal(t, "sk-test-xxxxx", got.APIKey, "APIKey should be decrypted on read")
assert.Equal(t, "sk-test-xxxxx", got.APIKey, "APIKey should be decrypted on read with withKey=true")
masked, err := r.GetMasked("test-encrypted")
masked, err := r.Get("test-encrypted")
require.NoError(t, err)
assert.NotEqual(t, "sk-test-xxxxx", masked.APIKey)
assert.NotEqual(t, "sk-test-xxxxx", masked.APIKey, "Get without withKey should mask")
assert.Contains(t, masked.APIKey, "xxxx")
// Verify raw store value is encrypted

161
llmprovider/roles.go Normal file
View file

@ -0,0 +1,161 @@
package llmprovider
import (
"fmt"
"github.com/yaoapp/yao/setting"
)
// RolesNamespace is the setting namespace for LLM role assignments.
const RolesNamespace = "llm.roles"
// SetDefaults writes agent.yml system-level role defaults into setting.Global
// under ScopeSystem. roles maps role names (e.g. "default", "vision") to connectorIDs.
func (r *Registry) SetDefaults(roles map[string]string) error {
if setting.Global == nil {
return fmt.Errorf("setting registry not initialized")
}
data := make(map[string]interface{})
for role, cid := range roles {
p, err := r.Get(cid, true)
if err != nil {
// Builtin providers: Key == ConnectorID, use connectorID directly
data[role] = map[string]interface{}{
"provider": cid,
"model": "",
}
continue
}
data[role] = map[string]interface{}{
"provider": p.Key,
"model": defaultModel(p),
}
}
_, err := setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeSystem},
RolesNamespace,
data,
)
return err
}
// GetRole returns the connectorID for a role at system scope.
func (r *Registry) GetRole(role string) (string, error) {
return r.resolveRole(role, "", "")
}
// GetRoleByUser returns the connectorID for a role, merged user > system.
func (r *Registry) GetRoleByUser(role, userID string) (string, error) {
return r.resolveRole(role, userID, "")
}
// GetRoleByTeam returns the connectorID for a role, merged team > system.
func (r *Registry) GetRoleByTeam(role, teamID string) (string, error) {
return r.resolveRole(role, "", teamID)
}
// GetRoleBy returns the connectorID for a role, scoped by identity (team > user).
func (r *Registry) GetRoleBy(role string, id Identity) (string, error) {
if id.GetTeamID() != "" {
return r.GetRoleByTeam(role, id.GetTeamID())
}
return r.GetRoleByUser(role, id.GetUserID())
}
// ListRoles returns all role assignments at system scope.
func (r *Registry) ListRoles() (map[string]RoleTarget, error) {
return r.listRoles("", "")
}
// ListRolesByUser returns all role assignments, merged user > system.
func (r *Registry) ListRolesByUser(userID string) (map[string]RoleTarget, error) {
return r.listRoles(userID, "")
}
// ListRolesByTeam returns all role assignments, merged team > system.
func (r *Registry) ListRolesByTeam(teamID string) (map[string]RoleTarget, error) {
return r.listRoles("", teamID)
}
// ListRolesBy returns all role assignments, scoped by identity (team > user).
func (r *Registry) ListRolesBy(id Identity) (map[string]RoleTarget, error) {
if id.GetTeamID() != "" {
return r.ListRolesByTeam(id.GetTeamID())
}
return r.ListRolesByUser(id.GetUserID())
}
// ---------------------------------------------------------------------------
// internal
// ---------------------------------------------------------------------------
func (r *Registry) resolveRole(role, userID, teamID string) (string, error) {
if setting.Global == nil {
return "", fmt.Errorf("setting registry not initialized")
}
merged, err := setting.Global.GetMerged(userID, teamID, RolesNamespace)
if err != nil {
return "", fmt.Errorf("role %q not configured: %w", role, err)
}
target, ok := merged[role]
if !ok {
return "", fmt.Errorf("role %q not configured", role)
}
cid := r.extractConnectorID(target)
if cid == "" {
return "", fmt.Errorf("role %q has invalid target", role)
}
return cid, nil
}
func (r *Registry) listRoles(userID, teamID string) (map[string]RoleTarget, error) {
if setting.Global == nil {
return nil, fmt.Errorf("setting registry not initialized")
}
merged, err := setting.Global.GetMerged(userID, teamID, RolesNamespace)
if err != nil {
return nil, fmt.Errorf("failed to load roles: %w", err)
}
result := make(map[string]RoleTarget)
for role, target := range merged {
rt := parseRoleTarget(target)
if rt.Provider != "" {
result[role] = rt
}
}
return result, nil
}
func (r *Registry) extractConnectorID(target interface{}) string {
rt := parseRoleTarget(target)
if rt.Provider == "" {
return ""
}
p, err := r.Get(rt.Provider, true)
if err != nil {
// Builtin providers: Key == ConnectorID
return rt.Provider
}
return p.ConnectorID
}
func parseRoleTarget(v interface{}) RoleTarget {
switch t := v.(type) {
case map[string]interface{}:
provider, _ := t["provider"].(string)
model, _ := t["model"].(string)
return RoleTarget{Provider: provider, Model: model}
case RoleTarget:
return t
default:
return RoleTarget{}
}
}

217
llmprovider/roles_test.go Normal file
View file

@ -0,0 +1,217 @@
package llmprovider_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/gou/store"
"github.com/yaoapp/yao/llmprovider"
"github.com/yaoapp/yao/setting"
)
func setupRegistryWithSetting(t *testing.T) *llmprovider.Registry {
t.Helper()
r := setupRegistry(t)
err := setting.Init()
require.NoError(t, err)
t.Cleanup(func() {
s, _ := store.Get("__yao.store")
if s != nil {
s.Del("setting:*")
}
c, _ := store.Get("__yao.cache")
if c != nil {
c.Del("setting:*")
}
})
return r
}
func createTestProviderForRole(t *testing.T, r *llmprovider.Registry, key string) *llmprovider.Provider {
t.Helper()
p := llmprovider.Provider{
Key: key,
Name: "Test " + key,
Type: "openai",
APIURL: "https://api.openai.com",
APIKey: "sk-test-role",
Enabled: true,
Models: []llmprovider.ModelInfo{
{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"vision", "tool_calls", "streaming"}, Enabled: true},
},
Owner: llmprovider.ProviderOwner{Type: "system"},
}
created, err := r.Create(&p)
require.NoError(t, err)
return created
}
func TestSetDefaults(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "sd-provider")
err := r.SetDefaults(map[string]string{
"default": p.Key,
})
require.NoError(t, err)
merged, err := setting.Global.GetMerged("", "", "llm.roles")
require.NoError(t, err)
def, ok := merged["default"].(map[string]interface{})
require.True(t, ok)
assert.Equal(t, p.Key, def["provider"])
assert.Equal(t, "gpt-4o", def["model"])
}
func TestGetRole(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "role-provider")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
cid, err := r.GetRole("default")
require.NoError(t, err)
assert.Equal(t, p.ConnectorID, cid)
}
func TestGetRoleByUser(t *testing.T) {
r := setupRegistryWithSetting(t)
sysP := createTestProviderForRole(t, r, "sys-prov")
err := r.SetDefaults(map[string]string{"default": sysP.Key})
require.NoError(t, err)
userP := createTestProviderForRole(t, r, "user-prov")
_, err = setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeUser, UserID: "u1"},
"llm.roles",
map[string]interface{}{
"default": map[string]interface{}{
"provider": userP.Key,
"model": "gpt-4o",
},
},
)
require.NoError(t, err)
cid, err := r.GetRoleByUser("default", "u1")
require.NoError(t, err)
assert.Equal(t, userP.ConnectorID, cid, "user scope should override system")
cidSys, err := r.GetRole("default")
require.NoError(t, err)
assert.Equal(t, sysP.ConnectorID, cidSys, "system scope should still return system provider")
}
func TestGetRoleByTeam(t *testing.T) {
r := setupRegistryWithSetting(t)
sysP := createTestProviderForRole(t, r, "sys-team-prov")
err := r.SetDefaults(map[string]string{"default": sysP.Key})
require.NoError(t, err)
teamP := createTestProviderForRole(t, r, "team-prov")
_, err = setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeTeam, TeamID: "t1"},
"llm.roles",
map[string]interface{}{
"default": map[string]interface{}{
"provider": teamP.Key,
"model": "gpt-4o",
},
},
)
require.NoError(t, err)
cid, err := r.GetRoleByTeam("default", "t1")
require.NoError(t, err)
assert.Equal(t, teamP.ConnectorID, cid, "team scope should override system")
}
func TestGetRoleNotConfigured(t *testing.T) {
_ = setupRegistryWithSetting(t)
_, err := llmprovider.Global.GetRole("nonexistent")
assert.Error(t, err)
assert.Contains(t, err.Error(), "not configured")
}
func TestListRoles(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "list-prov")
err := r.SetDefaults(map[string]string{
"default": p.Key,
"vision": p.Key,
})
require.NoError(t, err)
roles, err := r.ListRoles()
require.NoError(t, err)
assert.Contains(t, roles, "default")
assert.Contains(t, roles, "vision")
assert.Equal(t, p.Key, roles["default"].Provider)
}
func TestListRolesByUser(t *testing.T) {
r := setupRegistryWithSetting(t)
sysP := createTestProviderForRole(t, r, "list-sys")
userP := createTestProviderForRole(t, r, "list-user")
err := r.SetDefaults(map[string]string{"default": sysP.Key, "vision": sysP.Key})
require.NoError(t, err)
_, err = setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeUser, UserID: "u2"},
"llm.roles",
map[string]interface{}{
"default": map[string]interface{}{
"provider": userP.Key,
"model": "gpt-4o",
},
},
)
require.NoError(t, err)
roles, err := r.ListRolesByUser("u2")
require.NoError(t, err)
assert.Equal(t, userP.Key, roles["default"].Provider, "user override for default")
assert.Equal(t, sysP.Key, roles["vision"].Provider, "system fallback for vision")
}
func TestProcessGetRole(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "proc-role")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
proc := process.New("llmprovider.getrole", "default")
result, err := proc.Exec()
require.NoError(t, err)
assert.Equal(t, p.ConnectorID, result)
}
func TestProcessListRoles(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "proc-list-role")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
proc := process.New("llmprovider.listroles")
result, err := proc.Exec()
require.NoError(t, err)
m, ok := result.(map[string]interface{})
require.True(t, ok)
def, ok := m["default"].(map[string]interface{})
require.True(t, ok)
assert.Equal(t, p.Key, def["provider"])
}

414
llmprovider/scope_test.go Normal file
View file

@ -0,0 +1,414 @@
package llmprovider_test
import (
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/yao/llmprovider"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/setting"
)
// ---------------------------------------------------------------------------
// Identity interface
// ---------------------------------------------------------------------------
func TestAuthorizedInfoSatisfiesIdentity(t *testing.T) {
info := &oauthTypes.AuthorizedInfo{UserID: "u1", TeamID: "t1"}
var id llmprovider.Identity = info
assert.Equal(t, "u1", id.GetUserID())
assert.Equal(t, "t1", id.GetTeamID())
}
func TestAuthorizedInfoNilSafe(t *testing.T) {
var info *oauthTypes.AuthorizedInfo
assert.Equal(t, "", info.GetUserID())
assert.Equal(t, "", info.GetTeamID())
}
// ---------------------------------------------------------------------------
// ListModels owner filtering
// ---------------------------------------------------------------------------
func TestListModelsByUserFiltersOwner(t *testing.T) {
r := setupRegistryWithSetting(t)
createOwnedProvider(t, r, "user-alice-prov", llmprovider.ProviderOwner{Type: "user", UserID: "alice"})
createOwnedProvider(t, r, "user-bob-prov", llmprovider.ProviderOwner{Type: "user", UserID: "bob"})
createOwnedProvider(t, r, "team-x-prov", llmprovider.ProviderOwner{Type: "team", TeamID: "x"})
opts := r.ListModelsByUser("alice")
labels := optLabels(opts)
assert.Contains(t, labels, "Test user-alice-prov / GPT-4o", "should include alice's model")
assert.NotContains(t, labels, "Test user-bob-prov / GPT-4o", "should exclude bob's model")
assert.NotContains(t, labels, "Test team-x-prov / GPT-4o", "should exclude team model")
}
func TestListModelsByTeamFiltersOwner(t *testing.T) {
r := setupRegistryWithSetting(t)
createOwnedProvider(t, r, "team-alpha-prov", llmprovider.ProviderOwner{Type: "team", TeamID: "alpha"})
createOwnedProvider(t, r, "team-beta-prov", llmprovider.ProviderOwner{Type: "team", TeamID: "beta"})
createOwnedProvider(t, r, "user-u1-prov", llmprovider.ProviderOwner{Type: "user", UserID: "u1"})
opts := r.ListModelsByTeam("alpha")
labels := optLabels(opts)
assert.Contains(t, labels, "Test team-alpha-prov / GPT-4o", "should include alpha's model")
assert.NotContains(t, labels, "Test team-beta-prov / GPT-4o", "should exclude beta's model")
assert.NotContains(t, labels, "Test user-u1-prov / GPT-4o", "should exclude user model")
}
func TestListModelsByIncludesBuiltin(t *testing.T) {
r := setupRegistryWithSetting(t)
createOwnedProvider(t, r, "user-x-prov", llmprovider.ProviderOwner{Type: "user", UserID: "x"})
all := r.ListModels()
byUser := r.ListModelsByUser("x")
builtinAll := countBuiltin(all)
builtinScoped := countBuiltin(byUser)
assert.Equal(t, builtinAll, builtinScoped, "ByUser should include all builtin providers")
}
func TestListModelsBy_TeamRouting(t *testing.T) {
r := setupRegistryWithSetting(t)
createOwnedProvider(t, r, "team-rt-prov", llmprovider.ProviderOwner{Type: "team", TeamID: "rt"})
createOwnedProvider(t, r, "user-rt-prov", llmprovider.ProviderOwner{Type: "user", UserID: "rt"})
info := &oauthTypes.AuthorizedInfo{UserID: "rt", TeamID: "rt"}
opts := r.ListModelsBy(info)
labels := optLabels(opts)
assert.Contains(t, labels, "Test team-rt-prov / GPT-4o", "team takes priority when TeamID is set")
assert.NotContains(t, labels, "Test user-rt-prov / GPT-4o", "user model should be excluded when TeamID is set")
}
func TestListModelsBy_UserFallback(t *testing.T) {
r := setupRegistryWithSetting(t)
createOwnedProvider(t, r, "user-fb-prov", llmprovider.ProviderOwner{Type: "user", UserID: "fb"})
info := &oauthTypes.AuthorizedInfo{UserID: "fb"}
opts := r.ListModelsBy(info)
labels := optLabels(opts)
assert.Contains(t, labels, "Test user-fb-prov / GPT-4o", "should include user model when no TeamID")
}
// ---------------------------------------------------------------------------
// ListModels per-model expansion
// ---------------------------------------------------------------------------
func TestListModelsExpandsMultipleModels(t *testing.T) {
r := setupRegistryWithSetting(t)
p := llmprovider.Provider{
Key: "multi-model-prov",
Name: "MultiModel",
Type: "openai",
APIURL: "https://api.openai.com",
APIKey: "sk-test",
Models: []llmprovider.ModelInfo{
{ID: "gpt-4o", Name: "GPT-4o", Enabled: true},
{ID: "gpt-4o-mini", Name: "GPT-4o Mini", Enabled: true},
{ID: "gpt-disabled", Name: "Disabled", Enabled: false},
},
Enabled: true,
Owner: llmprovider.ProviderOwner{Type: "user", UserID: "multi-u"},
}
_, err := r.Create(&p)
require.NoError(t, err)
opts := r.ListModelsByUser("multi-u")
labels := optLabels(opts)
values := optValues(opts)
assert.Contains(t, labels, "MultiModel / GPT-4o")
assert.Contains(t, labels, "MultiModel / GPT-4o Mini")
assert.NotContains(t, labels, "MultiModel / Disabled", "disabled model should not appear")
// Values should be "providerCID:modelID" format
for _, v := range values {
if strings.Contains(v, "multi-model-prov") {
assert.Contains(t, v, ":", "dynamic model option should use colon-separated format")
}
}
}
func TestGetModelWithModelLevelCID(t *testing.T) {
r := setupRegistryWithSetting(t)
p := llmprovider.Provider{
Key: "mlcid-prov",
Name: "MLTest",
Type: "openai",
APIURL: "https://api.openai.com",
APIKey: "sk-test",
Models: []llmprovider.ModelInfo{
{ID: "gpt-4o", Name: "GPT-4o", Enabled: true},
{ID: "gpt-4o-mini", Name: "GPT-4o Mini", Enabled: true},
},
Enabled: true,
Owner: llmprovider.ProviderOwner{Type: "team", TeamID: "mlcid-t1"},
}
created, err := r.Create(&p)
require.NoError(t, err)
modelCID := created.ConnectorID + ":gpt-4o"
conn, err := r.GetModel(modelCID)
require.NoError(t, err)
assert.NotNil(t, conn)
s := conn.Setting()
model, _ := s["model"].(string)
assert.Equal(t, "gpt-4o", model, "model-level connector should have the correct model")
modelCID2 := created.ConnectorID + ":gpt-4o-mini"
conn2, err := r.GetModel(modelCID2)
require.NoError(t, err)
s2 := conn2.Setting()
model2, _ := s2["model"].(string)
assert.Equal(t, "gpt-4o-mini", model2, "second model should have its own connector")
}
// ---------------------------------------------------------------------------
// GetModel ConnectorID reverse lookup
// ---------------------------------------------------------------------------
func TestGetModelByConnectorIDReverseLookup(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createOwnedProvider(t, r, "rev-prov", llmprovider.ProviderOwner{Type: "user", UserID: "u99"})
cid := p.ConnectorID
assert.NotEqual(t, p.Key, cid, "dynamic provider ConnectorID should differ from Key")
_ = connector.Unregister(cid)
conn, err := r.GetModel(cid)
require.NoError(t, err, "GetModel should find provider via ConnectorID reverse lookup")
assert.NotNil(t, conn)
s := conn.Setting()
host, _ := s["host"].(string)
assert.Equal(t, "https://api.openai.com", host)
}
func TestGetByConnectorID(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createOwnedProvider(t, r, "bycid-prov", llmprovider.ProviderOwner{Type: "team", TeamID: "t55"})
found, err := r.GetByConnectorID(p.ConnectorID)
require.NoError(t, err)
assert.Equal(t, p.Key, found.Key)
}
func TestGetByConnectorIDNotFound(t *testing.T) {
r := setupRegistryWithSetting(t)
_, err := r.GetByConnectorID("nonexistent-cid")
assert.Error(t, err)
assert.Contains(t, err.Error(), "not found")
}
// ---------------------------------------------------------------------------
// GetRoleBy / ListRolesBy / GetRoleModelBy — Identity routing
// ---------------------------------------------------------------------------
func TestGetRoleBy_TeamPriority(t *testing.T) {
r := setupRegistryWithSetting(t)
sysP := createTestProviderForRole(t, r, "grb-sys")
teamP := createTestProviderForRole(t, r, "grb-team")
err := r.SetDefaults(map[string]string{"default": sysP.Key})
require.NoError(t, err)
_, err = setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeTeam, TeamID: "grb-t1"},
"llm.roles",
map[string]interface{}{
"default": map[string]interface{}{
"provider": teamP.Key,
"model": "gpt-4o",
},
},
)
require.NoError(t, err)
info := &oauthTypes.AuthorizedInfo{UserID: "u1", TeamID: "grb-t1"}
cid, err := r.GetRoleBy("default", info)
require.NoError(t, err)
assert.Equal(t, teamP.ConnectorID, cid, "should resolve via team scope when TeamID is set")
}
func TestGetRoleBy_UserFallback(t *testing.T) {
r := setupRegistryWithSetting(t)
sysP := createTestProviderForRole(t, r, "grbu-sys")
userP := createTestProviderForRole(t, r, "grbu-user")
err := r.SetDefaults(map[string]string{"default": sysP.Key})
require.NoError(t, err)
_, err = setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeUser, UserID: "grbu-u1"},
"llm.roles",
map[string]interface{}{
"default": map[string]interface{}{
"provider": userP.Key,
"model": "gpt-4o",
},
},
)
require.NoError(t, err)
info := &oauthTypes.AuthorizedInfo{UserID: "grbu-u1"}
cid, err := r.GetRoleBy("default", info)
require.NoError(t, err)
assert.Equal(t, userP.ConnectorID, cid, "should resolve via user scope when no TeamID")
}
func TestListRolesBy(t *testing.T) {
r := setupRegistryWithSetting(t)
sysP := createTestProviderForRole(t, r, "lrb-sys")
teamP := createTestProviderForRole(t, r, "lrb-team")
err := r.SetDefaults(map[string]string{"default": sysP.Key, "vision": sysP.Key})
require.NoError(t, err)
_, err = setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeTeam, TeamID: "lrb-t1"},
"llm.roles",
map[string]interface{}{
"default": map[string]interface{}{
"provider": teamP.Key,
"model": "gpt-4o",
},
},
)
require.NoError(t, err)
info := &oauthTypes.AuthorizedInfo{TeamID: "lrb-t1"}
roles, err := r.ListRolesBy(info)
require.NoError(t, err)
assert.Equal(t, teamP.Key, roles["default"].Provider, "team override for default")
assert.Equal(t, sysP.Key, roles["vision"].Provider, "system fallback for vision")
}
func TestGetRoleModelBy(t *testing.T) {
r := setupRegistryWithSetting(t)
sysP := createTestProviderForRole(t, r, "grmb-sys")
userP := createTestProviderForRole(t, r, "grmb-user")
err := r.SetDefaults(map[string]string{"default": sysP.Key})
require.NoError(t, err)
_, err = setting.Global.Set(
setting.ScopeID{Scope: setting.ScopeUser, UserID: "grmb-u1"},
"llm.roles",
map[string]interface{}{
"default": map[string]interface{}{
"provider": userP.Key,
"model": "gpt-4o",
},
},
)
require.NoError(t, err)
info := &oauthTypes.AuthorizedInfo{UserID: "grmb-u1"}
conn, err := r.GetRoleModelBy("default", info)
require.NoError(t, err)
assert.NotNil(t, conn)
s := conn.Setting()
model, _ := s["model"].(string)
assert.Equal(t, "gpt-4o", model)
}
func TestGetDefaultModelBy(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "gdmb-prov")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
info := &oauthTypes.AuthorizedInfo{UserID: "gdmb-u1"}
conn, err := r.GetDefaultModelBy(info)
require.NoError(t, err)
assert.NotNil(t, conn)
}
func TestGetRoleCapabilitiesBy(t *testing.T) {
r := setupRegistryWithSetting(t)
p := createTestProviderForRole(t, r, "grcb-prov")
err := r.SetDefaults(map[string]string{"default": p.Key})
require.NoError(t, err)
info := &oauthTypes.AuthorizedInfo{UserID: "grcb-u1"}
caps, err := r.GetRoleCapabilitiesBy("default", info)
require.NoError(t, err)
assert.NotNil(t, caps)
}
// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------
func createOwnedProvider(t *testing.T, r *llmprovider.Registry, key string, owner llmprovider.ProviderOwner) *llmprovider.Provider {
t.Helper()
p := llmprovider.Provider{
Key: key,
Name: "Test " + key,
Type: "openai",
APIURL: "https://api.openai.com",
APIKey: "sk-test-owned",
Enabled: true,
Models: []llmprovider.ModelInfo{
{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"vision", "tool_calls", "streaming"}, Enabled: true},
},
Owner: owner,
}
created, err := r.Create(&p)
require.NoError(t, err)
return created
}
func optLabels(opts []connector.Option) []string {
labels := make([]string, len(opts))
for i, o := range opts {
labels[i] = o.Label
}
return labels
}
func optValues(opts []connector.Option) []string {
values := make([]string, len(opts))
for i, o := range opts {
values[i] = o.Value
}
return values
}
func countBuiltin(opts []connector.Option) int {
n := 0
for _, o := range opts {
for _, ai := range connector.AIConnectors {
if o.Value == ai.Value {
n++
break
}
}
}
return n
}

View file

@ -35,19 +35,45 @@ func defaultModel(p *Provider) string {
// marshalDSL builds a connector DSL JSON from the flat Provider fields.
func marshalDSL(p *Provider) ([]byte, error) {
opts := map[string]interface{}{
"host": p.APIURL,
"key": p.APIKey,
"model": defaultModel(p),
}
if caps := aggregateCapabilities(p); len(caps) > 0 {
opts["capabilities"] = caps
}
dsl := map[string]interface{}{
"type": p.Type,
"name": p.Name,
"label": p.Name,
"options": map[string]interface{}{
"host": p.APIURL,
"key": p.APIKey,
"model": defaultModel(p),
},
"type": p.Type,
"name": p.Name,
"label": p.Name,
"options": opts,
}
return json.Marshal(dsl)
}
// aggregateCapabilities merges all model capabilities into a single map.
// Falls back to type-based defaults when no model declares explicit caps.
func aggregateCapabilities(p *Provider) map[string]bool {
caps := make(map[string]bool)
for _, m := range p.Models {
for _, c := range m.Capabilities {
caps[c] = true
}
}
if len(caps) == 0 {
switch p.Type {
case "openai", "anthropic":
caps["streaming"] = true
caps["tool_calls"] = true
caps["temperature_adjustable"] = true
}
}
return caps
}
// ensureConnector makes sure the provider's connector is registered in the runtime.
// Builtin providers are managed by engine.Load and skipped here.
func ensureConnector(p *Provider) error {
@ -80,6 +106,75 @@ func ensureConnector(p *Provider) error {
return nil
}
// ensureModelConnector registers a per-model connector for a dynamic provider.
// The connector ID format is "{providerConnectorID}:{modelID}".
func ensureModelConnector(p *Provider, m *ModelInfo) error {
if p.Source == ProviderSourceBuiltIn {
return nil
}
if !p.Enabled || !m.Enabled {
return nil
}
baseCID := p.ConnectorID
if baseCID == "" {
baseCID = connectorID(p)
}
cid := baseCID + ":" + m.ID
if _, err := connector.Select(cid); err == nil {
return nil
}
dslJSON, err := marshalModelDSL(p, m)
if err != nil {
return fmt.Errorf("ensureModelConnector %s:%s: %w", p.Key, m.ID, err)
}
_, err = connector.LoadSourceSync(dslJSON, cid, "__registry/"+baseCID+"/"+m.ID+".conn.yao")
if err != nil {
return fmt.Errorf("ensureModelConnector %s:%s: LoadSourceSync: %w", p.Key, m.ID, err)
}
return nil
}
// marshalModelDSL builds a connector DSL for a specific model within a provider.
func marshalModelDSL(p *Provider, m *ModelInfo) ([]byte, error) {
caps := make(map[string]bool)
for _, c := range m.Capabilities {
caps[c] = true
}
if len(caps) == 0 {
switch p.Type {
case "openai", "anthropic":
caps["streaming"] = true
caps["tool_calls"] = true
caps["temperature_adjustable"] = true
}
}
opts := map[string]interface{}{
"host": p.APIURL,
"key": p.APIKey,
"model": m.ID,
}
if len(caps) > 0 {
opts["capabilities"] = caps
}
name := m.Name
if name == "" {
name = m.ID
}
dsl := map[string]interface{}{
"type": p.Type,
"name": name,
"label": name,
"options": opts,
}
return json.Marshal(dsl)
}
// unregisterConnector removes the provider's connector from the runtime.
func unregisterConnector(p *Provider) error {
if p.Source == ProviderSourceBuiltIn {

View file

@ -1,5 +1,12 @@
package llmprovider
// Identity abstracts a caller's user/team context for scope-aware lookups.
// Implemented by oauthTypes.AuthorizedInfo and any struct with UserID/TeamID.
type Identity interface {
GetUserID() string
GetTeamID() string
}
// Provider represents a configured LLM provider (one vendor connection with multiple models).
// Fields align with the frontend ProviderConfig interface.
type Provider struct {
@ -75,9 +82,6 @@ type ProviderTestResult struct {
LatencyMs int64 `json:"latency_ms,omitempty"`
}
// RoleAssignment maps model roles to specific provider+model pairs.
type RoleAssignment map[string]RoleTarget
// RoleTarget identifies a provider and model for a given role.
type RoleTarget struct {
Provider string `json:"provider"`

View file

@ -1,11 +1,14 @@
package llm
import (
"fmt"
"strings"
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/connector"
agentllm "github.com/yaoapp/yao/agent/llm"
"github.com/yaoapp/yao/llmprovider"
"github.com/yaoapp/yao/openapi/oauth/authorized"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response"
)
@ -34,7 +37,6 @@ func Attach(group *gin.RouterGroup, oauth oauthTypes.OAuth) {
func listProviders(c *gin.Context) {
allProviders := make([]Provider, 0)
// Parse filter parameters from query string
filtersParam := c.Query("filters")
var filters []string
if filtersParam != "" {
@ -44,49 +46,68 @@ func listProviders(c *gin.Context) {
}
}
for _, opt := range connector.AIConnectors {
connType := getConnectorType(opt.Value)
if connType == "openai" || connType == "anthropic" {
conn, ok := connector.Connectors[opt.Value]
if !ok {
continue
}
fmt.Printf("[llm/providers] filtersParam=%q\n", filtersParam)
capabilities := getCapabilitiesFromConn(conn)
info := authorized.GetInfo(c)
fmt.Printf("[llm/providers] identity: UserID=%q TeamID=%q\n", info.GetUserID(), info.GetTeamID())
// Apply capability filters
if len(filters) > 0 && !matchesFilters(capabilities, filters) {
continue
}
allProviders = append(allProviders, Provider{
Label: opt.Label,
Value: opt.Value,
Type: connType,
Builtin: conn.GetMetaInfo().Builtin,
Capabilities: capabilities,
})
}
var opts []connector.Option
if llmprovider.Global != nil {
opts = llmprovider.Global.ListModelsBy(info)
} else {
opts = connector.AIConnectors
}
fmt.Printf("[llm/providers] ListModelsBy returned %d options\n", len(opts))
for i, o := range opts {
fmt.Printf("[llm/providers] [%d] label=%q value=%q\n", i, o.Label, o.Value)
}
for _, opt := range opts {
var conn connector.Connector
var err error
if llmprovider.Global != nil {
conn, err = llmprovider.Global.GetModel(opt.Value)
} else {
conn, err = connector.Select(opt.Value)
}
if err != nil {
fmt.Printf("[llm/providers] GetModel(%q) FAILED: %v\n", opt.Value, err)
continue
}
connType := connectorType(conn)
if connType != "openai" && connType != "anthropic" {
fmt.Printf("[llm/providers] SKIP %q: type=%q (not openai/anthropic)\n", opt.Value, connType)
continue
}
capabilities := getCapabilitiesFromConn(conn)
if len(filters) > 0 && !matchesFilters(capabilities, filters) {
fmt.Printf("[llm/providers] SKIP %q: caps filter %v not matched (streaming=%v)\n", opt.Value, filters, capabilities["streaming"])
continue
}
allProviders = append(allProviders, Provider{
Label: opt.Label,
Value: opt.Value,
Type: connType,
Builtin: conn.GetMetaInfo().Builtin,
Capabilities: capabilities,
})
}
fmt.Printf("[llm/providers] returning %d providers\n", len(allProviders))
response.RespondWithSuccess(c, response.StatusOK, allProviders)
}
// getConnectorType retrieves the connector type by checking the global connector map
func getConnectorType(id string) string {
conn, ok := connector.Connectors[id]
if !ok {
return "unknown"
}
// connectorType returns the type string for a connector.
func connectorType(conn connector.Connector) string {
if conn.Is(connector.OPENAI) {
return "openai"
}
if conn.Is(connector.ANTHROPIC) {
return "anthropic"
}
return "unknown"
}

View file

@ -635,6 +635,22 @@ type AuthorizedInfo struct {
Constraints DataConstraints `json:"constraints,omitempty"`
}
// GetUserID implements llmprovider.Identity.
func (auth *AuthorizedInfo) GetUserID() string {
if auth == nil {
return ""
}
return auth.UserID
}
// GetTeamID implements llmprovider.Identity.
func (auth *AuthorizedInfo) GetTeamID() string {
if auth == nil {
return ""
}
return auth.TeamID
}
// AuthorizedToMap converts AuthorizedInfo to map[string]interface{}
// This is useful for passing authorized information to runtime bridges (e.g., V8)
func (auth *AuthorizedInfo) AuthorizedToMap() map[string]interface{} {

View file

@ -16,7 +16,7 @@ import (
"github.com/yaoapp/yao/setting"
)
const llmRolesNS = "llm.roles"
var llmRolesNS = llmprovider.RolesNamespace
func llmEnsureEncKey() {
if llmprovider.Global != nil && config.Conf.DB.AESKey != "" {