opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor assistant capabilities and update data handling

Browse source 
- Enhanced the getConnectorCapabilities method to prioritize model capabilities and connector settings, improving capability retrieval logic.
- Deprecated the tools field in the Assistant model, transitioning to MCP for tool management, and updated related methods accordingly.
- Introduced new fields for connector options and prompt presets in the Assistant model, allowing for more flexible configurations.
- Updated the GetAssistant method to support field selection, improving data retrieval efficiency and flexibility.
- Refactored tests and documentation to reflect changes in the assistant structure and capabilities, ensuring clarity and maintainability.
This commit is contained in:
Max 2025-12-02 12:36:47 +08:00
parent 3dd63e5530
commit 339a486eb4
19 changed files with 1019 additions and 480 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
agent/assistant/agent.go
View file

@ -323,33 +323,56 @@ func (ast *Assistant) GetConnector(ctx *context.Context) (connector.Connector, *
}
// Get connector capabilities from settings
capabilities := ast.getConnectorCapabilities(connectorID)
capabilities := ast.getConnectorCapabilities(conn)
return conn, capabilities, nil
}
// getConnectorCapabilities get the capabilities of a connector from settings
func (ast *Assistant) getConnectorCapabilities(connectorID string) *openai.Capabilities {
// Get model capabilities from global configuration
modelCaps, exists := modelCapabilities[connectorID]
if !exists {
// Return default capabilities if model not found in configuration
falseVal := false
// Priority: 1. modelCapabilities mapping, 2. connector's Setting()["capabilities"]
func (ast *Assistant) getConnectorCapabilities(conn connector.Connector) *openai.Capabilities {
if conn == nil {
return &openai.Capabilities{
Vision: falseVal,
Vision: false,
ToolCalls: false,
Audio: false,
Reasoning: false,
Streaming: false,
JSON: false,
Multimodal: false,
TemperatureAdjustable: true, // Default to true for non-reasoning models
TemperatureAdjustable: true,
}
}
// Return capabilities directly
// Note: TemperatureAdjustable is automatically set in connector.Setting() based on Reasoning flag
return &modelCaps
// Get connector ID
connectorID := conn.ID()
// Priority 1: Check global modelCapabilities mapping
if modelCaps, exists := modelCapabilities[connectorID]; exists {
return &modelCaps
}
// Priority 2: Get capabilities from connector's Setting() method
// Modern connectors (post-upgrade) provide default capabilities via Setting()
settings := conn.Setting()
if caps, ok := settings["capabilities"]; ok {
if capabilities, ok := caps.(*openai.Capabilities); ok {
return capabilities
}
}
// Fallback: Return minimal default capabilities
// This should rarely happen with upgraded connectors
return &openai.Capabilities{
Vision: false,
ToolCalls: false,
Audio: false,
Reasoning: false,
Streaming: false,
JSON: false,
Multimodal: false,
TemperatureAdjustable: true, // Default to true for non-reasoning models
}
}
// Info get the assistant information

40
agent/assistant/assistant.go
View file

@ -4,7 +4,6 @@ import (
"fmt"
"path"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/i18n"
store "github.com/yaoapp/yao/agent/store/types"
@ -105,7 +104,6 @@ func (ast *Assistant) Map() map[string]interface{} {
"prompts": ast.Prompts,
"kb": ast.KB,
"mcp": ast.MCP,
"tools": ast.Tools,
"workflow": ast.Workflow,
"tags": ast.Tags,
"mentionable": ast.Mentionable,
@ -247,20 +245,6 @@ func (ast *Assistant) Clone() *Assistant {
copy(clone.Prompts, ast.Prompts)
}
// Deep copy tools
if ast.Tools != nil {
clone.Tools = &store.ToolCalls{}
if ast.Tools.Tools != nil {
clone.Tools.Tools = make([]store.Tool, len(ast.Tools.Tools))
copy(clone.Tools.Tools, ast.Tools.Tools)
}
if ast.Tools.Prompts != nil {
clone.Tools.Prompts = make([]store.Prompt, len(ast.Tools.Prompts))
copy(clone.Tools.Prompts, ast.Tools.Prompts)
}
}
// Deep copy workflow
if ast.Workflow != nil {
clone.Workflow = &store.Workflow{}
@ -328,29 +312,7 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
ast.Connector = v
}
if v, has := data["tools"]; has {
switch tools := v.(type) {
case []store.Tool:
ast.Tools = &store.ToolCalls{
Tools: tools,
Prompts: ast.Prompts,
}
case *store.ToolCalls:
ast.Tools = tools
default:
raw, err := jsoniter.Marshal(tools)
if err != nil {
return err
}
ast.Tools = &store.ToolCalls{}
err = jsoniter.Unmarshal(raw, &ast.Tools)
if err != nil {
return err
}
}
}
// Note: tools field is deprecated, now handled by MCP
if v, ok := data["type"].(string); ok {
ast.Type = v

33
agent/assistant/load.go
View file

@ -181,7 +181,8 @@ func LoadStore(id string) (*Assistant, error) {
return nil, fmt.Errorf("storage is not set")
}
storeModel, err := storage.GetAssistant(id)
// Request all fields when loading assistant from store
storeModel, err := storage.GetAssistant(id, store.AssistantFullFields)
if err != nil {
return nil, err
}
@ -509,32 +510,10 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
}
}
// tools
if tools, has := data["tools"]; has {
switch vv := tools.(type) {
case []store.Tool:
assistant.Tools = &store.ToolCalls{
Tools: vv,
Prompts: assistant.Prompts,
}
case store.ToolCalls:
assistant.Tools = &vv
default:
raw, err := jsoniter.Marshal(tools)
if err != nil {
return nil, fmt.Errorf("tools format error %s", err.Error())
}
var tools store.ToolCalls
err = jsoniter.Unmarshal(raw, &tools)
if err != nil {
return nil, fmt.Errorf("tools format error %s", err.Error())
}
assistant.Tools = &tools
}
}
// tools - deprecated, now handled by MCP
// if tools, has := data["tools"]; has {
// ... removed ...
// }
// kb
if kb, has := data["kb"]; has {

3
agent/store/mongo/mongo.go
View file

@ -76,7 +76,8 @@ func (m *Mongo) GetAssistants(filter types.AssistantFilter, locale ...string) (*
}
// GetAssistant retrieves a single assistant by ID
func (m *Mongo) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
// fields: Optional list of fields to retrieve. If empty, a default set of fields will be returned.
func (m *Mongo) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
return nil, nil
}

3
agent/store/redis/redis.go
View file

@ -76,7 +76,8 @@ func (r *Redis) GetAssistants(filter types.AssistantFilter, locale ...string) (*
}
// GetAssistant retrieves a single assistant by ID
func (r *Redis) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
// fields: Optional list of fields to retrieve. If empty, a default set of fields will be returned.
func (r *Redis) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
return nil, nil
}

36
agent/store/types/convert.go
View file

@ -217,6 +217,9 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
if path, ok := data["path"].(string); ok {
model.Path = path
}
if source, ok := data["source"].(string); ok {
model.Source = source
}
if description, ok := data["description"].(string); ok {
model.Description = description
}
@ -277,6 +280,28 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
}
}
// PromptPresets
if promptPresets, ok := data["prompt_presets"]; ok && promptPresets != nil {
raw, err := jsoniter.Marshal(promptPresets)
if err == nil {
var pp map[string][]Prompt
if err := jsoniter.Unmarshal(raw, &pp); err == nil {
model.PromptPresets = pp
}
}
}
// ConnectorOptions
if connectorOptions, ok := data["connector_options"]; ok && connectorOptions != nil {
raw, err := jsoniter.Marshal(connectorOptions)
if err == nil {
var co ConnectorOptions
if err := jsoniter.Unmarshal(raw, &co); err == nil {
model.ConnectorOptions = &co
}
}
}
// KB
if kb, ok := data["kb"]; ok && kb != nil {
kbConverted, err := ToKnowledgeBase(kb)
@ -301,17 +326,6 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
}
}
// Tools
if tools, ok := data["tools"]; ok && tools != nil {
raw, err := jsoniter.Marshal(tools)
if err == nil {
var tc ToolCalls
if err := jsoniter.Unmarshal(raw, &tc); err == nil {
model.Tools = &tc
}
}
}
// Placeholder
if placeholder, ok := data["placeholder"]; ok && placeholder != nil {
raw, err := jsoniter.Marshal(placeholder)

235
agent/store/types/convert_test.go
View file

@ -428,24 +428,38 @@ func TestToAssistantModel(t *testing.T) {
"name": "Test Assistant",
"avatar": "https://example.com/avatar.png",
"connector": "openai",
"path": "/path/to/assistant",
"description": "Test description",
"share": "team",
"built_in": true,
"readonly": false,
"public": true,
"mentionable": true,
"automated": false,
"sort": 100,
"created_at": int64(1609459200),
"updated_at": int64(1609459300),
"tags": []string{"tag1", "tag2"},
"connector_options": map[string]interface{}{
"optional": true,
"connectors": []string{"openai", "anthropic"},
"filters": []string{"vision", "tool_calls"},
},
"path": "/path/to/assistant",
"description": "Test description",
"share": "team",
"built_in": true,
"readonly": false,
"public": true,
"mentionable": true,
"automated": false,
"sort": 100,
"created_at": int64(1609459200),
"updated_at": int64(1609459300),
"tags": []string{"tag1", "tag2"},
"options": map[string]interface{}{
"temperature": 0.7,
},
"prompts": []map[string]interface{}{
{"role": "system", "content": "You are helpful"},
},
"prompt_presets": map[string]interface{}{
"chat": []map[string]interface{}{
{"role": "system", "content": "You are a chat assistant"},
},
"task": []map[string]interface{}{
{"role": "system", "content": "You are a task assistant"},
},
},
"source": "function hook() { return 'test'; }",
"kb": map[string]interface{}{
"collections": []string{"col1"},
},
@ -455,9 +469,6 @@ func TestToAssistantModel(t *testing.T) {
"workflow": map[string]interface{}{
"workflows": []string{"wf1"},
},
"tools": map[string]interface{}{
"calls": []string{"tool1"},
},
"placeholder": map[string]interface{}{
"title": "Enter message",
},
@ -489,9 +500,25 @@ func TestToAssistantModel(t *testing.T) {
if result.Connector != "openai" {
t.Errorf("Expected Connector 'openai', got '%s'", result.Connector)
}
if result.ConnectorOptions == nil {
t.Error("Expected ConnectorOptions to be set")
} else {
if !result.ConnectorOptions.Optional {
t.Error("Expected ConnectorOptions.Optional to be true")
}
if len(result.ConnectorOptions.Connectors) != 2 {
t.Errorf("Expected 2 connectors in options, got %d", len(result.ConnectorOptions.Connectors))
}
if len(result.ConnectorOptions.Filters) != 2 {
t.Errorf("Expected 2 filters, got %d", len(result.ConnectorOptions.Filters))
}
}
if result.Path != "/path/to/assistant" {
t.Errorf("Expected Path, got '%s'", result.Path)
}
if result.Source != "function hook() { return 'test'; }" {
t.Errorf("Expected Source, got '%s'", result.Source)
}
if result.Description != "Test description" {
t.Errorf("Expected Description, got '%s'", result.Description)
}
@ -531,6 +558,23 @@ func TestToAssistantModel(t *testing.T) {
if len(result.Prompts) != 1 {
t.Errorf("Expected 1 prompt, got %d", len(result.Prompts))
}
if result.PromptPresets == nil {
t.Error("Expected PromptPresets to be set")
} else {
if len(result.PromptPresets) != 2 {
t.Errorf("Expected 2 prompt presets, got %d", len(result.PromptPresets))
}
if chatPrompts, ok := result.PromptPresets["chat"]; !ok {
t.Error("Expected 'chat' prompt preset")
} else if len(chatPrompts) != 1 {
t.Errorf("Expected 1 chat prompt, got %d", len(chatPrompts))
}
if taskPrompts, ok := result.PromptPresets["task"]; !ok {
t.Error("Expected 'task' prompt preset")
} else if len(taskPrompts) != 1 {
t.Errorf("Expected 1 task prompt, got %d", len(taskPrompts))
}
}
if result.KB == nil {
t.Error("Expected KB to be set")
}
@ -540,9 +584,6 @@ func TestToAssistantModel(t *testing.T) {
if result.Workflow == nil {
t.Error("Expected Workflow to be set")
}
if result.Tools == nil {
t.Error("Expected Tools to be set")
}
if result.Placeholder == nil {
t.Error("Expected Placeholder to be set")
}
@ -583,7 +624,6 @@ func TestToAssistantModel(t *testing.T) {
"kb": nil,
"mcp": nil,
"workflow": nil,
"tools": nil,
"placeholder": nil,
"locales": nil,
}
@ -659,6 +699,163 @@ func TestToAssistantModel(t *testing.T) {
})
}
// TestToAssistantModelNewFields tests the newly added fields
func TestToAssistantModelNewFields(t *testing.T) {
t.Run("ConnectorOptions", func(t *testing.T) {
data := map[string]interface{}{
"connector_options": map[string]interface{}{
"optional": true,
"connectors": []string{"openai", "anthropic", "azure"},
"filters": []string{"vision", "tool_calls", "audio"},
},
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ConnectorOptions == nil {
t.Fatal("Expected ConnectorOptions to be set")
}
if !result.ConnectorOptions.Optional {
t.Error("Expected Optional to be true")
}
if len(result.ConnectorOptions.Connectors) != 3 {
t.Errorf("Expected 3 connectors, got %d", len(result.ConnectorOptions.Connectors))
}
if len(result.ConnectorOptions.Filters) != 3 {
t.Errorf("Expected 3 filters, got %d", len(result.ConnectorOptions.Filters))
}
})
t.Run("PromptPresets", func(t *testing.T) {
data := map[string]interface{}{
"prompt_presets": map[string]interface{}{
"chat": []map[string]interface{}{
{"role": "system", "content": "You are a helpful chat assistant"},
{"role": "user", "content": "Example question"},
},
"task": []map[string]interface{}{
{"role": "system", "content": "You are a task completion assistant"},
},
"analyze": []map[string]interface{}{
{"role": "system", "content": "You are a data analysis assistant"},
},
},
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.PromptPresets == nil {
t.Fatal("Expected PromptPresets to be set")
}
if len(result.PromptPresets) != 3 {
t.Errorf("Expected 3 prompt preset modes, got %d", len(result.PromptPresets))
}
if chatPrompts, ok := result.PromptPresets["chat"]; !ok {
t.Error("Expected 'chat' mode in prompt presets")
} else if len(chatPrompts) != 2 {
t.Errorf("Expected 2 prompts in chat mode, got %d", len(chatPrompts))
}
if taskPrompts, ok := result.PromptPresets["task"]; !ok {
t.Error("Expected 'task' mode in prompt presets")
} else if len(taskPrompts) != 1 {
t.Errorf("Expected 1 prompt in task mode, got %d", len(taskPrompts))
}
if analyzePrompts, ok := result.PromptPresets["analyze"]; !ok {
t.Error("Expected 'analyze' mode in prompt presets")
} else if len(analyzePrompts) != 1 {
t.Errorf("Expected 1 prompt in analyze mode, got %d", len(analyzePrompts))
}
})
t.Run("Source", func(t *testing.T) {
hookScript := `
function beforeChat(context) {
console.log('Hook called');
return context;
}
`
data := map[string]interface{}{
"source": hookScript,
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.Source != hookScript {
t.Errorf("Expected Source to match, got '%s'", result.Source)
}
})
t.Run("AllNewFields", func(t *testing.T) {
data := map[string]interface{}{
"connector_options": map[string]interface{}{
"optional": true,
"connectors": []string{"openai"},
"filters": []string{"vision"},
},
"prompt_presets": map[string]interface{}{
"chat": []map[string]interface{}{
{"role": "system", "content": "Chat mode"},
},
},
"source": "function test() {}",
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ConnectorOptions == nil {
t.Error("Expected ConnectorOptions to be set")
}
if result.PromptPresets == nil {
t.Error("Expected PromptPresets to be set")
}
if result.Source == "" {
t.Error("Expected Source to be set")
}
})
t.Run("NilNewFields", func(t *testing.T) {
data := map[string]interface{}{
"connector_options": nil,
"prompt_presets": nil,
"source": nil,
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ConnectorOptions != nil {
t.Error("Expected ConnectorOptions to be nil")
}
if result.PromptPresets != nil {
t.Error("Expected PromptPresets to be nil")
}
if result.Source != "" {
t.Error("Expected Source to be empty")
}
})
}
// TestToAssistantModelComplexTypes tests complex type conversions in ToAssistantModel
func TestToAssistantModelComplexTypes(t *testing.T) {
t.Run("CompleteLocales", func(t *testing.T) {

108
agent/store/types/fields.go
View file

@ -4,39 +4,43 @@ import "github.com/yaoapp/kun/log"
// AssistantAllowedFields defines the whitelist of fields that can be selected for assistants
var AssistantAllowedFields = map[string]bool{
"id": true,
"assistant_id": true,
"type": true,
"name": true,
"avatar": true,
"connector": true,
"description": true,
"path": true,
"sort": true,
"built_in": true,
"placeholder": true,
"options": true,
"prompts": true,
"workflow": true,
"kb": true,
"mcp": true,
"tools": true,
"tags": true,
"readonly": true,
"public": true,
"share": true,
"locales": true,
"automated": true,
"mentionable": true,
"created_at": true,
"updated_at": true,
"__yao_created_by": true,
"__yao_updated_by": true,
"__yao_team_id": true,
"__yao_tenant_id": true,
"id": true,
"assistant_id": true,
"type": true,
"name": true,
"avatar": true,
"connector": true,
"connector_options": true,
"description": true,
"path": true,
"sort": true,
"built_in": true,
"placeholder": true,
"options": true,
"prompts": true,
"prompt_presets": true,
"workflow": true,
"kb": true,
"mcp": true,
"source": true,
"tags": true,
"readonly": true,
"public": true,
"share": true,
"locales": true,
"uses": true,
"automated": true,
"mentionable": true,
"created_at": true,
"updated_at": true,
"__yao_created_by": true,
"__yao_updated_by": true,
"__yao_team_id": true,
"__yao_tenant_id": true,
}
// AssistantDefaultFields defines the default fields to select for assistants when no specific fields are requested
// These are lightweight fields suitable for list views and basic information display
var AssistantDefaultFields = []string{
"assistant_id",
"type",
@ -44,6 +48,7 @@ var AssistantDefaultFields = []string{
"avatar",
"connector",
"description",
"tags", // Tags for categorization (lightweight)
"sort",
"built_in",
"readonly",
@ -51,8 +56,51 @@ var AssistantDefaultFields = []string{
"share",
"automated",
"mentionable",
"kb", // Knowledge base configuration (lightweight)
"mcp", // MCP servers configuration (lightweight)
"created_at",
"updated_at",
"__yao_created_by", // Permission: creator user ID
"__yao_updated_by", // Permission: updater user ID
"__yao_team_id", // Permission: team ID
"__yao_tenant_id", // Permission: tenant ID
}
// AssistantFullFields defines all available fields including complex/large fields
// Use this when you need complete assistant data for backend processing
var AssistantFullFields = []string{
"assistant_id",
"type",
"name",
"avatar",
"connector",
"connector_options",
"description",
"path",
"sort",
"built_in",
"placeholder",
"options",
"prompts",
"prompt_presets",
"workflow",
"kb",
"mcp",
"source",
"tags",
"readonly",
"public",
"share",
"locales",
"uses",
"automated",
"mentionable",
"created_at",
"updated_at",
"__yao_created_by",
"__yao_updated_by",
"__yao_team_id",
"__yao_tenant_id",
}
// ValidateAssistantFields validates and filters assistant select fields against the whitelist

97
agent/store/types/fields_test.go
View file

@ -118,12 +118,15 @@ func TestAssistantAllowedFields(t *testing.T) {
complexFields := []string{
"options",
"prompts",
"prompt_presets",
"workflow",
"kb",
"mcp",
"tools",
"placeholder",
"locales",
"uses",
"connector_options",
"source",
}
for _, field := range complexFields {
if !AssistantAllowedFields[field] {
@ -139,6 +142,12 @@ func TestAssistantDefaultFields(t *testing.T) {
"assistant_id",
"name",
"type",
"kb", // Knowledge base is essential for assistant functionality
"mcp", // MCP servers are essential for assistant functionality
"__yao_created_by", // Permission fields are essential for access control
"__yao_updated_by",
"__yao_team_id",
"__yao_tenant_id",
}
defaultFieldsMap := make(map[string]bool)
@ -155,15 +164,17 @@ func TestAssistantDefaultFields(t *testing.T) {
t.Run("DoesNotContainSensitiveFields", func(t *testing.T) {
// Default fields should not include complex/large fields by default
// Note: kb, mcp, and tags are lightweight and included in defaults
sensitiveFields := []string{
"options",
"prompts",
"prompt_presets",
"workflow",
"kb",
"mcp",
"tools",
"placeholder",
"locales",
"uses",
"connector_options",
"source",
}
defaultFieldsMap := make(map[string]bool)
@ -178,3 +189,81 @@ func TestAssistantDefaultFields(t *testing.T) {
}
})
}
func TestAssistantFullFields(t *testing.T) {
t.Run("ContainsAllAllowedFields", func(t *testing.T) {
// Full fields should contain all fields from allowed fields
fullFieldsMap := make(map[string]bool)
for _, field := range AssistantFullFields {
fullFieldsMap[field] = true
}
for field := range AssistantAllowedFields {
if field == "id" {
// "id" is an alias for "assistant_id", skip
continue
}
if !fullFieldsMap[field] {
t.Errorf("Allowed field %s is missing from full fields", field)
}
}
})
t.Run("AllFieldsAreAllowed", func(t *testing.T) {
// All fields in full list should be in allowed fields
for _, field := range AssistantFullFields {
if !AssistantAllowedFields[field] {
t.Errorf("Full field %s is not in allowed fields", field)
}
}
})
t.Run("ContainsComplexFields", func(t *testing.T) {
// Full fields should include all complex/large fields
complexFields := []string{
"options",
"prompts",
"prompt_presets",
"workflow",
"kb",
"mcp",
"placeholder",
"locales",
"uses",
"connector_options",
"source",
}
fullFieldsMap := make(map[string]bool)
for _, field := range AssistantFullFields {
fullFieldsMap[field] = true
}
for _, field := range complexFields {
if !fullFieldsMap[field] {
t.Errorf("Complex field %s is missing from full fields", field)
}
}
})
t.Run("ContainsPermissionFields", func(t *testing.T) {
// Full fields should include permission fields
permissionFields := []string{
"__yao_created_by",
"__yao_updated_by",
"__yao_team_id",
"__yao_tenant_id",
}
fullFieldsMap := make(map[string]bool)
for _, field := range AssistantFullFields {
fullFieldsMap[field] = true
}
for _, field := range permissionFields {
if !fullFieldsMap[field] {
t.Errorf("Permission field %s is missing from full fields", field)
}
}
})
}

4
agent/store/types/store.go
View file

@ -91,8 +91,10 @@ type Store interface {
// GetAssistant retrieves a single assistant by ID
// assistantID: Assistant ID
// fields: List of fields to select, empty/nil means default fields (AssistantDefaultFields)
// locale: Optional locale for i18n translations
// Returns: Assistant information and potential error
GetAssistant(assistantID string, locale ...string) (*AssistantModel, error)
GetAssistant(assistantID string, fields []string, locale ...string) (*AssistantModel, error)
// DeleteAssistants deletes assistants based on filter conditions
// filter: Filter conditions

77
agent/store/types/types.go
View file

@ -221,34 +221,59 @@ type Placeholder struct {
Prompts []string `json:"prompts,omitempty"`
}
// ModelCapability defines the available model capability filters
type ModelCapability string
// Model capability constants for filtering connectors
const (
CapVision ModelCapability = "vision"
CapAudio ModelCapability = "audio"
CapToolCalls ModelCapability = "tool_calls"
CapReasoning ModelCapability = "reasoning"
CapStreaming ModelCapability = "streaming"
CapJSON ModelCapability = "json"
CapMultimodal ModelCapability = "multimodal"
CapTemperatureAdjustable ModelCapability = "temperature_adjustable"
)
// ConnectorOptions the connector selection options
// Allows defining optional connector selection with filtering capabilities
type ConnectorOptions struct {
Optional bool `json:"optional,omitempty"` // Whether connector is optional for user selection
Connectors []string `json:"connectors,omitempty"` // List of available connectors, empty means all connectors are available
Filters []ModelCapability `json:"filters,omitempty"` // Filter by model capabilities, conditions can be stacked
}
// AssistantModel the assistant database model
type AssistantModel struct {
ID string `json:"assistant_id"` // Assistant ID
Type string `json:"type,omitempty"` // Assistant Type, default is assistant
Name string `json:"name,omitempty"` // Assistant Name
Avatar string `json:"avatar,omitempty"` // Assistant Avatar
Connector string `json:"connector"` // AI Connector
Path string `json:"path,omitempty"` // Assistant Path
BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant
Sort int `json:"sort,omitempty"` // Assistant Sort
Description string `json:"description,omitempty"` // Assistant Description
Tags []string `json:"tags,omitempty"` // Assistant Tags
Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly
Public bool `json:"public,omitempty"` // Whether this assistant is shared across all teams in the platform
Share string `json:"share,omitempty"` // Assistant sharing scope (private/team)
Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable
Automated bool `json:"automated,omitempty"` // Whether this assistant is automated
Options map[string]interface{} `json:"options,omitempty"` // AI Options
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts
KB *KnowledgeBase `json:"kb,omitempty"` // Knowledge base configuration
MCP *MCPServers `json:"mcp,omitempty"` // MCP servers configuration
Tools *ToolCalls `json:"tools,omitempty"` // Assistant Tools
Workflow *Workflow `json:"workflow,omitempty"` // Workflow configuration
Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp
ID string `json:"assistant_id"` // Assistant ID
Type string `json:"type,omitempty"` // Assistant Type, default is assistant
Name string `json:"name,omitempty"` // Assistant Name
Avatar string `json:"avatar,omitempty"` // Assistant Avatar
Connector string `json:"connector"` // AI Connector (default connector)
ConnectorOptions *ConnectorOptions `json:"connector_options,omitempty"` // Connector selection options for user to choose from
Path string `json:"path,omitempty"` // Assistant Path
BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant
Sort int `json:"sort,omitempty"` // Assistant Sort
Description string `json:"description,omitempty"` // Assistant Description
Tags []string `json:"tags,omitempty"` // Assistant Tags
Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly
Public bool `json:"public,omitempty"` // Whether this assistant is shared across all teams in the platform
Share string `json:"share,omitempty"` // Assistant sharing scope (private/team)
Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable
Automated bool `json:"automated,omitempty"` // Whether this assistant is automated
Options map[string]interface{} `json:"options,omitempty"` // AI Options
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts (default prompts)
PromptPresets map[string][]Prompt `json:"prompt_presets,omitempty"` // Prompt presets organized by mode (e.g., "chat", "task", etc.)
KB *KnowledgeBase `json:"kb,omitempty"` // Knowledge base configuration
MCP *MCPServers `json:"mcp,omitempty"` // MCP servers configuration
Workflow *Workflow `json:"workflow,omitempty"` // Workflow configuration
Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder
Source string `json:"source,omitempty"` // Hook script source code
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp
// Permission management fields (not exposed in JSON API responses)
YaoCreatedBy string `json:"-"` // User who created the assistant (not exposed in JSON)

85
agent/store/xun/assistant.go
View file

@ -102,6 +102,11 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
} else {
data["path"] = nil
}
if assistant.Source != "" {
data["source"] = assistant.Source
} else {
data["source"] = nil
}
// Share field: nullable: false with default "private"
// Apply default if empty
@ -152,14 +157,15 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
// Handle interface{} fields - they should already be in the correct format
jsonFields := map[string]interface{}{
"prompts": assistant.Prompts,
"kb": assistant.KB,
"mcp": assistant.MCP,
"workflow": assistant.Workflow,
"tools": assistant.Tools,
"placeholder": assistant.Placeholder,
"locales": assistant.Locales,
"uses": assistant.Uses,
"prompts": assistant.Prompts,
"prompt_presets": assistant.PromptPresets,
"connector_options": assistant.ConnectorOptions,
"kb": assistant.KB,
"mcp": assistant.MCP,
"workflow": assistant.Workflow,
"placeholder": assistant.Placeholder,
"locales": assistant.Locales,
"uses": assistant.Uses,
}
for field, value := range jsonFields {
@ -218,14 +224,14 @@ func (conv *Xun) UpdateAssistant(assistantID string, updates map[string]interfac
data := make(map[string]interface{})
// List of fields that need JSON marshaling
jsonFields := []string{"options", "tags", "prompts", "kb", "mcp", "workflow", "tools", "placeholder", "locales", "uses"}
jsonFields := []string{"options", "tags", "prompts", "prompt_presets", "connector_options", "kb", "mcp", "workflow", "placeholder", "locales", "uses"}
jsonFieldSet := make(map[string]bool)
for _, field := range jsonFields {
jsonFieldSet[field] = true
}
// List of nullable string fields
nullableStringFields := []string{"name", "avatar", "description", "path", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id"}
nullableStringFields := []string{"name", "avatar", "description", "path", "source", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id"}
nullableFieldSet := make(map[string]bool)
for _, field := range nullableStringFields {
nullableFieldSet[field] = true
@ -418,7 +424,7 @@ func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (
// Convert rows to types.AssistantModel slice
assistants := make([]*types.AssistantModel, 0, len(rows))
jsonFields := []string{"tags", "options", "prompts", "workflow", "kb", "mcp", "tools", "placeholder", "locales", "uses"}
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "mcp", "placeholder", "locales", "uses"}
for _, row := range rows {
data := row.ToMap()
@ -456,11 +462,27 @@ func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (
}
// GetAssistant retrieves a single assistant by ID
func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
row, err := conv.query.New().
func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
qb := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
First()
Where("assistant_id", assistantID)
// Apply select fields with security validation
// If no fields specified, use default fields
fieldsToSelect := fields
if len(fieldsToSelect) == 0 {
fieldsToSelect = types.AssistantDefaultFields
}
// ValidateAssistantFields will validate fields against whitelist
sanitized := types.ValidateAssistantFields(fieldsToSelect)
selectFields := make([]interface{}, len(sanitized))
for i, field := range sanitized {
selectFields[i] = field
}
qb.Select(selectFields...)
row, err := qb.First()
if err != nil {
return nil, err
}
@ -475,7 +497,7 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
}
// Parse JSON fields
jsonFields := []string{"tags", "options", "prompts", "workflow", "kb", "mcp", "tools", "placeholder", "locales", "uses"}
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "mcp", "placeholder", "locales", "uses"}
conv.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel
@ -486,6 +508,7 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
Avatar: getString(data, "avatar"),
Connector: getString(data, "connector"),
Path: getString(data, "path"),
Source: getString(data, "source"),
BuiltIn: getBool(data, "built_in"),
Sort: getInt(data, "sort"),
Description: getString(data, "description"),
@ -529,6 +552,26 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
}
}
if promptPresets, has := data["prompt_presets"]; has && promptPresets != nil {
raw, err := jsoniter.Marshal(promptPresets)
if err == nil {
var pp map[string][]types.Prompt
if err := jsoniter.Unmarshal(raw, &pp); err == nil {
model.PromptPresets = pp
}
}
}
if connectorOptions, has := data["connector_options"]; has && connectorOptions != nil {
raw, err := jsoniter.Marshal(connectorOptions)
if err == nil {
var co types.ConnectorOptions
if err := jsoniter.Unmarshal(raw, &co); err == nil {
model.ConnectorOptions = &co
}
}
}
if kb, has := data["kb"]; has && kb != nil {
kbConverted, err := types.ToKnowledgeBase(kb)
if err == nil {
@ -550,16 +593,6 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
}
}
if tools, has := data["tools"]; has && tools != nil {
raw, err := jsoniter.Marshal(tools)
if err == nil {
var tc types.ToolCalls
if err := jsoniter.Unmarshal(raw, &tc); err == nil {
model.Tools = &tc
}
}
}
if placeholder, has := data["placeholder"]; has && placeholder != nil {
raw, err := jsoniter.Marshal(placeholder)
if err == nil {

318
agent/store/xun/assistant_test.go
View file

@ -103,8 +103,8 @@ func TestSaveAssistant(t *testing.T) {
t.Errorf("Expected ID %s, got %s", id, updatedID)
}
// Verify update
retrieved, err := store.GetAssistant(id)
// Verify update - request all fields to see the update
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve updated assistant: %v", err)
}
@ -183,8 +183,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save complex assistant: %v", err)
}
// Retrieve and verify
retrieved, err := store.GetAssistant(id)
// Retrieve and verify - request all fields for complex data
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve complex assistant: %v", err)
}
@ -237,8 +237,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save assistant with MCP: %v", err)
}
// Retrieve and verify MCP configuration
retrieved, err := store.GetAssistant(id)
// Retrieve and verify MCP configuration - mcp is in default fields
retrieved, err := store.GetAssistant(id, []string{})
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -316,8 +316,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant with MCP: %v", err)
}
// Retrieve and verify
retrieved, err := store.GetAssistant(id)
// Retrieve and verify - mcp is in default fields
retrieved, err := store.GetAssistant(id, []string{})
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -353,8 +353,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save assistant with uses: %v", err)
}
// Retrieve and verify uses configuration
retrieved, err := store.GetAssistant(id)
// Retrieve and verify uses configuration - uses is NOT in default fields, need to request all
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -396,8 +396,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save assistant without uses: %v", err)
}
// Retrieve and verify uses is nil
retrieved, err := store.GetAssistant(id)
// Retrieve and verify uses is nil - request all fields to check uses
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -425,8 +425,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save assistant with partial uses: %v", err)
}
// Retrieve and verify
retrieved, err := store.GetAssistant(id)
// Retrieve and verify - request all fields for uses
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -451,6 +451,194 @@ func TestSaveAssistant(t *testing.T) {
t.Errorf("Expected fetch to be empty, got '%s'", retrieved.Uses.Fetch)
}
})
t.Run("ConnectorOptions", func(t *testing.T) {
// Test assistant with connector options
assistant := &types.AssistantModel{
Name: "Connector Options Test",
Type: "assistant",
Connector: "openai",
Share: "private",
ConnectorOptions: &types.ConnectorOptions{
Optional: true,
Connectors: []string{"openai", "anthropic"},
Filters: []types.ModelCapability{types.CapVision, types.CapToolCalls},
},
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with connector options: %v", err)
}
// Retrieve and verify - connector_options is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.ConnectorOptions == nil {
t.Fatal("Expected connector options to be set")
}
if !retrieved.ConnectorOptions.Optional {
t.Error("Expected optional to be true")
}
if len(retrieved.ConnectorOptions.Connectors) != 2 {
t.Errorf("Expected 2 connectors, got %d", len(retrieved.ConnectorOptions.Connectors))
}
if len(retrieved.ConnectorOptions.Filters) != 2 {
t.Errorf("Expected 2 filters, got %d", len(retrieved.ConnectorOptions.Filters))
}
if retrieved.ConnectorOptions.Filters[0] != types.CapVision {
t.Errorf("Expected first filter to be vision, got '%s'", retrieved.ConnectorOptions.Filters[0])
}
t.Logf("Successfully saved and retrieved connector options for assistant %s", id)
})
t.Run("PromptPresets", func(t *testing.T) {
// Test assistant with prompt presets
assistant := &types.AssistantModel{
Name: "Prompt Presets Test",
Type: "assistant",
Connector: "openai",
Share: "private",
PromptPresets: map[string][]types.Prompt{
"chat": {
{Role: "system", Content: "You are a friendly chatbot"},
{Role: "user", Content: "Hello!"},
},
"task": {
{Role: "system", Content: "You are a task executor"},
},
},
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with prompt presets: %v", err)
}
// Retrieve and verify - prompt_presets is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.PromptPresets == nil {
t.Fatal("Expected prompt presets to be set")
}
if len(retrieved.PromptPresets) != 2 {
t.Errorf("Expected 2 preset groups, got %d", len(retrieved.PromptPresets))
}
chatPrompts, ok := retrieved.PromptPresets["chat"]
if !ok {
t.Fatal("Expected 'chat' preset to exist")
}
if len(chatPrompts) != 2 {
t.Errorf("Expected 2 chat prompts, got %d", len(chatPrompts))
}
if chatPrompts[0].Role != "system" {
t.Errorf("Expected system role, got '%s'", chatPrompts[0].Role)
}
taskPrompts, ok := retrieved.PromptPresets["task"]
if !ok {
t.Fatal("Expected 'task' preset to exist")
}
if len(taskPrompts) != 1 {
t.Errorf("Expected 1 task prompt, got %d", len(taskPrompts))
}
t.Logf("Successfully saved and retrieved prompt presets for assistant %s", id)
})
t.Run("SourceField", func(t *testing.T) {
// Test assistant with source code
sourceCode := `function onMessage(msg) {
console.log("Received:", msg);
return { status: "ok" };
}`
assistant := &types.AssistantModel{
Name: "Source Field Test",
Type: "assistant",
Connector: "openai",
Share: "private",
Source: sourceCode,
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with source: %v", err)
}
// Retrieve and verify - source is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Source != sourceCode {
t.Errorf("Expected source code to match, got '%s'", retrieved.Source)
}
t.Logf("Successfully saved and retrieved source code for assistant %s", id)
})
t.Run("AllNewFieldsTogether", func(t *testing.T) {
// Test assistant with all new fields together
assistant := &types.AssistantModel{
Name: "All New Fields Test",
Type: "assistant",
Connector: "openai",
Share: "private",
ConnectorOptions: &types.ConnectorOptions{
Optional: false,
Connectors: []string{"openai"},
Filters: []types.ModelCapability{types.CapVision},
},
PromptPresets: map[string][]types.Prompt{
"default": {
{Role: "system", Content: "Default system prompt"},
},
},
Source: "// Hook code here",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with all new fields: %v", err)
}
// Retrieve and verify all new fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.ConnectorOptions == nil {
t.Error("Expected connector options to be set")
}
if retrieved.PromptPresets == nil {
t.Error("Expected prompt presets to be set")
}
if retrieved.Source == "" {
t.Error("Expected source to be set")
}
t.Logf("Successfully saved and retrieved all new fields for assistant %s", id)
})
}
// TestDeleteAssistant tests deleting a single assistant
@ -487,7 +675,7 @@ func TestDeleteAssistant(t *testing.T) {
}
// Verify deletion
_, err = store.GetAssistant(id)
_, err = store.GetAssistant(id, nil)
if err == nil {
t.Error("Expected error when getting deleted assistant")
}
@ -534,8 +722,8 @@ func TestGetAssistant(t *testing.T) {
t.Fatalf("Failed to create assistant: %v", err)
}
// Retrieve it
retrieved, err := store.GetAssistant(id)
// Retrieve it with default fields (tags are now in default fields)
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -562,7 +750,7 @@ func TestGetAssistant(t *testing.T) {
})
t.Run("GetNonExistentAssistant", func(t *testing.T) {
_, err := store.GetAssistant("nonexistent-id")
_, err := store.GetAssistant("nonexistent-id", nil)
if err == nil {
t.Error("Expected error when getting non-existent assistant")
}
@ -1313,8 +1501,8 @@ func TestAssistantPermissionFields(t *testing.T) {
t.Fatalf("Failed to save assistant with permission fields: %v", err)
}
// Retrieve and verify
retrieved, err := store.GetAssistant(id)
// Retrieve and verify - default fields include permission fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -1361,8 +1549,8 @@ func TestAssistantPermissionFields(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify update
retrieved, err := store.GetAssistant(id)
// Verify update - default fields include permission fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get updated assistant: %v", err)
}
@ -1393,7 +1581,7 @@ func TestAssistantPermissionFields(t *testing.T) {
}
// Retrieve and verify fields are empty
retrieved, err := store.GetAssistant(id)
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -1448,7 +1636,7 @@ func TestEmptyStringAsNull(t *testing.T) {
}
// Retrieve and verify empty strings are returned (not stored as empty strings)
retrieved, err := store.GetAssistant(id)
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -1493,8 +1681,8 @@ func TestEmptyStringAsNull(t *testing.T) {
t.Fatalf("Failed to save assistant: %v", err)
}
// Retrieve and verify values are preserved
retrieved, err := store.GetAssistant(id)
// Retrieve and verify values are preserved - path is sensitive, need full fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -1576,8 +1764,8 @@ func TestGetAssistantWithLocale(t *testing.T) {
},
}
// Test English locale
retrievedEN, err := store.GetAssistant(id, "en")
// Test English locale - request all fields for placeholder
retrievedEN, err := store.GetAssistant(id, types.AssistantFullFields, "en")
if err != nil {
t.Fatalf("Failed to get assistant with EN locale: %v", err)
}
@ -1604,8 +1792,8 @@ func TestGetAssistantWithLocale(t *testing.T) {
t.Errorf("Expected first prompt 'How can I help you?', got '%s'", retrievedEN.Placeholder.Prompts[0])
}
// Test Chinese locale
retrievedZH, err := store.GetAssistant(id, "zh-cn")
// Test Chinese locale - request all fields for placeholder
retrievedZH, err := store.GetAssistant(id, types.AssistantFullFields, "zh-cn")
if err != nil {
t.Fatalf("Failed to get assistant with ZH locale: %v", err)
}
@ -1623,8 +1811,8 @@ func TestGetAssistantWithLocale(t *testing.T) {
t.Errorf("Expected placeholder title '与我聊天', got '%s'", retrievedZH.Placeholder.Title)
}
// Test without locale (should return original {{...}} values)
retrievedNoLocale, err := store.GetAssistant(id)
// Test without locale (should return original {{...}} values) - request all fields for placeholder
retrievedNoLocale, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to get assistant without locale: %v", err)
}
@ -2039,8 +2227,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify update
retrieved, err := store.GetAssistant(id)
// Verify update - need full fields to see tags
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2087,8 +2275,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify all updates
retrieved, err := store.GetAssistant(id)
// Verify all updates - use default fields (includes name, description, sort, mentionable)
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2144,8 +2332,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update JSON fields: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - need full fields for tags, options, prompts
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2193,8 +2381,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update KB and MCP: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - KB and MCP are in default fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2247,8 +2435,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update MCP: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - MCP is in default fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2325,8 +2513,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update uses: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - uses is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2361,8 +2549,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update uses again: %v", err)
}
// Verify second update
retrieved2, err := store.GetAssistant(id)
// Verify second update - uses is NOT in default fields
retrieved2, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2384,8 +2572,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to set uses to nil: %v", err)
}
// Verify uses is nil
retrieved3, err := store.GetAssistant(id)
// Verify uses is nil - uses is NOT in default fields
retrieved3, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2422,8 +2610,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update permission fields: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - permission fields are in default fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2467,8 +2655,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update with empty strings: %v", err)
}
// Verify empty strings are stored as NULL
retrieved, err := store.GetAssistant(id)
// Verify empty strings are stored as NULL - default fields include avatar, description
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2553,8 +2741,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to create assistant: %v", err)
}
// Get original updated_at
original, err := store.GetAssistant(id)
// Get original updated_at - default fields include updated_at
original, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2572,8 +2760,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Get updated assistant
updated, err := store.GetAssistant(id)
// Get updated assistant - default fields include description, updated_at
updated, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve updated assistant: %v", err)
}
@ -2607,8 +2795,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to create assistant: %v", err)
}
// Get original
original, err := store.GetAssistant(id)
// Get original - default fields
original, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2625,8 +2813,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify system fields unchanged, but name updated
retrieved, err := store.GetAssistant(id)
// Verify system fields unchanged, but name updated - default fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2693,9 +2881,9 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
t.Errorf("Expected at least 3 assistants, got %d", len(response.Data))
}
// Step 3: Update one assistant
// Step 3: Update one assistant - need full fields for tags
updatedID := assistantIDs[1]
updatedAssistant, err := store.GetAssistant(updatedID)
updatedAssistant, err := store.GetAssistant(updatedID, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to get assistant for update: %v", err)
}
@ -2708,8 +2896,8 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify update
verifyAssistant, err := store.GetAssistant(updatedID)
// Verify update - default fields include description
verifyAssistant, err := store.GetAssistant(updatedID, nil)
if err != nil {
t.Fatalf("Failed to verify update: %v", err)
}
@ -2725,7 +2913,7 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
}
// Verify deletion
_, err = store.GetAssistant(assistantIDs[0])
_, err = store.GetAssistant(assistantIDs[0], nil)
if err == nil {
t.Error("Expected error when getting deleted assistant")
}

