From bbba6b9fdb4380cff96bb721d8402543b254c5cd Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 3 Dec 2025 10:50:24 +0800 Subject: [PATCH] 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. --- agent/assistant/load.go | 23 ++++++ agent/store/types/convert.go | 36 +++++++++ agent/store/types/convert_test.go | 119 ++++++++++++++++++++++++++++++ 3 files changed, 178 insertions(+) diff --git a/agent/assistant/load.go b/agent/assistant/load.go index c6c8d524..3a059eb0 100644 --- a/agent/assistant/load.go +++ b/agent/assistant/load.go @@ -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) diff --git a/agent/store/types/convert.go b/agent/store/types/convert.go index 03a3f62d..0d3b3974 100644 --- a/agent/store/types/convert.go +++ b/agent/store/types/convert.go @@ -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 { diff --git a/agent/store/types/convert_test.go b/agent/store/types/convert_test.go index 800bef93..b0b8b347 100644 --- a/agent/store/types/convert_test.go +++ b/agent/store/types/convert_test.go @@ -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) {