Add database support to Assistant model and related functionalities

- Introduced a new `Database` type to encapsulate database models and options.
- Updated the `Assistant` model to include a `DB` field for database configuration.
- Enhanced the `Map`, `Clone`, and `Update` methods to handle the new `DB` field appropriately.
- Implemented the `ToDatabase` conversion function to support various input types for database initialization.
- Added comprehensive tests to validate the integration of the database functionality within the assistant model.
- Updated relevant fields and methods to ensure consistent handling of the new database configuration.

agent/assistant/assistant.go

@@ -82,6 +82,7 @@ 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,
@@ -183,6 +184,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{}
@@ -381,6 +397,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 {
@@ -313,6 +351,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) {
@@ -464,6 +564,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"},
 			},
@@ -582,6 +685,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")
 		}
@@ -626,6 +732,7 @@ func TestToAssistantModel(t *testing.T) {
 			"options":      nil,
 			"prompts":      nil,
 			"kb":           nil,
+			"db":           nil,
 			"mcp":          nil,
 			"workflow":     nil,
 			"placeholder":  nil,

agent/store/types/fields.go

@@ -22,6 +22,7 @@ var AssistantAllowedFields = map[string]bool{
 	"disable_global_prompts": true,
 	"workflow":               true,
 	"kb":                     true,
+	"db":                     true,
 	"mcp":                    true,
 	"source":                 true,
 	"tags":                   true,
@@ -58,6 +59,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,6 +89,7 @@ var AssistantFullFields = []string{
 	"disable_global_prompts",
 	"workflow",
 	"kb",
+	"db",
 	"mcp",
 	"source",
 	"tags",

agent/store/types/fields_test.go

@@ -122,6 +122,7 @@ func TestAssistantAllowedFields(t *testing.T) {
 			"disable_global_prompts",
 			"workflow",
 			"kb",
+			"db",
 			"mcp",
 			"placeholder",
 			"locales",
@@ -144,6 +145,7 @@ 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
 			"__yao_created_by", // Permission fields are essential for access control
 			"__yao_updated_by",
@@ -165,7 +167,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, and tags are lightweight and included in defaults
 		sensitiveFields := []string{
 			"options",
 			"prompts",
@@ -228,6 +230,7 @@ func TestAssistantFullFields(t *testing.T) {
 			"disable_global_prompts",
 			"workflow",
 			"kb",
+			"db",
 			"mcp",
 			"placeholder",
 			"locales",

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"
@@ -267,6 +273,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

@@ -162,6 +162,7 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
 		"prompt_presets":    assistant.PromptPresets,
 		"connector_options": assistant.ConnectorOptions,
 		"kb":                assistant.KB,
+		"db":                assistant.DB,
 		"mcp":               assistant.MCP,
 		"workflow":          assistant.Workflow,
 		"placeholder":       assistant.Placeholder,
@@ -225,7 +226,7 @@ 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", "prompts", "prompt_presets", "connector_options", "kb", "db", "mcp", "workflow", "placeholder", "locales", "uses"}
 	jsonFieldSet := make(map[string]bool)
 	for _, field := range jsonFields {
 		jsonFieldSet[field] = true
@@ -498,7 +499,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", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "db", "mcp", "placeholder", "locales", "uses"}
 	conv.parseJSONFields(data, jsonFields)
 
 	// Convert map to types.AssistantModel
@@ -581,6 +582,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,115 @@ 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("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",