284
data/bindata.go
View file

File diff suppressed because one or more lines are too long

19
openapi/agent/assistant.go
View file

@ -192,6 +192,17 @@ func GetAssistant(c *gin.Context) {
return
}
// Parse select fields (optional - if not provided, returns default fields)
// Query parameter: ?select=field1,field2,field3
var fields []string
if selectParam := c.Query("select"); selectParam != "" {
fields = strings.Split(selectParam, ",")
// Trim whitespace from each field
for i, field := range fields {
fields[i] = strings.TrimSpace(field)
}
}
// Parse locale (optional - if not provided, returns raw data without i18n translation)
// This is useful for form editing scenarios where you need the original values
var assistant *agenttypes.AssistantModel
@ -200,10 +211,10 @@ func GetAssistant(c *gin.Context) {
if loc := c.Query("locale"); loc != "" {
// If locale is specified, get assistant with translation
locale := strings.ToLower(strings.TrimSpace(loc))
assistant, err = agentInstance.Store.GetAssistant(assistantID, locale)
assistant, err = agentInstance.Store.GetAssistant(assistantID, fields, locale)
} else {
// If no locale specified, get raw data without translation
assistant, err = agentInstance.Store.GetAssistant(assistantID)
assistant, err = agentInstance.Store.GetAssistant(assistantID, fields)
}
if err != nil {
log.Error("Failed to get assistant %s: %v", assistantID, err)
@ -521,8 +532,8 @@ func checkAssistantPermission(authInfo *types.AuthorizedInfo, assistantID string
return false, fmt.Errorf("agent store not initialized")
}
// Get assistant from store
assistant, err := agentInstance.Store.GetAssistant(assistantID)
// Get assistant from store - only need default fields for permission check
assistant, err := agentInstance.Store.GetAssistant(assistantID, nil)
if err != nil {
return false, fmt.Errorf("assistant not found: %s", assistantID)
}

7
openapi/agent/filter.go
View file

@ -126,7 +126,7 @@ func AuthQueryFilter(c *gin.Context, authInfo *types.AuthorizedInfo) func(query.
}
// FilterBuiltInFields filters sensitive fields for built-in assistants in a list
// For built-in assistants, code-level fields (prompts, workflow, tools, kb, mcp, options) should be cleared
// For built-in assistants, code-level fields (prompts, prompt_presets, workflow, kb, mcp, options, source) should be cleared
func FilterBuiltInFields(assistants []*agenttypes.AssistantModel) {
if assistants == nil {
return
@ -138,7 +138,7 @@ func FilterBuiltInFields(assistants []*agenttypes.AssistantModel) {
}
// FilterBuiltInAssistant filters sensitive fields for a single built-in assistant
// For built-in assistants, code-level fields (prompts, workflow, tools, kb, mcp, options) should be cleared
// For built-in assistants, code-level fields (prompts, prompt_presets, workflow, kb, mcp, options, source) should be cleared
// This function can be used for both single assistant and list of assistants
func FilterBuiltInAssistant(assistant *agenttypes.AssistantModel) {
if assistant == nil {
@ -148,10 +148,11 @@ func FilterBuiltInAssistant(assistant *agenttypes.AssistantModel) {
if assistant.BuiltIn {
// Clear code-level sensitive fields for built-in assistants
assistant.Prompts = nil
assistant.PromptPresets = nil
assistant.Workflow = nil
assistant.Tools = nil
assistant.KB = nil
assistant.MCP = nil
assistant.Options = nil
assistant.Source = ""
}
}

15
openapi/agent/models.go
View file

@ -130,6 +130,17 @@ func GetModelDetails(c *gin.Context) {
return
}
// For model API, we only need minimal fields: assistant_id, name, connector, created_at, and permission fields
modelFields := []string{
"assistant_id",
"name",
"connector",
"created_at",
"built_in",
"__yao_team_id",
"__yao_created_by",
}
// Parse locale (optional - for assistant name translation)
// Priority: 1. Query parameter "locale", 2. Header "Accept-Language", 3. Metadata
locale := context.GetLocale(c, nil)
@ -138,9 +149,9 @@ func GetModelDetails(c *gin.Context) {
var err error
if locale != "" {
assistant, err = agentInstance.Store.GetAssistant(assistantID, locale)
assistant, err = agentInstance.Store.GetAssistant(assistantID, modelFields, locale)
} else {
assistant, err = agentInstance.Store.GetAssistant(assistantID)
assistant, err = agentInstance.Store.GetAssistant(assistantID, modelFields)
}
if err != nil {

64
openapi/tests/agent/assistant_update_test.go
View file

@ -620,62 +620,7 @@ func TestUpdateAssistant(t *testing.T) {
t.Logf("Successfully updated assistant mcp settings: %s", assistantID)
})
t.Run("UpdateAssistantTools", func(t *testing.T) {
// Create a test assistant
assistantID := createTestAssistant("Tools Update Test")
defer func() {
deleteReq, _ := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, _ := http.DefaultClient.Do(deleteReq)
if resp != nil {
resp.Body.Close()
}
}()
// Update tools
updateData := map[string]interface{}{
"tools": []map[string]interface{}{
{
"name": "web_search",
"description": "Search the web for information",
"parameters": map[string]interface{}{
"query": map[string]interface{}{
"type": "string",
"description": "Search query",
"required": true,
},
},
},
{
"name": "calculator",
"description": "Perform calculations",
"parameters": map[string]interface{}{
"expression": map[string]interface{}{
"type": "string",
"description": "Mathematical expression",
"required": true,
},
},
},
},
}
jsonData, err := json.Marshal(updateData)
assert.NoError(t, err)
req, err := http.NewRequest("PUT", serverURL+baseURL+"/agent/assistants/"+assistantID, bytes.NewBuffer(jsonData))
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully update tools")
t.Logf("Successfully updated assistant tools: %s", assistantID)
})
// Note: UpdateAssistantTools test removed - tools field is deprecated and replaced by MCP
t.Run("UpdateAssistantWorkflow", func(t *testing.T) {
// Create a test assistant
@ -776,12 +721,7 @@ func TestUpdateAssistant(t *testing.T) {
},
},
},
"tools": []map[string]interface{}{
{
"name": "updated_tool",
"description": "Updated tool description",
},
},
// Note: tools field removed - now handled by MCP
"kb": map[string]interface{}{
"collections": []string{"updated-collection"},
"enabled": true,

24
yao/models/agent/assistant.mod.yao
View file

@ -57,6 +57,13 @@
"length": 200,
"nullable": false
},
{
"name": "connector_options",
"type": "json",
"label": "Connector Options",
"comment": "Connector selection options: optional flag, available connectors list, and capability filters",
"nullable": true
},
{
"name": "description",
"type": "string",
@ -108,7 +115,14 @@
"name": "prompts",
"type": "json",
"label": "Prompts",
"comment": "Assistant prompts",
"comment": "Assistant default prompts",
"nullable": true
},
{
"name": "prompt_presets",
"type": "json",
"label": "Prompt Presets",
"comment": "Prompt presets organized by mode (e.g., chat, task, etc.)",
"nullable": true
},
{
@ -133,10 +147,10 @@
"nullable": true
},
{
"name": "tools",
"type": "json",
"label": "Tools",
"comment": "Assistant tools",
"name": "source",
"type": "text",
"label": "Source",
"comment": "Hook script source code",
"nullable": true
},
{