Merge pull request #1364 from trheyi/main

Add database support to Assistant model and related functionalities

agent/assistant/assistant.go

@@ -82,9 +82,12 @@ func (ast *Assistant) Map() map[string]interface{} {
 		"disable_global_prompts": ast.DisableGlobalPrompts,
 		"source":                 ast.Source,
 		"kb":                     ast.KB,
+		"db":                     ast.DB,
 		"mcp":                    ast.MCP,
 		"workflow":               ast.Workflow,
 		"tags":                   ast.Tags,
+		"modes":                  ast.Modes,
+		"default_mode":           ast.DefaultMode,
 		"mentionable":            ast.Mentionable,
 		"automated":              ast.Automated,
 		"placeholder":            ast.Placeholder,
@@ -168,6 +171,15 @@ func (ast *Assistant) Clone() *Assistant {
 		copy(clone.Tags, ast.Tags)
 	}
 
+	// Deep copy modes
+	if ast.Modes != nil {
+		clone.Modes = make([]string, len(ast.Modes))
+		copy(clone.Modes, ast.Modes)
+	}
+
+	// Copy default_mode (simple string)
+	clone.DefaultMode = ast.DefaultMode
+
 	// Deep copy KB
 	if ast.KB != nil {
 		clone.KB = &store.KnowledgeBase{}
@@ -183,6 +195,21 @@ func (ast *Assistant) Clone() *Assistant {
 		}
 	}
 
+	// Deep copy DB
+	if ast.DB != nil {
+		clone.DB = &store.Database{}
+		if ast.DB.Models != nil {
+			clone.DB.Models = make([]string, len(ast.DB.Models))
+			copy(clone.DB.Models, ast.DB.Models)
+		}
+		if ast.DB.Options != nil {
+			clone.DB.Options = make(map[string]interface{})
+			for k, v := range ast.DB.Options {
+				clone.DB.Options[k] = v
+			}
+		}
+	}
+
 	// Deep copy MCP
 	if ast.MCP != nil {
 		clone.MCP = &store.MCPServers{}
@@ -347,6 +374,12 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
 	if v, ok := data["tags"].([]string); ok {
 		ast.Tags = v
 	}
+	if v, ok := data["modes"].([]string); ok {
+		ast.Modes = v
+	}
+	if v, ok := data["default_mode"].(string); ok {
+		ast.DefaultMode = v
+	}
 	if v, ok := data["options"].(map[string]interface{}); ok {
 		ast.Options = v
 	}
@@ -381,6 +414,15 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
 		ast.KB = kb
 	}
 
+	// DB
+	if v, has := data["db"]; has {
+		db, err := store.ToDatabase(v)
+		if err != nil {
+			return err
+		}
+		ast.DB = db
+	}
+
 	// MCP
 	if v, has := data["mcp"]; has {
 		mcp, err := store.ToMCPServers(v)

agent/store/types/convert.go

@@ -50,6 +50,44 @@ func ToKnowledgeBase(v interface{}) (*KnowledgeBase, error) {
 	}
 }
 
+// ToDatabase converts various types to Database
+func ToDatabase(v interface{}) (*Database, error) {
+	if v == nil {
+		return nil, nil
+	}
+
+	switch db := v.(type) {
+	case *Database:
+		return db, nil
+
+	case Database:
+		return &db, nil
+
+	case []string:
+		return &Database{Models: db}, nil
+
+	case []interface{}:
+		var models []string
+		for _, item := range db {
+			models = append(models, cast.ToString(item))
+		}
+		return &Database{Models: models}, nil
+
+	default:
+		raw, err := jsoniter.Marshal(db)
+		if err != nil {
+			return nil, fmt.Errorf("db format error: %s", err.Error())
+		}
+
+		var database Database
+		err = jsoniter.Unmarshal(raw, &database)
+		if err != nil {
+			return nil, fmt.Errorf("db format error: %s", err.Error())
+		}
+		return &database, nil
+	}
+}
+
 // ToMCPServers converts various types to MCPServers
 func ToMCPServers(v interface{}) (*MCPServers, error) {
 	if v == nil {
@@ -264,6 +302,22 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
 		}
 	}
 
+	// Modes (string array)
+	if modes, ok := data["modes"]; ok && modes != nil {
+		raw, err := jsoniter.Marshal(modes)
+		if err == nil {
+			var m []string
+			if err := jsoniter.Unmarshal(raw, &m); err == nil {
+				model.Modes = m
+			}
+		}
+	}
+
+	// DefaultMode (string)
+	if defaultMode, ok := data["default_mode"].(string); ok {
+		model.DefaultMode = defaultMode
+	}
+
 	// Options (map)
 	if options, ok := data["options"].(map[string]interface{}); ok {
 		model.Options = options
@@ -313,6 +367,14 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
 		}
 	}
 
