637a0c9cbd
- Updated the agent's Stream and response processing methods to return and handle *context.Response directly, eliminating the need for type assertions. - Simplified test cases by removing unnecessary type conversions and directly accessing response fields. - Enhanced the extraction of data from Next hook responses, ensuring more robust handling of custom data structures. - Improved overall code readability and maintainability by streamlining response handling logic across various components.
107 lines
3.1 KiB
Go
107 lines
3.1 KiB
Go
package assistant_test
|
|
|
|
import (
|
|
stdContext "context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/yaoapp/yao/agent/assistant"
|
|
"github.com/yaoapp/yao/agent/context"
|
|
"github.com/yaoapp/yao/agent/output/message"
|
|
"github.com/yaoapp/yao/agent/testutils"
|
|
"github.com/yaoapp/yao/openapi/oauth/types"
|
|
)
|
|
|
|
// newSearchAutoHookDisableTestContext creates a test context
|
|
func newSearchAutoHookDisableTestContext(chatID, assistantID string) *context.Context {
|
|
authorized := &types.AuthorizedInfo{
|
|
Subject: "test-user",
|
|
UserID: "test-123",
|
|
TenantID: "test-tenant",
|
|
}
|
|
|
|
ctx := context.New(stdContext.Background(), authorized, chatID)
|
|
ctx.ID = chatID
|
|
ctx.AssistantID = assistantID
|
|
ctx.Locale = "en-us"
|
|
ctx.Client = context.Client{
|
|
Type: "web",
|
|
IP: "127.0.0.1",
|
|
}
|
|
ctx.Referer = context.RefererAPI
|
|
ctx.Accept = context.AcceptWebCUI
|
|
ctx.IDGenerator = message.NewIDGenerator()
|
|
ctx.Metadata = make(map[string]interface{})
|
|
return ctx
|
|
}
|
|
|
|
func TestSearchAutoHookDisable(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
ast, err := assistant.LoadPath("/assistants/tests/search-auto-hook-disable")
|
|
require.NoError(t, err)
|
|
require.NotNil(t, ast)
|
|
|
|
t.Run("ShouldHaveSearchConfigEnabled", func(t *testing.T) {
|
|
// Search config is enabled in package.yao
|
|
assert.NotNil(t, ast.Search, "search config should be set")
|
|
assert.NotNil(t, ast.Uses, "uses config should be set")
|
|
assert.Equal(t, "builtin", ast.Uses.Search, "uses.search should be builtin in config")
|
|
})
|
|
|
|
t.Run("ShouldHaveHookScript", func(t *testing.T) {
|
|
// Hook script should be loaded
|
|
assert.NotNil(t, ast.HookScript, "hook script should be loaded")
|
|
})
|
|
|
|
t.Run("HookShouldDisableSearch", func(t *testing.T) {
|
|
// Create context
|
|
ctx := newSearchAutoHookDisableTestContext("test-chat-id", "tests.search-auto-hook-disable")
|
|
|
|
// Create messages
|
|
messages := []context.Message{
|
|
{
|
|
Role: "user",
|
|
Content: "Test message",
|
|
},
|
|
}
|
|
|
|
// Call Create hook directly
|
|
opts := &context.Options{}
|
|
response, _, err := ast.HookScript.Create(ctx, messages, opts)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, response)
|
|
|
|
// Verify hook returns uses.search = "disabled"
|
|
assert.NotNil(t, response.Uses, "hook should return uses")
|
|
assert.Equal(t, "disabled", response.Uses.Search, "hook should disable search")
|
|
})
|
|
|
|
t.Run("StreamShouldRespectHookDisable", func(t *testing.T) {
|
|
// Get agent via assistant.Get (required for Stream)
|
|
agent, err := assistant.Get("tests.search-auto-hook-disable")
|
|
require.NoError(t, err)
|
|
require.NotNil(t, agent)
|
|
|
|
// Create context
|
|
ctx := newSearchAutoHookDisableTestContext("test-search-hook-disable", "tests.search-auto-hook-disable")
|
|
|
|
// Create messages
|
|
messages := []context.Message{
|
|
{
|
|
Role: "user",
|
|
Content: "What is AI?",
|
|
},
|
|
}
|
|
|
|
// Execute stream - hook will disable search
|
|
response, err := agent.Stream(ctx, messages)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, response)
|
|
|
|
assert.NotNil(t, response.Completion, "should have completion")
|
|
t.Logf("✓ Stream executed with hook disabling search")
|
|
})
|
|
}
|