From d0d110b0eccd93bee6c5cdbf0fbd9151eba83baf Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 1 Jan 2025 12:31:47 +0800 Subject: [PATCH] Refactor assistant management and enhance storage retrieval in Neo API - Removed Weaviate store implementation, streamlining the codebase and focusing on Mongo and Redis backends. - Introduced GetAssistant method in both Mongo and Redis stores to retrieve a single assistant by ID, improving data access capabilities. - Updated LoadStore function to utilize the new storage retrieval logic, enhancing the assistant loading process. - Enhanced the Assistant struct to include a Script field for better management of assistant scripts. - Improved tests to cover the new GetAssistant functionality, ensuring robust error handling and data retrieval across different scenarios. --- neo/assistant/assistant.go | 220 ++++++++++++++++++++++++++++++-- neo/assistant/assistant_test.go | 196 ++++++++++++++++++++++++++++ neo/assistant/cache.go | 105 +++++++++++++++ neo/assistant/cache_test.go | 151 ++++++++++++++++++++++ neo/assistant/types.go | 3 + neo/neo.go | 4 - neo/store/mongo.go | 5 + neo/store/redis.go | 5 + neo/store/types.go | 5 + neo/store/weaviate.go | 59 --------- neo/store/xun.go | 27 ++++ neo/store/xun_test.go | 53 ++++++++ 12 files changed, 756 insertions(+), 77 deletions(-) create mode 100644 neo/assistant/assistant_test.go create mode 100644 neo/assistant/cache.go create mode 100644 neo/assistant/cache_test.go delete mode 100644 neo/store/weaviate.go diff --git a/neo/assistant/assistant.go b/neo/assistant/assistant.go index 0de36580..f3ad5640 100644 --- a/neo/assistant/assistant.go +++ b/neo/assistant/assistant.go @@ -1,26 +1,218 @@ package assistant -import "github.com/yaoapp/yao/neo/store" +import ( + "fmt" + "path/filepath" + "regexp" + "strings" + "time" -// loadedAssistant the loaded assistant -var loadedAssistant = map[string]*Assistant{} + jsoniter "github.com/json-iterator/go" + "github.com/yaoapp/gou/fs" + v8 "github.com/yaoapp/gou/runtime/v8" + "github.com/yaoapp/yao/neo/store" + "github.com/yaoapp/yao/share" + "gopkg.in/yaml.v3" +) -// LoadLocal create a new assistant from local -func LoadLocal(path string) *Assistant { - return nil +// loaded the loaded assistant +var loaded = NewCache(200) // 200 is the default capacity +var storage store.Store = nil + +// SetStorage set the storage +func SetStorage(s store.Store) { + storage = s } -// LoadZip create a new assistant from zip -func LoadZip(zip string) *Assistant { - return nil +// SetCache set the cache +func SetCache(capacity int) { + ClearCache() + loaded = NewCache(capacity) } -// LoadRemote create a new assistant from remote -func LoadRemote(url string) *Assistant { - return nil +// ClearCache clear the cache +func ClearCache() { + if loaded != nil { + loaded.Clear() + loaded = nil + } } // LoadStore create a new assistant from store -func LoadStore(store store.Store) *Assistant { - return nil +func LoadStore(id string) (*Assistant, error) { + assistant, exists := loaded.Get(id) + if exists { + return assistant, nil + } + + if storage == nil { + return nil, fmt.Errorf("storage is not set") + } + + data, err := storage.GetAssistant(id) + if err != nil { + return nil, err + } + + assistant, err = loadMap(data) + if err != nil { + return nil, err + } + + loaded.Put(assistant) + return assistant, nil +} + +// LoadPath load assistant from path +func LoadPath(path string) (*Assistant, error) { + app, err := fs.Get("app") + if err != nil { + return nil, err + } + + pkgfile := filepath.Join(path, "package.yao") + if has, _ := app.Exists(pkgfile); !has { + return nil, fmt.Errorf("package.yao not found in %s", path) + } + + pkg, err := app.ReadFile(pkgfile) + if err != nil { + return nil, err + } + + id := strings.ReplaceAll(strings.TrimPrefix(path, "/assistants/"), "/", ".") + var data map[string]interface{} + err = jsoniter.Unmarshal(pkg, &data) + if err != nil { + return nil, err + } + + // assistant_id + data["assistant_id"] = id + + // prompts + promptsfile := filepath.Join(path, "prompts.yml") + if has, _ := app.Exists(promptsfile); has { + prompts, err := loadPrompts(promptsfile, path) + if err != nil { + return nil, err + } + data["prompts"] = prompts + } + + // load script + scriptfile := filepath.Join(path, "src", "index.ts") + if has, _ := app.Exists(scriptfile); has { + script, err := loadScript(scriptfile, path) + if err != nil { + return nil, err + } + data["script"] = script + } + + // load functions + + // load flow + + return loadMap(data) +} + +func loadMap(data map[string]interface{}) (*Assistant, error) { + + assistant := &Assistant{} + + // assistant_id is required + id, ok := data["assistant_id"].(string) + if !ok { + return nil, fmt.Errorf("assistant_id is required") + } + assistant.ID = id + + // name is required + name, ok := data["name"].(string) + if !ok { + return nil, fmt.Errorf("name is required") + } + assistant.Name = name + + // avatar + if avatar, ok := data["avatar"].(string); ok { + assistant.Avatar = avatar + } + + // connector + if connector, ok := data["connector"].(string); ok { + assistant.Connector = connector + } + + // prompts + if v, ok := data["prompts"].(string); ok { + var prompts []Prompt + err := yaml.Unmarshal([]byte(v), &prompts) + if err != nil { + return nil, err + } + assistant.Prompts = prompts + } + + // script + if data["script"] != nil { + switch v := data["script"].(type) { + case string: + file := fmt.Sprintf("assistants/%s/src/index.ts", assistant.ID) + script, err := loadScriptSource(v, file) + if err != nil { + return nil, err + } + assistant.Script = script + case *v8.Script: + assistant.Script = v + } + } + + return assistant, nil +} + +func loadPrompts(file string, root string) (string, error) { + + app, err := fs.Get("app") + if err != nil { + return "", err + } + + prompts, err := app.ReadFile(file) + if err != nil { + return "", err + } + + re := regexp.MustCompile(`@assets/([^\s]+\.(md|yml|yaml|json|txt))`) + prompts = re.ReplaceAllFunc(prompts, func(s []byte) []byte { + asset := re.FindStringSubmatch(string(s))[1] + assetFile := filepath.Join(root, "assets", asset) + assetContent, err := app.ReadFile(assetFile) + if err != nil { + return []byte("") + } + // Add proper YAML formatting for content + lines := strings.Split(string(assetContent), "\n") + formattedContent := "|\n" + for _, line := range lines { + formattedContent += " " + line + "\n" + } + return []byte(formattedContent) + }) + + return string(prompts), nil +} + +func loadScript(file string, root string) (*v8.Script, error) { + return v8.Load(file, share.ID(root, file)) +} + +func loadScriptSource(source string, file string) (*v8.Script, error) { + script, err := v8.MakeScript([]byte(source), file, 5*time.Second, true) + if err != nil { + return nil, err + } + return script, nil } diff --git a/neo/assistant/assistant_test.go b/neo/assistant/assistant_test.go new file mode 100644 index 00000000..f3319c1c --- /dev/null +++ b/neo/assistant/assistant_test.go @@ -0,0 +1,196 @@ +package assistant + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/neo/store" + "github.com/yaoapp/yao/test" +) + +func prepare(t *testing.T) { + test.Prepare(t, config.Conf) +} + +func TestAssistant_LoadPath(t *testing.T) { + prepare(t) + defer test.Clean() + + assistant, err := LoadPath("/assistants/modi") + if err != nil { + t.Fatal(err) + } + + // Validate basic properties + assert.NotNil(t, assistant) + assert.Equal(t, "modi", assistant.ID) + assert.Equal(t, "Modi", assistant.Name) + assert.Equal(t, "https://api.dicebear.com/7.x/bottts/svg?seed=Modi", assistant.Avatar) + assert.Equal(t, "deepseek", assistant.Connector) + assert.NotNil(t, assistant.Prompts) + assert.NotNil(t, assistant.Script) + + // Test non-existent assistant + _, err = LoadPath("/assistants/non-existent") + assert.Error(t, err) +} + +func TestAssistant_LoadStore(t *testing.T) { + prepare(t) + defer test.Clean() + + // Test with nil storage + _, err := LoadStore("test-id") + assert.Error(t, err) + assert.Contains(t, err.Error(), "storage is not set") + + // Setup mock storage + mockStore := &mockStore{ + data: map[string]map[string]interface{}{ + "test-id": { + "assistant_id": "test-id", + "name": "Test Assistant", + "avatar": "test-avatar", + "connector": "test-connector", + }, + }, + } + SetStorage(mockStore) + defer SetStorage(nil) + + // Test loading from store + assistant, err := LoadStore("test-id") + assert.NoError(t, err) + assert.NotNil(t, assistant) + assert.Equal(t, "test-id", assistant.ID) + assert.Equal(t, "Test Assistant", assistant.Name) + assert.Equal(t, "test-avatar", assistant.Avatar) + assert.Equal(t, "test-connector", assistant.Connector) + + // Test cache functionality + assistant2, err := LoadStore("test-id") + assert.NoError(t, err) + assert.Equal(t, assistant, assistant2) // Should be the same instance from cache + + // Test non-existent assistant + _, err = LoadStore("non-existent") + assert.Error(t, err) +} + +func TestAssistant_Cache(t *testing.T) { + prepare(t) + defer test.Clean() + + // Clear any existing cache first + ClearCache() + + // Test cache operations + SetCache(2) // Set small cache size for testing + assert.Equal(t, 2, loaded.capacity, "Cache capacity should be 2") + + // Create test assistants + assistant1 := &Assistant{ID: "id1", Name: "Assistant 1"} + assistant2 := &Assistant{ID: "id2", Name: "Assistant 2"} + assistant3 := &Assistant{ID: "id3", Name: "Assistant 3"} + + // Test Put and Get + loaded.Put(assistant1) + assert.Equal(t, 1, loaded.Len(), "Cache should have 1 item") + + loaded.Put(assistant2) + assert.Equal(t, 2, loaded.Len(), "Cache should have 2 items") + + // Test cache hit + cached, exists := loaded.Get("id1") + assert.True(t, exists) + assert.Equal(t, assistant1, cached) + + // Test cache eviction (LRU) + // At this point: assistant1 is most recently used (due to Get), then assistant2 + loaded.Put(assistant3) // This should evict assistant2 since it's least recently used + assert.Equal(t, 2, loaded.Len(), "Cache should still have 2 items") + _, exists = loaded.Get("id2") + assert.False(t, exists, "assistant2 should have been evicted (least recently used)") + _, exists = loaded.Get("id1") + assert.True(t, exists, "assistant1 should still be in cache (was accessed recently)") + _, exists = loaded.Get("id3") + assert.True(t, exists, "assistant3 should be in cache (most recently added)") + + // Test clear cache + ClearCache() + assert.Nil(t, loaded) + + // Test setting new cache capacity + SetCache(100) + assert.NotNil(t, loaded) +} + +// mockStore implements store.Store interface for testing +type mockStore struct { + data map[string]map[string]interface{} +} + +func (m *mockStore) GetAssistant(id string) (map[string]interface{}, error) { + if data, ok := m.data[id]; ok { + return data, nil + } + return nil, fmt.Errorf("assistant not found: %s", id) +} + +// Add other required interface methods with empty implementations +func (m *mockStore) GetThread(id string) (map[string]interface{}, error) { return nil, nil } +func (m *mockStore) GetMessage(id string) (map[string]interface{}, error) { return nil, nil } +func (m *mockStore) GetFile(id string) (map[string]interface{}, error) { return nil, nil } +func (m *mockStore) CreateAssistant(data map[string]interface{}) (map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) CreateThread(data map[string]interface{}) (map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) CreateMessage(data map[string]interface{}) (map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) CreateFile(data map[string]interface{}) (map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) UpdateAssistant(id string, data map[string]interface{}) error { return nil } +func (m *mockStore) UpdateThread(id string, data map[string]interface{}) error { return nil } +func (m *mockStore) UpdateMessage(id string, data map[string]interface{}) error { return nil } +func (m *mockStore) UpdateFile(id string, data map[string]interface{}) error { return nil } +func (m *mockStore) DeleteAssistant(id string) error { return nil } +func (m *mockStore) DeleteThread(id string) error { return nil } +func (m *mockStore) DeleteMessage(id string) error { return nil } +func (m *mockStore) DeleteFile(id string) error { return nil } +func (m *mockStore) ListAssistants(query map[string]interface{}) ([]map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) ListThreads(query map[string]interface{}) ([]map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) ListMessages(query map[string]interface{}) ([]map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) ListFiles(query map[string]interface{}) ([]map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) DeleteAllChats(id string) error { return nil } +func (m *mockStore) DeleteChat(id string, chatID string) error { return nil } +func (m *mockStore) GetAssistants(filter store.AssistantFilter) (*store.AssistantResponse, error) { + return nil, nil +} +func (m *mockStore) GetChat(id string, chatID string) (*store.ChatInfo, error) { return nil, nil } +func (m *mockStore) GetChats(id string, filter store.ChatFilter) (*store.ChatGroupResponse, error) { + return nil, nil +} +func (m *mockStore) GetHistory(id string, chatID string) ([]map[string]interface{}, error) { + return nil, nil +} +func (m *mockStore) SaveAssistant(assistant map[string]interface{}) (interface{}, error) { + return nil, nil +} +func (m *mockStore) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { + return nil +} +func (m *mockStore) UpdateChatTitle(sid string, cid string, title string) error { return nil } diff --git a/neo/assistant/cache.go b/neo/assistant/cache.go new file mode 100644 index 00000000..47383af9 --- /dev/null +++ b/neo/assistant/cache.go @@ -0,0 +1,105 @@ +package assistant + +import ( + "container/list" + "sync" +) + +// Cache represents a thread-safe LRU cache for Assistant objects +type Cache struct { + capacity int + mu sync.RWMutex + list *list.List + items map[string]*list.Element +} + +// cacheItem represents an item in the cache +type cacheItem struct { + key string + value *Assistant +} + +// NewCache creates a new LRU cache with the given capacity +func NewCache(capacity int) *Cache { + return &Cache{ + capacity: capacity, + list: list.New(), + items: make(map[string]*list.Element), + } +} + +// Get retrieves an Assistant from the cache by its ID +func (c *Cache) Get(id string) (*Assistant, bool) { + c.mu.Lock() + defer c.mu.Unlock() + + if element, exists := c.items[id]; exists { + c.list.MoveToFront(element) + return element.Value.(*cacheItem).value, true + } + return nil, false +} + +// Put adds or updates an Assistant in the cache +func (c *Cache) Put(assistant *Assistant) { + if assistant == nil || assistant.ID == "" { + return + } + + c.mu.Lock() + defer c.mu.Unlock() + + // If item exists, update it and move to front + if element, exists := c.items[assistant.ID]; exists { + c.list.MoveToFront(element) + element.Value.(*cacheItem).value = assistant + return + } + + // If cache is at capacity, remove oldest item before adding new one + if c.list.Len() >= c.capacity { + c.removeOldest() + } + + // Add new item + element := c.list.PushFront(&cacheItem{ + key: assistant.ID, + value: assistant, + }) + c.items[assistant.ID] = element +} + +// Remove removes an Assistant from the cache +func (c *Cache) Remove(id string) { + c.mu.Lock() + defer c.mu.Unlock() + + if element, exists := c.items[id]; exists { + c.list.Remove(element) + delete(c.items, id) + } +} + +// Len returns the current number of items in the cache +func (c *Cache) Len() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.list.Len() +} + +// Clear removes all items from the cache +func (c *Cache) Clear() { + c.mu.Lock() + defer c.mu.Unlock() + + c.list.Init() + c.items = make(map[string]*list.Element) +} + +// removeOldest removes the least recently used item from the cache +func (c *Cache) removeOldest() { + if element := c.list.Back(); element != nil { + c.list.Remove(element) + delete(c.items, element.Value.(*cacheItem).key) + } +} diff --git a/neo/assistant/cache_test.go b/neo/assistant/cache_test.go new file mode 100644 index 00000000..9ae260ce --- /dev/null +++ b/neo/assistant/cache_test.go @@ -0,0 +1,151 @@ +package assistant + +import ( + "sync" + "testing" +) + +func TestCache_Basic(t *testing.T) { + cache := NewCache(2) + + // Test empty cache + if cache.Len() != 0 { + t.Errorf("Expected empty cache, got length %d", cache.Len()) + } + + // Test adding items + assistant1 := &Assistant{ID: "1", Name: "Test1"} + assistant2 := &Assistant{ID: "2", Name: "Test2"} + + cache.Put(assistant1) + cache.Put(assistant2) + + if cache.Len() != 2 { + t.Errorf("Expected cache length 2, got %d", cache.Len()) + } + + // Test getting items + if a, exists := cache.Get("1"); !exists || a.ID != "1" { + t.Error("Failed to get assistant1") + } + + if a, exists := cache.Get("2"); !exists || a.ID != "2" { + t.Error("Failed to get assistant2") + } +} + +func TestCache_LRU(t *testing.T) { + cache := NewCache(2) + + assistant1 := &Assistant{ID: "1", Name: "Test1"} + assistant2 := &Assistant{ID: "2", Name: "Test2"} + assistant3 := &Assistant{ID: "3", Name: "Test3"} + + // Add first two items + cache.Put(assistant1) + cache.Put(assistant2) + + // Access assistant1 to make it most recently used + cache.Get("1") + + // Add third item, should evict assistant2 + cache.Put(assistant3) + + // Check assistant2 was evicted + if _, exists := cache.Get("2"); exists { + t.Error("Assistant2 should have been evicted") + } + + // Check assistant1 and assistant3 are still present + if _, exists := cache.Get("1"); !exists { + t.Error("Assistant1 should still be in cache") + } + if _, exists := cache.Get("3"); !exists { + t.Error("Assistant3 should be in cache") + } +} + +func TestCache_Remove(t *testing.T) { + cache := NewCache(2) + + assistant1 := &Assistant{ID: "1", Name: "Test1"} + cache.Put(assistant1) + + // Test remove existing item + cache.Remove("1") + if cache.Len() != 0 { + t.Error("Cache should be empty after removing item") + } + + // Test remove non-existing item + cache.Remove("nonexistent") + if cache.Len() != 0 { + t.Error("Cache length should not change when removing non-existent item") + } +} + +func TestCache_Clear(t *testing.T) { + cache := NewCache(2) + + assistant1 := &Assistant{ID: "1", Name: "Test1"} + assistant2 := &Assistant{ID: "2", Name: "Test2"} + + cache.Put(assistant1) + cache.Put(assistant2) + + cache.Clear() + if cache.Len() != 0 { + t.Error("Cache should be empty after clear") + } +} + +func TestCache_Concurrent(t *testing.T) { + cache := NewCache(100) + var wg sync.WaitGroup + workers := 10 + iterations := 100 + + // Concurrent writes + for i := 0; i < workers; i++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() + for j := 0; j < iterations; j++ { + assistant := &Assistant{ + ID: string(rune('A' + workerID)), + Name: "Test", + } + cache.Put(assistant) + } + }(i) + } + + // Concurrent reads + for i := 0; i < workers; i++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() + for j := 0; j < iterations; j++ { + cache.Get(string(rune('A' + workerID))) + } + }(i) + } + + wg.Wait() +} + +func TestCache_NilInput(t *testing.T) { + cache := NewCache(2) + + // Test putting nil assistant + cache.Put(nil) + if cache.Len() != 0 { + t.Error("Cache should not store nil assistant") + } + + // Test putting assistant with empty ID + cache.Put(&Assistant{ID: "", Name: "Test"}) + if cache.Len() != 0 { + t.Error("Cache should not store assistant with empty ID") + } +} diff --git a/neo/assistant/types.go b/neo/assistant/types.go index a122ed06..30b8c801 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -4,6 +4,8 @@ import ( "context" "io" "mime/multipart" + + v8 "github.com/yaoapp/gou/runtime/v8" ) // API the assistant API interface @@ -40,6 +42,7 @@ type Assistant struct { Option map[string]interface{} `json:"option,omitempty"` // AI Option Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows + Script *v8.Script `json:"-" yaml:"-"` // Assistant Script API API `json:"-" yaml:"-"` // Assistant API } diff --git a/neo/neo.go b/neo/neo.go index 97d16767..f072dc1c 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -535,10 +535,6 @@ func (neo *DSL) createStore() error { } else if conn.Is(connector.MONGO) { neo.Store = store.NewMongo() return nil - - } else if conn.Is(connector.WEAVIATE) { - neo.Store = store.NewWeaviate() - return nil } return fmt.Errorf("%s store connector %s not support", neo.ID, neo.StoreSetting.Connector) diff --git a/neo/store/mongo.go b/neo/store/mongo.go index 88336b1f..9c7ec275 100644 --- a/neo/store/mongo.go +++ b/neo/store/mongo.go @@ -57,3 +57,8 @@ func (m *Mongo) DeleteAssistant(assistantID string) error { func (m *Mongo) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { return &AssistantResponse{}, nil } + +// GetAssistant retrieves a single assistant by ID +func (m *Mongo) GetAssistant(assistantID string) (map[string]interface{}, error) { + return map[string]interface{}{}, nil +} diff --git a/neo/store/redis.go b/neo/store/redis.go index ad249543..94ab30db 100644 --- a/neo/store/redis.go +++ b/neo/store/redis.go @@ -57,3 +57,8 @@ func (r *Redis) DeleteAssistant(assistantID string) error { func (r *Redis) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { return &AssistantResponse{}, nil } + +// GetAssistant retrieves a single assistant by ID +func (r *Redis) GetAssistant(assistantID string) (map[string]interface{}, error) { + return map[string]interface{}{}, nil +} diff --git a/neo/store/types.go b/neo/store/types.go index ccfa2039..628a7627 100644 --- a/neo/store/types.go +++ b/neo/store/types.go @@ -130,4 +130,9 @@ type Store interface { // filter: Filter conditions // Returns: Paginated assistant list and potential error GetAssistants(filter AssistantFilter) (*AssistantResponse, error) + + // GetAssistant retrieves a single assistant by ID + // assistantID: Assistant ID + // Returns: Assistant information and potential error + GetAssistant(assistantID string) (map[string]interface{}, error) } diff --git a/neo/store/weaviate.go b/neo/store/weaviate.go deleted file mode 100644 index 4b2dda07..00000000 --- a/neo/store/weaviate.go +++ /dev/null @@ -1,59 +0,0 @@ -package store - -// Weaviate represents a Weaviate-based conversation storage -type Weaviate struct{} - -// NewWeaviate create a new weaviate store -func NewWeaviate() Store { - return &Weaviate{} -} - -// GetChats retrieves a list of chats -func (w *Weaviate) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { - return &ChatGroupResponse{}, nil -} - -// GetChat retrieves a single chat's information -func (w *Weaviate) GetChat(sid string, cid string) (*ChatInfo, error) { - return &ChatInfo{}, nil -} - -// GetHistory retrieves chat history -func (w *Weaviate) GetHistory(sid string, cid string) ([]map[string]interface{}, error) { - return []map[string]interface{}{}, nil -} - -// SaveHistory saves chat history -func (w *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error { - return nil -} - -// DeleteChat deletes a single chat -func (w *Weaviate) DeleteChat(sid string, cid string) error { - return nil -} - -// DeleteAllChats deletes all chats -func (w *Weaviate) DeleteAllChats(sid string) error { - return nil -} - -// UpdateChatTitle updates chat title -func (w *Weaviate) UpdateChatTitle(sid string, cid string, title string) error { - return nil -} - -// SaveAssistant saves assistant information -func (w *Weaviate) SaveAssistant(assistant map[string]interface{}) (interface{}, error) { - return assistant["assistant_id"], nil -} - -// DeleteAssistant deletes an assistant -func (w *Weaviate) DeleteAssistant(assistantID string) error { - return nil -} - -// GetAssistants retrieves a list of assistants -func (w *Weaviate) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) { - return &AssistantResponse{}, nil -} diff --git a/neo/store/xun.go b/neo/store/xun.go index 2407bdc2..41ee1ade 100644 --- a/neo/store/xun.go +++ b/neo/store/xun.go @@ -44,6 +44,7 @@ type Xun struct { // SaveAssistant creates or updates an assistant // DeleteAssistant deletes an assistant by assistant_id // GetAssistants retrieves a paginated list of assistants with filtering +// GetAssistant retrieves a single assistant by assistant_id // NewXun create a new xun store func NewXun(setting Setting) (Store, error) { @@ -923,3 +924,29 @@ func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, erro Total: total, }, nil } + +// GetAssistant retrieves a single assistant by ID +func (conv *Xun) GetAssistant(assistantID string) (map[string]interface{}, error) { + row, err := conv.query.New(). + Table(conv.getAssistantTable()). + Where("assistant_id", assistantID). + First() + if err != nil { + return nil, err + } + + if row == nil { + return nil, fmt.Errorf("assistant %s not found", assistantID) + } + + data := row.ToMap() + if data == nil || len(data) == 0 { + return nil, fmt.Errorf("assistant %s not found", assistantID) + } + + // Parse JSON fields + jsonFields := []string{"tags", "options", "prompts", "flows", "files", "functions", "permissions"} + conv.parseJSONFields(data, jsonFields) + + return data, nil +} diff --git a/neo/store/xun_test.go b/neo/store/xun_test.go index a8613c3a..a41ae3fa 100644 --- a/neo/store/xun_test.go +++ b/neo/store/xun_test.go @@ -483,6 +483,20 @@ func TestXunAssistantCRUD(t *testing.T) { assistantID := v.(string) assert.NotEmpty(t, assistantID) + // Test GetAssistant for the first assistant + assistantData, err := store.GetAssistant(assistantID) + assert.Nil(t, err) + assert.NotNil(t, assistantData) + assert.Equal(t, "Test Assistant", assistantData["name"]) + assert.Equal(t, "assistant", assistantData["type"]) + assert.Equal(t, "https://example.com/avatar.png", assistantData["avatar"]) + assert.Equal(t, "openai", assistantData["connector"]) + assert.Equal(t, "Test Description", assistantData["description"]) + assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, assistantData["tags"]) + assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, assistantData["options"]) + assert.Equal(t, int64(1), assistantData["mentionable"]) + assert.Equal(t, int64(1), assistantData["automated"]) + // Test case 2: JSON fields as native types assistant2 := map[string]interface{}{ "name": "Test Assistant 2", @@ -507,6 +521,24 @@ func TestXunAssistantCRUD(t *testing.T) { assistant2ID := v.(string) assert.NotEmpty(t, assistant2ID) + // Test GetAssistant for the second assistant + assistant2Data, err := store.GetAssistant(assistant2ID) + assert.Nil(t, err) + assert.NotNil(t, assistant2Data) + assert.Equal(t, "Test Assistant 2", assistant2Data["name"]) + assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, assistant2Data["tags"]) + assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, assistant2Data["options"]) + assert.Equal(t, []interface{}{"prompt1", "prompt2"}, assistant2Data["prompts"]) + assert.Equal(t, []interface{}{"flow1", "flow2"}, assistant2Data["flows"]) + assert.Equal(t, []interface{}{"file1", "file2"}, assistant2Data["files"]) + assert.Equal(t, []interface{}{ + map[string]interface{}{"name": "func1"}, + map[string]interface{}{"name": "func2"}, + }, assistant2Data["functions"]) + assert.Equal(t, map[string]interface{}{"read": true, "write": true}, assistant2Data["permissions"]) + assert.Equal(t, int64(1), assistant2Data["mentionable"]) + assert.Equal(t, int64(1), assistant2Data["automated"]) + // Test case 3: Test with nil JSON fields assistant3 := map[string]interface{}{ "name": "Test Assistant 3", @@ -530,6 +562,27 @@ func TestXunAssistantCRUD(t *testing.T) { assistant3ID := v.(string) assert.NotEmpty(t, assistant3ID) + // Test GetAssistant for the third assistant + assistant3Data, err := store.GetAssistant(assistant3ID) + assert.Nil(t, err) + assert.NotNil(t, assistant3Data) + assert.Equal(t, "Test Assistant 3", assistant3Data["name"]) + assert.Nil(t, assistant3Data["tags"]) + assert.Nil(t, assistant3Data["options"]) + assert.Nil(t, assistant3Data["prompts"]) + assert.Nil(t, assistant3Data["flows"]) + assert.Nil(t, assistant3Data["files"]) + assert.Nil(t, assistant3Data["functions"]) + assert.Nil(t, assistant3Data["permissions"]) + assert.Equal(t, int64(1), assistant3Data["mentionable"]) + assert.Equal(t, int64(1), assistant3Data["automated"]) + + // Test GetAssistant with non-existent ID + nonExistentData, err := store.GetAssistant("non-existent-id") + assert.Error(t, err) + assert.Nil(t, nonExistentData) + assert.Contains(t, err.Error(), "not found") + // Test GetAssistants to verify JSON fields are properly stored resp, err := store.GetAssistants(AssistantFilter{}) assert.Nil(t, err)