+	// DB
+	if db, ok := data["db"]; ok && db != nil {
+		dbConverted, err := ToDatabase(db)
+		if err == nil {
+			model.DB = dbConverted
+		}
+	}
+
 	// MCP
 	if mcp, ok := data["mcp"]; ok && mcp != nil {
 		mcpConverted, err := ToMCPServers(mcp)

agent/store/types/convert_test.go

@@ -5,6 +5,106 @@ import (
 	"time"
 )
 
+// TestToDatabase tests the ToDatabase conversion function
+func TestToDatabase(t *testing.T) {
+	t.Run("NilInput", func(t *testing.T) {
+		result, err := ToDatabase(nil)
+		if err != nil {
+			t.Errorf("Expected no error, got: %v", err)
+		}
+		if result != nil {
+			t.Errorf("Expected nil result, got: %v", result)
+		}
+	})
+
+	t.Run("DatabasePointer", func(t *testing.T) {
+		db := &Database{Models: []string{"model1", "model2"}}
+		result, err := ToDatabase(db)
+		if err != nil {
+			t.Errorf("Expected no error, got: %v", err)
+		}
+		if result != db {
+			t.Errorf("Expected same pointer")
+		}
+	})
+
+	t.Run("DatabaseValue", func(t *testing.T) {
+		db := Database{Models: []string{"model1", "model2"}}
+		result, err := ToDatabase(db)
+		if err != nil {
+			t.Errorf("Expected no error, got: %v", err)
+		}
+		if len(result.Models) != 2 {
+			t.Errorf("Expected 2 models, got %d", len(result.Models))
+		}
+	})
+
+	t.Run("StringSlice", func(t *testing.T) {
+		models := []string{"model1", "model2", "model3"}
+		result, err := ToDatabase(models)
+		if err != nil {
+			t.Errorf("Expected no error, got: %v", err)
+		}
+		if len(result.Models) != 3 {
+			t.Errorf("Expected 3 models, got %d", len(result.Models))
+		}
+		if result.Models[0] != "model1" {
+			t.Errorf("Expected 'model1', got '%s'", result.Models[0])
+		}
+	})
+
+	t.Run("InterfaceSlice", func(t *testing.T) {
+		models := []interface{}{"model1", "model2", 123}
+		result, err := ToDatabase(models)
+		if err != nil {
+			t.Errorf("Expected no error, got: %v", err)
+		}
+		if len(result.Models) != 3 {
+			t.Errorf("Expected 3 models, got %d", len(result.Models))
+		}
+		if result.Models[2] != "123" {
+			t.Errorf("Expected '123', got '%s'", result.Models[2])
+		}
+	})
+
+	t.Run("MapInput", func(t *testing.T) {
+		data := map[string]interface{}{
+			"models": []string{"model1", "model2"},
+		}
+		result, err := ToDatabase(data)
+		if err != nil {
+			t.Errorf("Expected no error, got: %v", err)
+		}
+		if len(result.Models) != 2 {
+			t.Errorf("Expected 2 models, got %d", len(result.Models))
+		}
+	})
+
+	t.Run("InvalidInput", func(t *testing.T) {
+		// Test with data that can't be marshaled
+		invalidData := make(chan int)
+		_, err := ToDatabase(invalidData)
+		if err == nil {
+			t.Error("Expected error for invalid input")
+		}
+	})
+
+	t.Run("InvalidJSONUnmarshal", func(t *testing.T) {
+		// Test with data that marshals but can't unmarshal to Database
+		data := map[string]interface{}{
+			"invalid_field": "should cause unmarshal to fail gracefully",
+		}
+		result, err := ToDatabase(data)
+		// Should not error, just return empty Database
+		if err != nil {
+			t.Errorf("Expected no error, got: %v", err)
+		}
+		if result == nil {
+			t.Error("Expected non-nil result")
+		}
+	})
+}
+
 // TestToKnowledgeBase tests the ToKnowledgeBase conversion function
 func TestToKnowledgeBase(t *testing.T) {
 	t.Run("NilInput", func(t *testing.T) {
@@ -433,18 +533,20 @@ func TestToAssistantModel(t *testing.T) {
 				"connectors": []string{"openai", "anthropic"},
 				"filters":    []string{"vision", "tool_calls"},
 			},
-			"path":        "/path/to/assistant",
+			"path":         "/path/to/assistant",
-			"description": "Test description",
+			"description":  "Test description",
-			"share":       "team",
+			"share":        "team",
-			"built_in":    true,
+			"built_in":     true,
-			"readonly":    false,
+			"readonly":     false,
-			"public":      true,
+			"public":       true,
-			"mentionable": true,
+			"mentionable":  true,
-			"automated":   false,
+			"automated":    false,
-			"sort":        100,
+			"sort":         100,
-			"created_at":  int64(1609459200),
+			"created_at":   int64(1609459200),
-			"updated_at":  int64(1609459300),
+			"updated_at":   int64(1609459300),
-			"tags":        []string{"tag1", "tag2"},
+			"tags":         []string{"tag1", "tag2"},
+			"modes":        []string{"chat", "task"},
+			"default_mode": "chat",
 			"options": map[string]interface{}{
 				"temperature": 0.7,
 			},
@@ -464,6 +566,9 @@ func TestToAssistantModel(t *testing.T) {
 			"kb": map[string]interface{}{
 				"collections": []string{"col1"},
 			},
+			"db": map[string]interface{}{
+				"models": []string{"model1"},
+			},
 			"mcp": map[string]interface{}{
 				"servers": []string{"server1"},
 			},
@@ -553,6 +658,15 @@ func TestToAssistantModel(t *testing.T) {
 		if len(result.Tags) != 2 {
 			t.Errorf("Expected 2 tags, got %d", len(result.Tags))
 		}
+		if len(result.Modes) != 2 {
+			t.Errorf("Expected 2 modes, got %d", len(result.Modes))
+		}
+		if result.Modes[0] != "chat" {
+			t.Errorf("Expected first mode 'chat', got '%s'", result.Modes[0])
+		}
+		if result.DefaultMode != "chat" {
+			t.Errorf("Expected default_mode 'chat', got '%s'", result.DefaultMode)
+		}
 		if result.Options == nil {
 			t.Error("Expected Options to be set")
 		}
@@ -582,6 +696,9 @@ func TestToAssistantModel(t *testing.T) {
 		if result.KB == nil {
 			t.Error("Expected KB to be set")
 		}
+		if result.DB == nil {
+			t.Error("Expected DB to be set")
+		}
 		if result.MCP == nil {
 			t.Error("Expected MCP to be set")
 		}
@@ -623,9 +740,12 @@ func TestToAssistantModel(t *testing.T) {
 		data := map[string]interface{}{
 			"assistant_id": "test-id",
 			"tags":         nil,
+			"modes":        nil,
+			"default_mode": "",
 			"options":      nil,
 			"prompts":      nil,
 			"kb":           nil,
+			"db":           nil,
 			"mcp":          nil,
 			"workflow":     nil,
 			"placeholder":  nil,
@@ -644,6 +764,12 @@ func TestToAssistantModel(t *testing.T) {
 		if result.Tags != nil {
 			t.Error("Expected Tags to be nil")
 		}
+		if result.Modes != nil {
+			t.Error("Expected Modes to be nil")
+		}
+		if result.DefaultMode != "" {
+			t.Error("Expected DefaultMode to be empty")
+		}
 		if result.Options != nil {
 			t.Error("Expected Options to be nil")
 		}

agent/store/types/fields.go

@@ -22,9 +22,12 @@ var AssistantAllowedFields = map[string]bool{
 	"disable_global_prompts": true,
 	"workflow":               true,
 	"kb":                     true,
+	"db":                     true,
 	"mcp":                    true,
 	"source":                 true,
 	"tags":                   true,
+	"modes":                  true,
+	"default_mode":           true,
 	"readonly":               true,
 	"public":                 true,
 	"share":                  true,
@@ -49,7 +52,9 @@ var AssistantDefaultFields = []string{
 	"avatar",
 	"connector",
 	"description",
-	"tags", // Tags for categorization (lightweight)
+	"tags",         // Tags for categorization (lightweight)
+	"modes",        // Supported modes (lightweight)
+	"default_mode", // Default mode (lightweight)
 	"sort",
 	"built_in",
 	"readonly",
@@ -58,6 +63,7 @@ var AssistantDefaultFields = []string{
 	"automated",
 	"mentionable",
 	"kb",  // Knowledge base configuration (lightweight)
+	"db",  // Database configuration (lightweight)
 	"mcp", // MCP servers configuration (lightweight)
 	"created_at",
 	"updated_at",
@@ -87,9 +93,12 @@ var AssistantFullFields = []string{
 	"disable_global_prompts",
 	"workflow",
 	"kb",
+	"db",
 	"mcp",
 	"source",
 	"tags",
+	"modes",
+	"default_mode",
 	"readonly",
 	"public",
 	"share",

agent/store/types/fields_test.go

@@ -122,12 +122,15 @@ func TestAssistantAllowedFields(t *testing.T) {
 			"disable_global_prompts",
 			"workflow",
 			"kb",
+			"db",
 			"mcp",
 			"placeholder",
 			"locales",
 			"uses",
 			"connector_options",
 			"source",
+			"modes",
+			"default_mode",
 		}
 		for _, field := range complexFields {
 			if !AssistantAllowedFields[field] {
@@ -144,7 +147,10 @@ func TestAssistantDefaultFields(t *testing.T) {
 			"name",
 			"type",
 			"kb",               // Knowledge base is essential for assistant functionality
+			"db",               // Database is essential for assistant functionality
 			"mcp",              // MCP servers are essential for assistant functionality
+			"modes",            // Supported modes are essential for mode filtering
+			"default_mode",     // Default mode is essential for mode selection
 			"__yao_created_by", // Permission fields are essential for access control
 			"__yao_updated_by",
 			"__yao_team_id",
@@ -165,7 +171,7 @@ 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
+		// Note: kb, db, mcp, tags, modes, and default_mode are lightweight and included in defaults
 		sensitiveFields := []string{
 			"options",
 			"prompts",
@@ -228,12 +234,15 @@ func TestAssistantFullFields(t *testing.T) {
 			"disable_global_prompts",
 			"workflow",
 			"kb",
+			"db",
 			"mcp",
 			"placeholder",
 			"locales",
 			"uses",
 			"connector_options",
 			"source",
+			"modes",
+			"default_mode",
 		}
 
 		fullFieldsMap := make(map[string]bool)

agent/store/types/types.go

@@ -101,6 +101,12 @@ type KnowledgeBase struct {
 	Options     map[string]interface{} `json:"options,omitempty"`     // Additional options for knowledge base
 }
 
+// Database the database configuration
+type Database struct {
+	Models  []string               `json:"models,omitempty"`  // Database models
+	Options map[string]interface{} `json:"options,omitempty"` // Additional options for database
+}
+
 // MCPServers the MCP servers configuration
 // Supports multiple formats in the servers array:
 // - Simple string: "server_id"
@@ -257,6 +263,8 @@ type AssistantModel struct {
 	Sort                 int                    `json:"sort,omitempty"`                   // Assistant Sort
 	Description          string                 `json:"description,omitempty"`            // Assistant Description
 	Tags                 []string               `json:"tags,omitempty"`                   // Assistant Tags
+	Modes                []string               `json:"modes,omitempty"`                  // Supported modes (e.g., ["task", "chat"]), null means all modes are supported
+	DefaultMode          string                 `json:"default_mode,omitempty"`           // Default mode, can be empty
 	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)
@@ -267,6 +275,7 @@ type AssistantModel struct {
 	PromptPresets        map[string][]Prompt    `json:"prompt_presets,omitempty"`         // Prompt presets organized by mode (e.g., "chat", "task", etc.)
 	DisableGlobalPrompts bool                   `json:"disable_global_prompts,omitempty"` // Whether to disable global prompts, default is false
 	KB                   *KnowledgeBase         `json:"kb,omitempty"`                     // Knowledge base configuration
+	DB                   *Database              `json:"db,omitempty"`                     // Database configuration
 	MCP                  *MCPServers            `json:"mcp,omitempty"`                    // MCP servers configuration
 	Workflow             *Workflow              `json:"workflow,omitempty"`               // Workflow configuration
 	Placeholder          *Placeholder           `json:"placeholder,omitempty"`            // Assistant Placeholder

agent/store/xun/assistant.go

@@ -156,12 +156,28 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
 		data["tags"] = jsonStr
 	}
 
+	if assistant.Modes != nil {
+		jsonStr, err := jsoniter.MarshalToString(assistant.Modes)
+		if err != nil {
+			return "", fmt.Errorf("failed to marshal modes: %w", err)
+		}
+		data["modes"] = jsonStr
+	}
+
+	// DefaultMode is a simple string field
+	if assistant.DefaultMode != "" {
+		data["default_mode"] = assistant.DefaultMode
+	} else {
+		data["default_mode"] = nil
+	}
+
 	// Handle interface{} fields - they should already be in the correct format
 	jsonFields := map[string]interface{}{
 		"prompts":           assistant.Prompts,
 		"prompt_presets":    assistant.PromptPresets,
 		"connector_options": assistant.ConnectorOptions,
 		"kb":                assistant.KB,
+		"db":                assistant.DB,
 		"mcp":               assistant.MCP,
 		"workflow":          assistant.Workflow,
 		"placeholder":       assistant.Placeholder,
@@ -225,14 +241,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", "prompt_presets", "connector_options", "kb", "mcp", "workflow", "placeholder", "locales", "uses"}
+	jsonFields := []string{"options", "tags", "modes", "prompts", "prompt_presets", "connector_options", "kb", "db", "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", "source", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id"}
+	nullableStringFields := []string{"name", "avatar", "description", "path", "source", "default_mode", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id"}
 	nullableFieldSet := make(map[string]bool)
 	for _, field := range nullableStringFields {
 		nullableFieldSet[field] = true
@@ -498,7 +514,7 @@ func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...str
 	}
 
 	// Parse JSON fields
-	jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "mcp", "placeholder", "locales", "uses"}
+	jsonFields := []string{"tags", "modes", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "db", "mcp", "placeholder", "locales", "uses"}
 	conv.parseJSONFields(data, jsonFields)
 
 	// Convert map to types.AssistantModel
@@ -513,6 +529,7 @@ func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...str
 		BuiltIn:              getBool(data, "built_in"),
 		Sort:                 getInt(data, "sort"),
 		Description:          getString(data, "description"),
+		DefaultMode:          getString(data, "default_mode"),
 		Readonly:             getBool(data, "readonly"),
 		Public:               getBool(data, "public"),
 		Share:                getString(data, "share"),
@@ -537,6 +554,16 @@ func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...str
 		}
 	}
 
+	// Handle Modes
+	if modes, ok := data["modes"].([]interface{}); ok {
+		model.Modes = make([]string, len(modes))
+		for i, mode := range modes {
+			if s, ok := mode.(string); ok {
+				model.Modes[i] = s
+			}
+		}
+	}
+
 	// Handle Options
 	if options, ok := data["options"].(map[string]interface{}); ok {
 		model.Options = options
@@ -581,6 +608,13 @@ func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...str
 		}
 	}
 
+	if db, has := data["db"]; has && db != nil {
+		dbConverted, err := types.ToDatabase(db)
+		if err == nil {
+			model.DB = dbConverted
+		}
+	}
+
 	if mcp, has := data["mcp"]; has && mcp != nil {
 		mcpConverted, err := types.ToMCPServers(mcp)
 		if err == nil {

agent/store/xun/assistant_test.go

@@ -2398,6 +2398,196 @@ func TestUpdateAssistant(t *testing.T) {
 		}
 	})
 
+	t.Run("UpdateKBDBAndMCP", func(t *testing.T) {
+		// Create assistant
+		assistant := &types.AssistantModel{
+			Name:      "KB DB MCP Test",
+			Type:      "assistant",
+			Connector: "openai",
+			Share:     "private",
+		}
+
+		id, err := store.SaveAssistant(assistant)
+		if err != nil {
+			t.Fatalf("Failed to create assistant: %v", err)
+		}
+
+		// Update KB, DB and MCP
+		updates := map[string]interface{}{
+			"kb": map[string]interface{}{
+				"collections": []string{"collection1", "collection2"},
+			},
+			"db": map[string]interface{}{
+				"models": []string{"model1", "model2"},
+			},
+			"mcp": map[string]interface{}{
+				"servers": []string{"server1", "server2"},
+			},
+		}
+
+		err = store.UpdateAssistant(id, updates)
+		if err != nil {
+			t.Fatalf("Failed to update KB, DB and MCP: %v", err)
+		}
+
+		// Verify updates - KB, DB and MCP are in default fields
+		retrieved, err := store.GetAssistant(id, nil)
+		if err != nil {
+			t.Fatalf("Failed to retrieve assistant: %v", err)
+		}
+
+		if retrieved.KB == nil || len(retrieved.KB.Collections) != 2 {
+			t.Errorf("Expected 2 KB collections, got %v", retrieved.KB)
+		}
+		if retrieved.DB == nil || len(retrieved.DB.Models) != 2 {
+			t.Errorf("Expected 2 DB models, got %v", retrieved.DB)
+		}
+		if retrieved.DB.Models[0] != "model1" {
+			t.Errorf("Expected first model 'model1', got '%s'", retrieved.DB.Models[0])
+		}
+		if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 2 {
+			t.Errorf("Expected 2 MCP servers, got %v", retrieved.MCP)
+		}
+		if retrieved.MCP.Servers[0].ServerID != "server1" {
+			t.Errorf("Expected first server 'server1', got '%s'", retrieved.MCP.Servers[0].ServerID)
+		}
+	})
+
+	t.Run("UpdateDBWithOptions", func(t *testing.T) {
+		// Create assistant
+		assistant := &types.AssistantModel{
+			Name:      "DB Advanced Test",
+			Type:      "assistant",
+			Connector: "openai",
+			Share:     "private",
+		}
+
+		id, err := store.SaveAssistant(assistant)
+		if err != nil {
+			t.Fatalf("Failed to create assistant: %v", err)
+		}
+
+		// Update with DB using advanced configuration
+		updates := map[string]interface{}{
+			"db": map[string]interface{}{
+				"models": []string{"user", "product", "order"},
+				"options": map[string]interface{}{
+					"limit":  100,
+					"offset": 0,
+				},
+			},
+		}
+
+		err = store.UpdateAssistant(id, updates)
+		if err != nil {
+			t.Fatalf("Failed to update DB: %v", err)
+		}
+
+		// Verify updates - DB is in default fields
+		retrieved, err := store.GetAssistant(id, nil)
+		if err != nil {
+			t.Fatalf("Failed to retrieve assistant: %v", err)
+		}
+
+		if retrieved.DB == nil {
+			t.Fatal("Expected DB to be set")
+		}
+		if len(retrieved.DB.Models) != 3 {
+			t.Errorf("Expected 3 DB models, got %d", len(retrieved.DB.Models))
+		}
+		if retrieved.DB.Models[0] != "user" {
+			t.Errorf("Expected first model 'user', got '%s'", retrieved.DB.Models[0])
+		}
+		if retrieved.DB.Options == nil {
+			t.Error("Expected DB options to be set")
+		} else {
+			if limit, ok := retrieved.DB.Options["limit"].(float64); !ok || limit != 100 {
+				t.Errorf("Expected DB limit 100, got %v", retrieved.DB.Options["limit"])
+			}
+		}
+	})
+
+	t.Run("UpdateModesAndDefaultMode", func(t *testing.T) {
+		// Create assistant
+		assistant := &types.AssistantModel{
+			Name:      "Modes Test",
+			Type:      "assistant",
+			Connector: "openai",
+			Share:     "private",
+		}
+
+		id, err := store.SaveAssistant(assistant)
+		if err != nil {
+			t.Fatalf("Failed to create assistant: %v", err)
+		}
+
+		// Update with modes and default_mode
+		updates := map[string]interface{}{
+			"modes":        []string{"chat", "task", "analyze"},
+			"default_mode": "chat",
+		}
+
+		err = store.UpdateAssistant(id, updates)
+		if err != nil {
+			t.Fatalf("Failed to update modes: %v", err)
+		}
+
+		// Verify updates - modes and default_mode are in default fields
+		retrieved, err := store.GetAssistant(id, nil)
+		if err != nil {
+			t.Fatalf("Failed to retrieve assistant: %v", err)
+		}
+
+		if retrieved.Modes == nil || len(retrieved.Modes) != 3 {
+			t.Errorf("Expected 3 modes, got %v", retrieved.Modes)
+		}
+		if retrieved.Modes[0] != "chat" {
+			t.Errorf("Expected first mode 'chat', got '%s'", retrieved.Modes[0])
+		}
+		if retrieved.DefaultMode != "chat" {
+			t.Errorf("Expected default_mode 'chat', got '%s'", retrieved.DefaultMode)
+		}
+	})
+
+	t.Run("UpdateModesOnly", func(t *testing.T) {
+		// Create assistant with default_mode
+		assistant := &types.AssistantModel{
+			Name:        "Modes Only Test",
+			Type:        "assistant",
+			Connector:   "openai",
+			Share:       "private",
+			DefaultMode: "task",
+		}
+
+		id, err := store.SaveAssistant(assistant)
+		if err != nil {
+			t.Fatalf("Failed to create assistant: %v", err)
+		}
+
+		// Update only modes
+		updates := map[string]interface{}{
+			"modes": []string{"chat", "task"},
+		}
+
+		err = store.UpdateAssistant(id, updates)
+		if err != nil {
+			t.Fatalf("Failed to update modes: %v", err)
+		}
+
+		// Verify updates - default_mode should remain unchanged
+		retrieved, err := store.GetAssistant(id, nil)
+		if err != nil {
+			t.Fatalf("Failed to retrieve assistant: %v", err)
+		}
+
+		if len(retrieved.Modes) != 2 {
+			t.Errorf("Expected 2 modes, got %d", len(retrieved.Modes))
+		}
+		if retrieved.DefaultMode != "task" {
+			t.Errorf("Expected default_mode to remain 'task', got '%s'", retrieved.DefaultMode)
+		}
+	})
+
 	t.Run("UpdateMCPWithToolsAndResources", func(t *testing.T) {
 		// Create assistant
 		assistant := &types.AssistantModel{

yao/models/agent/assistant.mod.yao

@@ -147,6 +147,13 @@
       "comment": "Assistant knowledge base collections",
       "nullable": true
     },
+    {
+      "name": "db",
+      "type": "json",
+      "label": "Database",
+      "comment": "Assistant database models",
+      "nullable": true
+    },
     {
       "name": "mcp",
       "type": "json",
@@ -168,6 +175,21 @@
       "comment": "Assistant tags",
       "nullable": true
     },
+    {
+      "name": "modes",
+      "type": "json",
+      "label": "Modes",
+      "comment": "Supported modes (e.g., chat, task), null means all modes are supported",
+      "nullable": true
+    },
+    {
+      "name": "default_mode",
+      "type": "string",
+      "label": "Default Mode",
+      "comment": "Default mode for the assistant",
+      "length": 50,
+      "nullable": true
+    },
     {
       "name": "readonly",
       "type": "boolean",