632c0f5674
- Replaced all instances of `ctx.Space` with `ctx.Memory.Context` in the context management code, ensuring a more structured approach to handling temporary request-scoped data. - Updated related test cases to reflect the changes in context memory usage, enhancing the reliability and clarity of tests. - Removed the deprecated `Space` references and adjusted comments and documentation to align with the new memory management strategy.
351 lines
8.1 KiB
Go
351 lines
8.1 KiB
Go
package search_test
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/yaoapp/yao/agent/context"
|
|
"github.com/yaoapp/yao/agent/search"
|
|
"github.com/yaoapp/yao/agent/search/types"
|
|
"github.com/yaoapp/yao/agent/testutils"
|
|
)
|
|
|
|
func TestNewJSAPI(t *testing.T) {
|
|
api := search.NewJSAPI(nil, nil, nil)
|
|
require.NotNil(t, api)
|
|
}
|
|
|
|
func TestJSAPI_Web(t *testing.T) {
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
Web: &types.WebConfig{Provider: "tavily"},
|
|
}, &search.Uses{Web: "builtin"})
|
|
|
|
result := api.Web("test query", nil)
|
|
require.NotNil(t, result)
|
|
|
|
r, ok := result.(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeWeb, r.Type)
|
|
assert.Equal(t, "test query", r.Query)
|
|
assert.Equal(t, types.SourceHook, r.Source)
|
|
}
|
|
|
|
func TestJSAPI_Web_WithOptions(t *testing.T) {
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
Web: &types.WebConfig{Provider: "tavily"},
|
|
}, &search.Uses{Web: "builtin"})
|
|
|
|
opts := map[string]interface{}{
|
|
"limit": float64(5),
|
|
"sites": []interface{}{"github.com", "stackoverflow.com"},
|
|
"time_range": "week",
|
|
}
|
|
|
|
result := api.Web("golang concurrency", opts)
|
|
require.NotNil(t, result)
|
|
|
|
r, ok := result.(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeWeb, r.Type)
|
|
assert.Equal(t, "golang concurrency", r.Query)
|
|
}
|
|
|
|
func TestJSAPI_KB(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
|
}, nil)
|
|
|
|
result := api.KB("test query", nil)
|
|
require.NotNil(t, result)
|
|
|
|
r, ok := result.(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeKB, r.Type)
|
|
assert.Equal(t, "test query", r.Query)
|
|
assert.Equal(t, types.SourceHook, r.Source)
|
|
}
|
|
|
|
func TestJSAPI_KB_WithOptions(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
|
}, nil)
|
|
|
|
opts := map[string]interface{}{
|
|
"collections": []interface{}{"docs", "faq"},
|
|
"threshold": 0.8,
|
|
"limit": float64(10),
|
|
"graph": true,
|
|
}
|
|
|
|
result := api.KB("knowledge base query", opts)
|
|
require.NotNil(t, result)
|
|
|
|
r, ok := result.(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeKB, r.Type)
|
|
assert.Equal(t, "knowledge base query", r.Query)
|
|
}
|
|
|
|
func TestJSAPI_DB(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
DB: &types.DBConfig{Models: []string{"product"}},
|
|
}, &search.Uses{QueryDSL: "builtin"})
|
|
|
|
result := api.DB("test query", nil)
|
|
require.NotNil(t, result)
|
|
|
|
r, ok := result.(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeDB, r.Type)
|
|
assert.Equal(t, "test query", r.Query)
|
|
assert.Equal(t, types.SourceHook, r.Source)
|
|
}
|
|
|
|
func TestJSAPI_DB_WithOptions(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
DB: &types.DBConfig{Models: []string{"product"}},
|
|
}, &search.Uses{QueryDSL: "builtin"})
|
|
|
|
opts := map[string]interface{}{
|
|
"models": []interface{}{"product", "order"},
|
|
"select": []interface{}{"id", "name", "price"},
|
|
"limit": float64(20),
|
|
}
|
|
|
|
result := api.DB("database query", opts)
|
|
require.NotNil(t, result)
|
|
|
|
r, ok := result.(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeDB, r.Type)
|
|
assert.Equal(t, "database query", r.Query)
|
|
}
|
|
|
|
func TestJSAPI_All(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
|
DB: &types.DBConfig{Models: []string{"product"}},
|
|
}, nil)
|
|
|
|
requests := []interface{}{
|
|
map[string]interface{}{
|
|
"type": "kb",
|
|
"query": "KB query",
|
|
},
|
|
map[string]interface{}{
|
|
"type": "db",
|
|
"query": "DB query",
|
|
},
|
|
}
|
|
|
|
results := api.All(requests)
|
|
require.Len(t, results, 2)
|
|
|
|
// First result (KB)
|
|
r0, ok := results[0].(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeKB, r0.Type)
|
|
assert.Equal(t, "KB query", r0.Query)
|
|
|
|
// Second result (DB)
|
|
r1, ok := results[1].(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeDB, r1.Type)
|
|
assert.Equal(t, "DB query", r1.Query)
|
|
}
|
|
|
|
func TestJSAPI_Any(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
|
DB: &types.DBConfig{Models: []string{"product"}},
|
|
}, nil)
|
|
|
|
requests := []interface{}{
|
|
map[string]interface{}{
|
|
"type": "kb",
|
|
"query": "KB query",
|
|
},
|
|
map[string]interface{}{
|
|
"type": "db",
|
|
"query": "DB query",
|
|
},
|
|
}
|
|
|
|
results := api.Any(requests)
|
|
require.Len(t, results, 2)
|
|
|
|
// At least one result should be present
|
|
hasResult := false
|
|
for _, r := range results {
|
|
if r != nil {
|
|
hasResult = true
|
|
break
|
|
}
|
|
}
|
|
assert.True(t, hasResult)
|
|
}
|
|
|
|
func TestJSAPI_Race(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
KB: &types.KBConfig{Collections: []string{"docs"}},
|
|
DB: &types.DBConfig{Models: []string{"product"}},
|
|
}, nil)
|
|
|
|
requests := []interface{}{
|
|
map[string]interface{}{
|
|
"type": "kb",
|
|
"query": "KB query",
|
|
},
|
|
map[string]interface{}{
|
|
"type": "db",
|
|
"query": "DB query",
|
|
},
|
|
}
|
|
|
|
results := api.Race(requests)
|
|
require.Len(t, results, 2)
|
|
|
|
// At least one result should be present
|
|
hasResult := false
|
|
for _, r := range results {
|
|
if r != nil {
|
|
hasResult = true
|
|
break
|
|
}
|
|
}
|
|
assert.True(t, hasResult)
|
|
}
|
|
|
|
func TestJSAPI_All_Empty(t *testing.T) {
|
|
api := search.NewJSAPI(nil, nil, nil)
|
|
results := api.All([]interface{}{})
|
|
assert.Len(t, results, 0)
|
|
}
|
|
|
|
func TestJSAPI_Any_Empty(t *testing.T) {
|
|
api := search.NewJSAPI(nil, nil, nil)
|
|
results := api.Any([]interface{}{})
|
|
assert.Len(t, results, 0)
|
|
}
|
|
|
|
func TestJSAPI_Race_Empty(t *testing.T) {
|
|
api := search.NewJSAPI(nil, nil, nil)
|
|
results := api.Race([]interface{}{})
|
|
assert.Len(t, results, 0)
|
|
}
|
|
|
|
func TestJSAPI_Web_WithRerank(t *testing.T) {
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
Web: &types.WebConfig{Provider: "tavily"},
|
|
}, &search.Uses{Web: "builtin"})
|
|
|
|
opts := map[string]interface{}{
|
|
"limit": float64(10),
|
|
"rerank": map[string]interface{}{
|
|
"top_n": float64(5),
|
|
},
|
|
}
|
|
|
|
result := api.Web("test query", opts)
|
|
require.NotNil(t, result)
|
|
|
|
r, ok := result.(*types.Result)
|
|
require.True(t, ok)
|
|
assert.Equal(t, types.SearchTypeWeb, r.Type)
|
|
}
|
|
|
|
func TestJSAPI_All_InvalidRequests(t *testing.T) {
|
|
api := search.NewJSAPI(nil, &types.Config{
|
|
Web: &types.WebConfig{Provider: "tavily"},
|
|
}, &search.Uses{Web: "builtin"})
|
|
|
|
// Mix of invalid and valid requests
|
|
requests := []interface{}{
|
|
"invalid", // Not a map
|
|
map[string]interface{}{
|
|
"query": "no type", // Missing type
|
|
},
|
|
map[string]interface{}{
|
|
"type": "web", // Missing query
|
|
},
|
|
map[string]interface{}{
|
|
"type": "web",
|
|
"query": "valid query",
|
|
},
|
|
}
|
|
|
|
results := api.All(requests)
|
|
// Only the valid request should produce a result
|
|
assert.Len(t, results, 1)
|
|
}
|
|
|
|
func TestSetJSAPIFactory(t *testing.T) {
|
|
// Reset factory
|
|
context.SearchAPIFactory = nil
|
|
|
|
// Set factory with nil getter (uses defaults)
|
|
search.SetJSAPIFactory(nil)
|
|
|
|
// Verify factory is set
|
|
require.NotNil(t, context.SearchAPIFactory)
|
|
|
|
// Create a mock context
|
|
ctx := context.New(nil, nil, "test-chat")
|
|
|
|
// Get search API
|
|
searchAPI := context.SearchAPIFactory(ctx)
|
|
require.NotNil(t, searchAPI)
|
|
}
|
|
|
|
func TestSetJSAPIFactory_WithGetter(t *testing.T) {
|
|
// Reset factory
|
|
context.SearchAPIFactory = nil
|
|
|
|
// Set factory with custom getter
|
|
search.SetJSAPIFactory(func(assistantID string) (*types.Config, *search.Uses) {
|
|
if assistantID == "test-assistant" {
|
|
return &types.Config{
|
|
Web: &types.WebConfig{Provider: "tavily"},
|
|
}, &search.Uses{Web: "builtin"}
|
|
}
|
|
return nil, nil
|
|
})
|
|
|
|
// Verify factory is set
|
|
require.NotNil(t, context.SearchAPIFactory)
|
|
|
|
// Create a context with assistant ID
|
|
ctx := context.New(nil, nil, "test-chat")
|
|
ctx.AssistantID = "test-assistant"
|
|
|
|
// Get search API
|
|
searchAPI := context.SearchAPIFactory(ctx)
|
|
require.NotNil(t, searchAPI)
|
|
}
|
|
|
|
func TestJSAPI_ImplementsSearchAPI(t *testing.T) {
|
|
// Verify JSAPI implements context.SearchAPI interface
|
|
var _ context.SearchAPI = search.NewJSAPI(nil, nil, nil)
|
|
}
|