Add modes and database support to Assistant model

- Enhanced the `loadMap` function to include handling for `modes` and `default_mode` fields, allowing for flexible operational modes and a specified primary mode.
- Introduced a new `DB` field in the Assistant model to support database configuration, improving data management capabilities.
- Implemented the `ToModes` conversion function to facilitate various input types for modes, ensuring robust handling and validation.
- Added comprehensive tests for the `ToModes` function to validate its functionality across different input scenarios, enhancing overall reliability.
- Updated relevant methods to ensure consistent integration of the new fields and functionalities within the Assistant model.
This commit is contained in:
Max 2025-12-03 10:50:24 +08:00
parent 8d762964b2
commit bbba6b9fdb
3 changed files with 178 additions and 0 deletions

View file

@ -425,6 +425,20 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.Automated = v
}
// modes
if v, has := data["modes"]; has {
modes, err := store.ToModes(v)
if err != nil {
return nil, err
}
assistant.Modes = modes
}
// default_mode
if v, ok := data["default_mode"].(string); ok {
assistant.DefaultMode = v
}
// DisableGlobalPrompts
if v, ok := data["disable_global_prompts"].(bool); ok {
assistant.DisableGlobalPrompts = v
@ -593,6 +607,15 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.KB = knowledgeBase
}
// db
if db, has := data["db"]; has {
database, err := store.ToDatabase(db)
if err != nil {
return nil, err
}
assistant.DB = database
}
// mcp
if mcp, has := data["mcp"]; has {
mcpServers, err := store.ToMCPServers(mcp)

View file

@ -556,6 +556,42 @@ func ToConnectorOptions(v interface{}) (*ConnectorOptions, error) {
}
}
// ToModes converts various types to []string for modes
func ToModes(v interface{}) ([]string, error) {
if v == nil {
return nil, nil
}
switch modes := v.(type) {
case []string:
return modes, nil
case []interface{}:
var result []string
for _, item := range modes {
result = append(result, cast.ToString(item))
}
return result, nil
case string:
// Single string becomes a slice with one element
return []string{modes}, nil
default:
raw, err := jsoniter.Marshal(modes)
if err != nil {
return nil, fmt.Errorf("modes format error: %s", err.Error())
}
var result []string
err = jsoniter.Unmarshal(raw, &result)
if err != nil {
return nil, fmt.Errorf("modes format error: %s", err.Error())
}
return result, nil
}
}
// ToPromptPresets converts various types to map[string][]Prompt
func ToPromptPresets(v interface{}) (map[string][]Prompt, error) {
if v == nil {

View file

@ -1449,6 +1449,125 @@ func TestToConnectorOptions(t *testing.T) {
})
}
// TestToModes tests the ToModes conversion function
func TestToModes(t *testing.T) {
t.Run("NilInput", func(t *testing.T) {
result, err := ToModes(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("StringSlice", func(t *testing.T) {
modes := []string{"chat", "task", "analyze"}
result, err := ToModes(modes)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result) != 3 {
t.Errorf("Expected 3 modes, got %d", len(result))
}
if result[0] != "chat" {
t.Errorf("Expected 'chat', got '%s'", result[0])
}
if result[1] != "task" {
t.Errorf("Expected 'task', got '%s'", result[1])
}
if result[2] != "analyze" {
t.Errorf("Expected 'analyze', got '%s'", result[2])
}
})
t.Run("InterfaceSlice", func(t *testing.T) {
modes := []interface{}{"chat", "task", 123}
result, err := ToModes(modes)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result) != 3 {
t.Errorf("Expected 3 modes, got %d", len(result))
}
if result[0] != "chat" {
t.Errorf("Expected 'chat', got '%s'", result[0])
}
if result[2] != "123" {
t.Errorf("Expected '123', got '%s'", result[2])
}
})
t.Run("SingleString", func(t *testing.T) {
mode := "chat"
result, err := ToModes(mode)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result) != 1 {
t.Errorf("Expected 1 mode, got %d", len(result))
}
if result[0] != "chat" {
t.Errorf("Expected 'chat', got '%s'", result[0])
}
})
t.Run("EmptySlice", func(t *testing.T) {
modes := []string{}
result, err := ToModes(modes)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result) != 0 {
t.Errorf("Expected 0 modes, got %d", len(result))
}
})
t.Run("InvalidInput", func(t *testing.T) {
// Test with data that can't be marshaled
invalidData := make(chan int)
_, err := ToModes(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 []string
data := map[string]interface{}{
"invalid": "structure",
}
_, err := ToModes(data)
if err == nil {
t.Error("Expected error for invalid unmarshal")
}
})
t.Run("MixedTypes", func(t *testing.T) {
modes := []interface{}{"chat", 456, "task", true}
result, err := ToModes(modes)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result) != 4 {
t.Errorf("Expected 4 modes, got %d", len(result))
}
// cast.ToString should convert all to strings
if result[0] != "chat" {
t.Errorf("Expected 'chat', got '%s'", result[0])
}
if result[1] != "456" {
t.Errorf("Expected '456', got '%s'", result[1])
}
if result[2] != "task" {
t.Errorf("Expected 'task', got '%s'", result[2])
}
if result[3] != "true" {
t.Errorf("Expected 'true', got '%s'", result[3])
}
})
}
// TestToPromptPresets tests the ToPromptPresets conversion function
func TestToPromptPresets(t *testing.T) {
t.Run("NilInput", func(t *testing.T) {