Refactor Assistant cache tests and enhance cache functionality
- Updated cache tests to use the testify assertion library for improved readability and maintainability. - Added new tests for cache operations including basic functionality, LRU eviction, removal, clearing, and concurrent access. - Enhanced the cache implementation to unregister scripts upon removal and clearing, ensuring proper resource management. - Introduced an `All` method to retrieve all assistants in the cache, improving accessibility of cached items. - Streamlined the script registration and unregistration process within the Assistant struct, enhancing script management.
This commit is contained in:
Max
parent
d0981e53af
commit
f0e445862f
6 changed files with 673 additions and 154 deletions
|
|
@ -75,6 +75,13 @@ func (c *Cache) Remove(id string) {
|
|||
defer c.mu.Unlock()
|
||||
|
||||
if element, exists := c.items[id]; exists {
|
||||
item := element.Value.(*cacheItem)
|
||||
|
||||
// Unregister scripts before removing from cache
|
||||
if item.value != nil && len(item.value.Scripts) > 0 {
|
||||
item.value.UnregisterScripts()
|
||||
}
|
||||
|
||||
c.list.Remove(element)
|
||||
delete(c.items, id)
|
||||
}
|
||||
|
|
@ -87,11 +94,32 @@ func (c *Cache) Len() int {
|
|||
return c.list.Len()
|
||||
}
|
||||
|
||||
// All returns all assistants in the cache
|
||||
func (c *Cache) All() []*Assistant {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
|
||||
assistants := make([]*Assistant, 0, c.list.Len())
|
||||
for element := c.list.Front(); element != nil; element = element.Next() {
|
||||
item := element.Value.(*cacheItem)
|
||||
assistants = append(assistants, item.value)
|
||||
}
|
||||
return assistants
|
||||
}
|
||||
|
||||
// Clear removes all items from the cache
|
||||
func (c *Cache) Clear() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
// Unregister all scripts before clearing cache
|
||||
for element := c.list.Front(); element != nil; element = element.Next() {
|
||||
item := element.Value.(*cacheItem)
|
||||
if item.value != nil && len(item.value.Scripts) > 0 {
|
||||
item.value.UnregisterScripts()
|
||||
}
|
||||
}
|
||||
|
||||
c.list.Init()
|
||||
c.items = make(map[string]*list.Element)
|
||||
}
|
||||
|
|
@ -99,7 +127,14 @@ func (c *Cache) Clear() {
|
|||
// removeOldest removes the least recently used item from the cache
|
||||
func (c *Cache) removeOldest() {
|
||||
if element := c.list.Back(); element != nil {
|
||||
item := element.Value.(*cacheItem)
|
||||
|
||||
// Unregister scripts before removing from cache
|
||||
if item.value != nil && len(item.value.Scripts) > 0 {
|
||||
item.value.UnregisterScripts()
|
||||
}
|
||||
|
||||
c.list.Remove(element)
|
||||
delete(c.items, element.Value.(*cacheItem).key)
|
||||
delete(c.items, item.key)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,146 +1,253 @@
|
|||
package assistant
|
||||
package assistant_test
|
||||
|
||||
// func TestCache_Basic(t *testing.T) {
|
||||
// cache := NewCache(2)
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
// // Test empty cache
|
||||
// if cache.Len() != 0 {
|
||||
// t.Errorf("Expected empty cache, got length %d", cache.Len())
|
||||
// }
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/gou/process"
|
||||
"github.com/yaoapp/yao/agent/assistant"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
)
|
||||
|
||||
// // Test adding items
|
||||
// assistant1 := &Assistant{ID: "1", Name: "Test1"}
|
||||
// assistant2 := &Assistant{ID: "2", Name: "Test2"}
|
||||
func TestCacheBasic(t *testing.T) {
|
||||
cache := assistant.NewCache(2)
|
||||
|
||||
// cache.Put(assistant1)
|
||||
// cache.Put(assistant2)
|
||||
// Test empty cache
|
||||
assert.Equal(t, 0, cache.Len(), "Expected empty cache")
|
||||
|
||||
// if cache.Len() != 2 {
|
||||
// t.Errorf("Expected cache length 2, got %d", cache.Len())
|
||||
// }
|
||||
// Create test assistants
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// // Test getting items
|
||||
// if a, exists := cache.Get("1"); !exists || a.ID != "1" {
|
||||
// t.Error("Failed to get assistant1")
|
||||
// }
|
||||
ast1, err := assistant.Get("tests.mcpload")
|
||||
assert.NoError(t, err)
|
||||
|
||||
// if a, exists := cache.Get("2"); !exists || a.ID != "2" {
|
||||
// t.Error("Failed to get assistant2")
|
||||
// }
|
||||
// }
|
||||
ast2, err := assistant.Get("tests.create")
|
||||
assert.NoError(t, err)
|
||||
|
||||
// func TestCache_LRU(t *testing.T) {
|
||||
// cache := NewCache(2)
|
||||
// Test adding items
|
||||
cache.Put(ast1)
|
||||
cache.Put(ast2)
|
||||
|
||||
// assistant1 := &Assistant{ID: "1", Name: "Test1"}
|
||||
// assistant2 := &Assistant{ID: "2", Name: "Test2"}
|
||||
// assistant3 := &Assistant{ID: "3", Name: "Test3"}
|
||||
assert.Equal(t, 2, cache.Len(), "Expected cache length 2")
|
||||
|
||||
// // Add first two items
|
||||
// cache.Put(assistant1)
|
||||
// cache.Put(assistant2)
|
||||
// Test getting items
|
||||
cached1, exists := cache.Get("tests.mcpload")
|
||||
assert.True(t, exists, "Should find tests.mcpload")
|
||||
assert.Equal(t, "tests.mcpload", cached1.ID)
|
||||
|
||||
// // Access assistant1 to make it most recently used
|
||||
// cache.Get("1")
|
||||
cached2, exists := cache.Get("tests.create")
|
||||
assert.True(t, exists, "Should find tests.create")
|
||||
assert.Equal(t, "tests.create", cached2.ID)
|
||||
}
|
||||
|
||||
// // Add third item, should evict assistant2
|
||||
// cache.Put(assistant3)
|
||||
func TestCacheLRU(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// // Check assistant2 was evicted
|
||||
// if _, exists := cache.Get("2"); exists {
|
||||
// t.Error("Assistant2 should have been evicted")
|
||||
// }
|
||||
cache := assistant.NewCache(2)
|
||||
|
||||
// // 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")
|
||||
// }
|
||||
// }
|
||||
ast1, _ := assistant.Get("tests.mcpload")
|
||||
ast2, _ := assistant.Get("tests.create")
|
||||
ast3, _ := assistant.Get("tests.next")
|
||||
|
||||
// func TestCache_Remove(t *testing.T) {
|
||||
// cache := NewCache(2)
|
||||
// Add first two items
|
||||
cache.Put(ast1)
|
||||
cache.Put(ast2)
|
||||
|
||||
// assistant1 := &Assistant{ID: "1", Name: "Test1"}
|
||||
// cache.Put(assistant1)
|
||||
// Access ast1 to make it most recently used
|
||||
cache.Get("tests.mcpload")
|
||||
|
||||
// // Test remove existing item
|
||||
// cache.Remove("1")
|
||||
// if cache.Len() != 0 {
|
||||
// t.Error("Cache should be empty after removing item")
|
||||
// }
|
||||
// Add third item, should evict ast2
|
||||
cache.Put(ast3)
|
||||
|
||||
// // Test remove non-existing item
|
||||
// cache.Remove("nonexistent")
|
||||
// if cache.Len() != 0 {
|
||||
// t.Error("Cache length should not change when removing non-existent item")
|
||||
// }
|
||||
// }
|
||||
// Check ast2 was evicted
|
||||
_, exists := cache.Get("tests.create")
|
||||
assert.False(t, exists, "tests.create should have been evicted")
|
||||
|
||||
// func TestCache_Clear(t *testing.T) {
|
||||
// cache := NewCache(2)
|
||||
// Check ast1 and ast3 are still present
|
||||
_, exists = cache.Get("tests.mcpload")
|
||||
assert.True(t, exists, "tests.mcpload should still be in cache")
|
||||
|
||||
// assistant1 := &Assistant{ID: "1", Name: "Test1"}
|
||||
// assistant2 := &Assistant{ID: "2", Name: "Test2"}
|
||||
_, exists = cache.Get("tests.next")
|
||||
assert.True(t, exists, "tests.next should be in cache")
|
||||
}
|
||||
|
||||
// cache.Put(assistant1)
|
||||
// cache.Put(assistant2)
|
||||
func TestCacheRemove(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// cache.Clear()
|
||||
// if cache.Len() != 0 {
|
||||
// t.Error("Cache should be empty after clear")
|
||||
// }
|
||||
// }
|
||||
cache := assistant.NewCache(2)
|
||||
|
||||
// func TestCache_Concurrent(t *testing.T) {
|
||||
// cache := NewCache(100)
|
||||
// var wg sync.WaitGroup
|
||||
// workers := 10
|
||||
// iterations := 100
|
||||
ast1, _ := assistant.Get("tests.mcpload")
|
||||
cache.Put(ast1)
|
||||
|
||||
// // 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)
|
||||
// }
|
||||
// Verify scripts are registered
|
||||
_, exists := process.Handlers["agents.tests.mcpload.tools"]
|
||||
assert.True(t, exists, "Handler should be registered before removal")
|
||||
|
||||
// // 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)
|
||||
// }
|
||||
// Test remove existing item
|
||||
cache.Remove("tests.mcpload")
|
||||
assert.Equal(t, 0, cache.Len(), "Cache should be empty after removing item")
|
||||
|
||||
// wg.Wait()
|
||||
// }
|
||||
// Verify scripts are unregistered
|
||||
_, exists = process.Handlers["agents.tests.mcpload.tools"]
|
||||
assert.False(t, exists, "Handler should be unregistered after removal")
|
||||
|
||||
// func TestCache_NilInput(t *testing.T) {
|
||||
// cache := NewCache(2)
|
||||
// Test remove non-existing item (should not panic)
|
||||
cache.Remove("nonexistent")
|
||||
assert.Equal(t, 0, cache.Len(), "Cache length should not change")
|
||||
}
|
||||
|
||||
// // Test putting nil assistant
|
||||
// cache.Put(nil)
|
||||
// if cache.Len() != 0 {
|
||||
// t.Error("Cache should not store nil assistant")
|
||||
// }
|
||||
func TestCacheClear(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// // 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")
|
||||
// }
|
||||
// }
|
||||
cache := assistant.NewCache(3)
|
||||
|
||||
ast1, _ := assistant.Get("tests.mcpload")
|
||||
ast2, _ := assistant.Get("tests.create")
|
||||
ast3, _ := assistant.Get("tests.next")
|
||||
|
||||
cache.Put(ast1)
|
||||
cache.Put(ast2)
|
||||
cache.Put(ast3)
|
||||
|
||||
assert.Equal(t, 3, cache.Len(), "Cache should have 3 items")
|
||||
|
||||
// Verify scripts are registered
|
||||
_, exists := process.Handlers["agents.tests.mcpload.tools"]
|
||||
assert.True(t, exists, "Handler should be registered before clear")
|
||||
|
||||
// Clear cache
|
||||
cache.Clear()
|
||||
assert.Equal(t, 0, cache.Len(), "Cache should be empty after clear")
|
||||
|
||||
// Verify all scripts are unregistered
|
||||
_, exists = process.Handlers["agents.tests.mcpload.tools"]
|
||||
assert.False(t, exists, "Handler should be unregistered after clear")
|
||||
}
|
||||
|
||||
func TestCacheLRUEviction(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
cache := assistant.NewCache(2)
|
||||
|
||||
ast1, _ := assistant.Get("tests.mcpload")
|
||||
ast2, _ := assistant.Get("tests.create")
|
||||
ast3, _ := assistant.Get("tests.next")
|
||||
|
||||
cache.Put(ast1)
|
||||
cache.Put(ast2)
|
||||
|
||||
// Verify both are registered
|
||||
_, exists1 := process.Handlers["agents.tests.mcpload.tools"]
|
||||
assert.True(t, exists1, "Handler 1 should be registered")
|
||||
|
||||
// Add third item to trigger LRU eviction of oldest (ast1)
|
||||
cache.Put(ast3)
|
||||
|
||||
// Verify ast1's handler was unregistered due to eviction
|
||||
_, exists := process.Handlers["agents.tests.mcpload.tools"]
|
||||
assert.False(t, exists, "Handler should be unregistered after LRU eviction")
|
||||
|
||||
// Verify ast2 and ast3 are still in cache
|
||||
_, exists = cache.Get("tests.create")
|
||||
assert.True(t, exists, "tests.create should still be in cache")
|
||||
|
||||
_, exists = cache.Get("tests.next")
|
||||
assert.True(t, exists, "tests.next should be in cache")
|
||||
}
|
||||
|
||||
func TestCacheConcurrent(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
cache := assistant.NewCache(10)
|
||||
var wg sync.WaitGroup
|
||||
workers := 5
|
||||
iterations := 20
|
||||
|
||||
// Load some assistants for concurrent testing
|
||||
assistants := []string{
|
||||
"tests.mcpload",
|
||||
"tests.create",
|
||||
"tests.next",
|
||||
}
|
||||
|
||||
// Concurrent writes
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func(workerID int) {
|
||||
defer wg.Done()
|
||||
for j := 0; j < iterations; j++ {
|
||||
astID := assistants[j%len(assistants)]
|
||||
ast, _ := assistant.Get(astID)
|
||||
if ast != nil {
|
||||
cache.Put(ast)
|
||||
}
|
||||
}
|
||||
}(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++ {
|
||||
astID := assistants[j%len(assistants)]
|
||||
cache.Get(astID)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Verify cache is in valid state
|
||||
assert.True(t, cache.Len() >= 0, "Cache should have valid length")
|
||||
assert.True(t, cache.Len() <= 10, "Cache should not exceed capacity")
|
||||
}
|
||||
|
||||
func TestCacheNilInput(t *testing.T) {
|
||||
cache := assistant.NewCache(2)
|
||||
|
||||
// Test putting nil assistant
|
||||
cache.Put(nil)
|
||||
assert.Equal(t, 0, cache.Len(), "Cache should not store nil assistant")
|
||||
|
||||
// Test putting assistant with empty ID
|
||||
emptyAST := &assistant.Assistant{}
|
||||
cache.Put(emptyAST)
|
||||
assert.Equal(t, 0, cache.Len(), "Cache should not store assistant with empty ID")
|
||||
}
|
||||
|
||||
func TestCacheAll(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
cache := assistant.NewCache(5)
|
||||
|
||||
ast1, _ := assistant.Get("tests.mcpload")
|
||||
ast2, _ := assistant.Get("tests.create")
|
||||
ast3, _ := assistant.Get("tests.next")
|
||||
|
||||
cache.Put(ast1)
|
||||
cache.Put(ast2)
|
||||
cache.Put(ast3)
|
||||
|
||||
all := cache.All()
|
||||
assert.Equal(t, 3, len(all), "All() should return 3 assistants")
|
||||
|
||||
// Verify all expected assistants are present
|
||||
ids := make(map[string]bool)
|
||||
for _, ast := range all {
|
||||
ids[ast.ID] = true
|
||||
}
|
||||
|
||||
assert.True(t, ids["tests.mcpload"], "Should contain tests.mcpload")
|
||||
assert.True(t, ids["tests.create"], "Should contain tests.create")
|
||||
assert.True(t, ids["tests.next"], "Should contain tests.next")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -730,44 +730,12 @@ func (ast *Assistant) initialize() error {
|
|||
}
|
||||
ast.openai = api
|
||||
|
||||
// Check if the assistant has an init hook
|
||||
if ast.HookScript != nil {
|
||||
scriptCtx, err := ast.HookScript.NewContext("", nil)
|
||||
if err != nil {
|
||||
return err
|
||||
// Register scripts as process handlers
|
||||
if len(ast.Scripts) > 0 {
|
||||
if err := ast.RegisterScripts(); err != nil {
|
||||
return fmt.Errorf("failed to register scripts: %w", err)
|
||||
}
|
||||
defer scriptCtx.Close()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func loadTools(file string) (*store.ToolCalls, int64, error) {
|
||||
|
||||
app, err := fs.Get("app")
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
content, err := app.ReadFile(file)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
ts, err := app.ModTime(file)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
if len(content) == 0 {
|
||||
return &store.ToolCalls{Tools: []store.Tool{}, Prompts: []store.Prompt{}}, ts.UnixNano(), nil
|
||||
}
|
||||
|
||||
var tools store.ToolCalls
|
||||
err = application.Parse(file, content, &tools)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
return &tools, ts.UnixNano(), nil
|
||||
}
|
||||
|
|
|
|||
120
agent/assistant/load_process_test.go
Normal file
120
agent/assistant/load_process_test.go
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
package assistant_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/gou/process"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
)
|
||||
|
||||
func TestLoadProcessIntegration(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// After testutils.Prepare, all assistants should be loaded and scripts registered
|
||||
// Test calling mcpload assistant's tools.Hello function
|
||||
|
||||
t.Run("CallHelloAfterLoad", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.Hello", map[string]interface{}{
|
||||
"name": "TestUser",
|
||||
})
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NoError(t, err)
|
||||
|
||||
result := proc.Value()
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultStr, ok := result.(string)
|
||||
assert.True(t, ok, "Result should be a string")
|
||||
assert.Contains(t, resultStr, "Hello, TestUser")
|
||||
assert.Contains(t, resultStr, "mcpload assistant")
|
||||
})
|
||||
|
||||
t.Run("CallPingAfterLoad", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.Ping", map[string]interface{}{
|
||||
"message": "integration test",
|
||||
})
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NoError(t, err)
|
||||
|
||||
result := proc.Value()
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultMap, ok := result.(map[string]interface{})
|
||||
assert.True(t, ok, "Result should be a map")
|
||||
assert.Equal(t, "integration test", resultMap["message"])
|
||||
assert.Contains(t, resultMap["echo"], "Pong")
|
||||
assert.NotEmpty(t, resultMap["timestamp"])
|
||||
})
|
||||
|
||||
t.Run("CallCalculateAfterLoad", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.Calculate", map[string]interface{}{
|
||||
"operation": "add",
|
||||
"a": float64(100),
|
||||
"b": float64(50),
|
||||
})
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NoError(t, err)
|
||||
|
||||
result := proc.Value()
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultMap, ok := result.(map[string]interface{})
|
||||
assert.True(t, ok, "Result should be a map")
|
||||
assert.Equal(t, float64(150), resultMap["result"])
|
||||
assert.Equal(t, "add", resultMap["operation"])
|
||||
assert.Equal(t, float64(100), resultMap["a"])
|
||||
assert.Equal(t, float64(50), resultMap["b"])
|
||||
})
|
||||
|
||||
t.Run("CallNonExistentScript", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.nonexistent.Method")
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NotNil(t, err, "Should return error for non-existent script")
|
||||
assert.Contains(t, err.Error(), "Exception|404")
|
||||
})
|
||||
|
||||
t.Run("CallNonExistentMethod", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.NonExistentMethod")
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NotNil(t, err, "Should return error for non-existent method")
|
||||
assert.Contains(t, err.Error(), "Exception|500")
|
||||
})
|
||||
}
|
||||
|
||||
func TestLoadProcessMultipleAssistants(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Test that multiple assistants can have their scripts registered
|
||||
// and process calls work correctly for different assistants
|
||||
|
||||
t.Run("MCPLoadAssistant", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.Hello", map[string]interface{}{
|
||||
"name": "User1",
|
||||
})
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NoError(t, err)
|
||||
|
||||
result := proc.Value()
|
||||
resultStr, ok := result.(string)
|
||||
assert.True(t, ok)
|
||||
assert.Contains(t, resultStr, "mcpload assistant")
|
||||
})
|
||||
|
||||
// If there are other test assistants with scripts, they can be tested here
|
||||
// For now, we verify that the handler is properly isolated per assistant
|
||||
t.Run("VerifyIsolation", func(t *testing.T) {
|
||||
// Verify that the mcpload handler is correctly registered
|
||||
handler, exists := process.Handlers["agents.tests.mcpload.tools"]
|
||||
assert.True(t, exists, "Handler should be registered")
|
||||
assert.NotNil(t, handler)
|
||||
})
|
||||
}
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
package assistant
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
|
@ -8,13 +9,34 @@ import (
|
|||
"time"
|
||||
|
||||
"github.com/yaoapp/gou/application"
|
||||
"github.com/yaoapp/gou/process"
|
||||
v8 "github.com/yaoapp/gou/runtime/v8"
|
||||
"github.com/yaoapp/kun/exception"
|
||||
"github.com/yaoapp/yao/agent/assistant/hook"
|
||||
)
|
||||
|
||||
// scriptsMutex protects concurrent v8.Load calls and Scripts map access
|
||||
var scriptsMutex sync.Mutex
|
||||
|
||||
// Execute execute the script
|
||||
func (s *Script) Execute(ctx context.Context, method string, args ...interface{}) (interface{}, error) {
|
||||
if s == nil || s.Script == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
scriptCtx, err := s.NewContext("", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer scriptCtx.Close()
|
||||
|
||||
// Call the method with provided arguments as-is
|
||||
result, err := scriptCtx.CallWith(ctx, method, args...)
|
||||
|
||||
// Return error as-is (including "not defined" errors)
|
||||
return result, err
|
||||
}
|
||||
|
||||
// LoadScripts loads all scripts from a src directory path
|
||||
// It scans for .ts and .js files (excluding index.ts which is the hook script)
|
||||
// Returns the HookScript and a map of other scripts
|
||||
|
|
@ -289,3 +311,60 @@ func loadScriptsField(scriptsData interface{}) (map[string]*Script, error) {
|
|||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// RegisterScripts registers all scripts as process handlers
|
||||
// Handler naming: agents.<assistantID>.<scriptID>
|
||||
func (ast *Assistant) RegisterScripts() error {
|
||||
if len(ast.Scripts) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
assistantID := ast.ID
|
||||
handlers := make(map[string]process.Handler)
|
||||
|
||||
for scriptID, script := range ast.Scripts {
|
||||
// Create handler for this script
|
||||
handlers[scriptID] = makeScriptHandler(script)
|
||||
}
|
||||
|
||||
// Register the handler group dynamically
|
||||
groupName := fmt.Sprintf("agents.%s", assistantID)
|
||||
process.RegisterDynamicGroup(groupName, handlers)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// UnregisterScripts unregisters all scripts from process handlers
|
||||
func (ast *Assistant) UnregisterScripts() error {
|
||||
if len(ast.Scripts) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
assistantID := ast.ID
|
||||
|
||||
for scriptID := range ast.Scripts {
|
||||
handlerID := fmt.Sprintf("agents.%s.%s", strings.ToLower(assistantID), strings.ToLower(scriptID))
|
||||
delete(process.Handlers, handlerID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// makeScriptHandler creates a process handler for a script
|
||||
func makeScriptHandler(script *Script) process.Handler {
|
||||
return func(p *process.Process) interface{} {
|
||||
// Extract method name from process
|
||||
method := p.Method
|
||||
|
||||
// Get arguments from process
|
||||
args := p.Args
|
||||
|
||||
// Execute the script
|
||||
result, err := script.Execute(p.Context, method, args...)
|
||||
if err != nil {
|
||||
exception.New(err.Error(), 500).Throw()
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
}
|
||||
|
|
|
|||
210
agent/assistant/scripts_process_test.go
Normal file
210
agent/assistant/scripts_process_test.go
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
package assistant_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/gou/process"
|
||||
"github.com/yaoapp/yao/agent/assistant"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
)
|
||||
|
||||
func TestScriptsProcessFlow(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Get the mcpload assistant
|
||||
assistantID := "tests.mcpload"
|
||||
ast, err := assistant.Get(assistantID)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, ast, "Assistant should be loaded")
|
||||
|
||||
// Check that scripts were loaded
|
||||
assert.NotNil(t, ast.Scripts)
|
||||
assert.Greater(t, len(ast.Scripts), 0, "Should have loaded at least one script")
|
||||
|
||||
// Verify tools.ts was loaded
|
||||
toolsScript, hasTools := ast.Scripts["tools"]
|
||||
assert.True(t, hasTools, "Should have loaded tools script")
|
||||
assert.NotNil(t, toolsScript)
|
||||
|
||||
// Register scripts as process handlers
|
||||
err = ast.RegisterScripts()
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Test 1: Call Hello function
|
||||
t.Run("CallHelloFunction", func(t *testing.T) {
|
||||
handlerID := "agents.tests.mcpload.tools"
|
||||
handler, exists := process.Handlers[handlerID]
|
||||
assert.True(t, exists, "Handler should be registered")
|
||||
|
||||
p := &process.Process{
|
||||
ID: handlerID + ".Hello",
|
||||
Method: "Hello",
|
||||
Args: []interface{}{map[string]interface{}{"name": "Yao"}},
|
||||
Context: context.Background(),
|
||||
}
|
||||
|
||||
result := handler(p)
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultStr, ok := result.(string)
|
||||
assert.True(t, ok, "Result should be a string")
|
||||
assert.Contains(t, resultStr, "Hello, Yao")
|
||||
})
|
||||
|
||||
// Test 2: Call Ping function
|
||||
t.Run("CallPingFunction", func(t *testing.T) {
|
||||
handlerID := "agents.tests.mcpload.tools"
|
||||
handler, exists := process.Handlers[handlerID]
|
||||
assert.True(t, exists, "Handler should be registered")
|
||||
|
||||
p := &process.Process{
|
||||
ID: handlerID + ".Ping",
|
||||
Method: "Ping",
|
||||
Args: []interface{}{map[string]interface{}{"message": "test"}},
|
||||
Context: context.Background(),
|
||||
}
|
||||
|
||||
result := handler(p)
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultMap, ok := result.(map[string]interface{})
|
||||
assert.True(t, ok, "Result should be a map")
|
||||
assert.Equal(t, "test", resultMap["message"])
|
||||
assert.Contains(t, resultMap["echo"], "Pong")
|
||||
})
|
||||
|
||||
// Test 3: Call Calculate function
|
||||
t.Run("CallCalculateFunction", func(t *testing.T) {
|
||||
handlerID := "agents.tests.mcpload.tools"
|
||||
handler, exists := process.Handlers[handlerID]
|
||||
assert.True(t, exists, "Handler should be registered")
|
||||
|
||||
p := &process.Process{
|
||||
ID: handlerID + ".Calculate",
|
||||
Method: "Calculate",
|
||||
Args: []interface{}{map[string]interface{}{
|
||||
"operation": "add",
|
||||
"a": float64(10),
|
||||
"b": float64(5),
|
||||
}},
|
||||
Context: context.Background(),
|
||||
}
|
||||
|
||||
result := handler(p)
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultMap, ok := result.(map[string]interface{})
|
||||
assert.True(t, ok, "Result should be a map")
|
||||
assert.Equal(t, float64(15), resultMap["result"])
|
||||
})
|
||||
|
||||
// Test 4: Unregister scripts
|
||||
t.Run("UnregisterScripts", func(t *testing.T) {
|
||||
err := ast.UnregisterScripts()
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify handlers are removed
|
||||
handlerID := "agents.tests.mcpload.tools"
|
||||
_, exists := process.Handlers[handlerID]
|
||||
assert.False(t, exists, "Handler should be unregistered")
|
||||
})
|
||||
}
|
||||
|
||||
func TestScriptsProcessUsing(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Get the mcpload assistant
|
||||
assistantID := "tests.mcpload"
|
||||
ast, err := assistant.Get(assistantID)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, ast)
|
||||
|
||||
// Register scripts
|
||||
err = ast.RegisterScripts()
|
||||
assert.NoError(t, err)
|
||||
defer ast.UnregisterScripts()
|
||||
|
||||
// Test 1: Call Hello using process.New().Execute()
|
||||
t.Run("ProcessHello", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.Hello", map[string]interface{}{
|
||||
"name": "Yao",
|
||||
})
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NoError(t, err)
|
||||
|
||||
result := proc.Value()
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultStr, ok := result.(string)
|
||||
assert.True(t, ok, "Result should be a string")
|
||||
assert.Contains(t, resultStr, "Hello, Yao")
|
||||
})
|
||||
|
||||
// Test 2: Call Ping using process.New().Execute()
|
||||
t.Run("ProcessPing", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.Ping", map[string]interface{}{
|
||||
"message": "test message",
|
||||
})
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NoError(t, err)
|
||||
|
||||
result := proc.Value()
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultMap, ok := result.(map[string]interface{})
|
||||
assert.True(t, ok, "Result should be a map")
|
||||
assert.Equal(t, "test message", resultMap["message"])
|
||||
assert.Contains(t, resultMap["echo"], "Pong")
|
||||
})
|
||||
|
||||
// Test 3: Call Calculate using process.New().Execute()
|
||||
t.Run("ProcessCalculate", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.Calculate", map[string]interface{}{
|
||||
"operation": "multiply",
|
||||
"a": float64(6),
|
||||
"b": float64(7),
|
||||
})
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NoError(t, err)
|
||||
|
||||
result := proc.Value()
|
||||
assert.NotNil(t, result)
|
||||
|
||||
resultMap, ok := result.(map[string]interface{})
|
||||
assert.True(t, ok, "Result should be a map")
|
||||
assert.Equal(t, float64(42), resultMap["result"])
|
||||
assert.Equal(t, "multiply", resultMap["operation"])
|
||||
})
|
||||
}
|
||||
|
||||
func TestScriptsProcessError(t *testing.T) {
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Get the mcpload assistant
|
||||
assistantID := "tests.mcpload"
|
||||
ast, err := assistant.Get(assistantID)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, ast)
|
||||
|
||||
// Register scripts
|
||||
err = ast.RegisterScripts()
|
||||
assert.NoError(t, err)
|
||||
defer ast.UnregisterScripts()
|
||||
|
||||
// Test calling non-existent method
|
||||
t.Run("CallNonExistentMethod", func(t *testing.T) {
|
||||
proc := process.New("agents.tests.mcpload.tools.NonExistent")
|
||||
|
||||
err := proc.Execute()
|
||||
assert.NotNil(t, err, "Should return error when calling non-existent method")
|
||||
assert.Contains(t, err.Error(), "Exception|500")
|
||||
})
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue