opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #1381 from trheyi/main

Browse source 
Add AI Agent Search Feature with JSAPI Integration & Smart Configuration
This commit is contained in:
Max 2025-12-13 17:26:19 +08:00 • committed by GitHub
commit 9684a9e563
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
85 changed files with 12271 additions and 578 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.github/workflows/pr-test.yml vendored
View file

@ -49,6 +49,11 @@ env:
DEEPSEEK_MODELS_V3: ${{ secrets.DEEPSEEK_MODELS_V3 }}
DEEPSEEK_MODELS_V3_1: ${{ secrets.DEEPSEEK_MODELS_V3_1 }}
# Search API Configuration
TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }}
SERPAPI_API_KEY: ${{ secrets.SERPAPI_API_KEY }}
SERPER_API_KEY: ${{ secrets.SERPER_API_KEY }}
# Claude API Configuration
CLAUDE_API_KEY: ${{ secrets.CLAUDE_API_KEY }}
CLAUDE_PROXY: ${{ secrets.CLAUDE_PROXY }}

5
.github/workflows/unit-test.yml vendored
View file

@ -53,6 +53,11 @@ env:
DEEPSEEK_MODELS_V3: ${{ secrets.DEEPSEEK_MODELS_V3 }}
DEEPSEEK_MODELS_V3_1: ${{ secrets.DEEPSEEK_MODELS_V3_1 }}
# Search API Configuration
TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }}
SERPAPI_API_KEY: ${{ secrets.SERPAPI_API_KEY }}
SERPER_API_KEY: ${{ secrets.SERPER_API_KEY }}
# Claude API Configuration
CLAUDE_API_KEY: ${{ secrets.CLAUDE_API_KEY }}
CLAUDE_PROXY: ${{ secrets.CLAUDE_PROXY }}

2
.gitignore vendored
View file

@ -41,6 +41,7 @@ xgen/v1.0/*
*-unit-test
docker/build/test
db
!agent/search/handlers/db
*.sh
data/bindata.go.bak
share/const.go.bak
@ -49,3 +50,4 @@ share/const.goe
openapi/*.md
coverage.html
agent/assistant/hook/*.test.md
agent/search/TODO.md

14
agent/assistant/agent.go
View file

@ -199,6 +199,16 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
return nil, err
}
// ================================================
// Execute Auto Search (if enabled)
// ================================================
if ast.shouldAutoSearch(ctx, createResponse) {
refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse, opts)
if refCtx != nil && len(refCtx.References) > 0 {
completionMessages = ast.injectSearchContext(completionMessages, refCtx)
}
}
// Begin step tracking for LLM call
ast.BeginStep(ctx, context.StepTypeLLM, map[string]interface{}{
"messages": completionMessages,
@ -367,6 +377,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
var nextResponse *context.NextHookResponse = nil
if ast.HookScript != nil {
ctx.Logger.HookStart("Next")
// Begin step tracking for hook_next
ast.BeginStep(ctx, context.StepTypeHookNext, map[string]interface{}{
"messages": fullMessages,
@ -393,6 +405,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
"response": nextResponse,
})
ctx.Logger.HookComplete("Next")
// Process Next hook response
finalResponse, err = ast.processNextResponse(&NextProcessContext{
Context: ctx,

135
agent/assistant/assistant.go
View file

@ -5,16 +5,18 @@ import (
"path"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/content"
"github.com/yaoapp/yao/agent/caller"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/search"
searchTypes "github.com/yaoapp/yao/agent/search/types"
store "github.com/yaoapp/yao/agent/store/types"
sui "github.com/yaoapp/yao/sui/core"
)
func init() {
// Initialize AgentGetterFunc to allow content package to call agents
content.AgentGetterFunc = func(agentID string) (content.AgentCaller, error) {
// Initialize AgentGetterFunc to allow content and search packages to call agents
caller.AgentGetterFunc = func(agentID string) (caller.AgentCaller, error) {
ast, err := Get(agentID)
if err != nil {
return nil, err
@ -22,6 +24,26 @@ func init() {
// Return a wrapper that implements AgentCaller interface
return &agentCallerWrapper{ast: ast}, nil
}
// Initialize Search JSAPI factory with config getter
search.SetJSAPIFactory(func(assistantID string) (*searchTypes.Config, *search.Uses) {
ast, err := Get(assistantID)
if err != nil || ast == nil {
return nil, nil
}
// Convert assistant.Uses to search.Uses
var uses *search.Uses
if ast.Uses != nil {
uses = &search.Uses{
Search: ast.Uses.Search,
Web: ast.Uses.Web,
Keyword: ast.Uses.Keyword,
QueryDSL: ast.Uses.QueryDSL,
Rerank: ast.Uses.Rerank,
}
}
return ast.Search, uses
})
}
// agentCallerWrapper wraps Assistant to implement AgentCaller interface
@ -115,6 +137,8 @@ func (ast *Assistant) Map() map[string]interface{} {
"automated": ast.Automated,
"placeholder": ast.Placeholder,
"locales": ast.Locales,
"uses": ast.Uses,
"search": ast.Search,
"created_at": store.ToMySQLTime(ast.CreatedAt),
"updated_at": store.ToMySQLTime(ast.UpdatedAt),
}
@ -183,7 +207,6 @@ func (ast *Assistant) Clone() *Assistant {
CreatedAt: ast.CreatedAt,
UpdatedAt: ast.UpdatedAt,
},
Search: ast.Search,
HookScript: ast.HookScript,
openai: ast.openai,
}
@ -346,6 +369,86 @@ func (ast *Assistant) Clone() *Assistant {
}
}
// Deep copy uses
if ast.Uses != nil {
clone.Uses = &agentContext.Uses{
Vision: ast.Uses.Vision,
Audio: ast.Uses.Audio,
Search: ast.Uses.Search,
Fetch: ast.Uses.Fetch,
Web: ast.Uses.Web,
Keyword: ast.Uses.Keyword,
QueryDSL: ast.Uses.QueryDSL,
Rerank: ast.Uses.Rerank,
}
}
// Deep copy search config
if ast.Search != nil {
clone.Search = &searchTypes.Config{}
if ast.Search.Web != nil {
clone.Search.Web = &searchTypes.WebConfig{
Provider: ast.Search.Web.Provider,
APIKeyEnv: ast.Search.Web.APIKeyEnv,
MaxResults: ast.Search.Web.MaxResults,
}
}
if ast.Search.KB != nil {
clone.Search.KB = &searchTypes.KBConfig{
Threshold: ast.Search.KB.Threshold,
Graph: ast.Search.KB.Graph,
}
if ast.Search.KB.Collections != nil {
clone.Search.KB.Collections = make([]string, len(ast.Search.KB.Collections))
copy(clone.Search.KB.Collections, ast.Search.KB.Collections)
}
}
if ast.Search.DB != nil {
clone.Search.DB = &searchTypes.DBConfig{
MaxResults: ast.Search.DB.MaxResults,
}
if ast.Search.DB.Models != nil {
clone.Search.DB.Models = make([]string, len(ast.Search.DB.Models))
copy(clone.Search.DB.Models, ast.Search.DB.Models)
}
}
if ast.Search.Keyword != nil {
clone.Search.Keyword = &searchTypes.KeywordConfig{
MaxKeywords: ast.Search.Keyword.MaxKeywords,
Language: ast.Search.Keyword.Language,
}
}
if ast.Search.QueryDSL != nil {
clone.Search.QueryDSL = &searchTypes.QueryDSLConfig{
Strict: ast.Search.QueryDSL.Strict,
}
}
if ast.Search.Rerank != nil {
clone.Search.Rerank = &searchTypes.RerankConfig{
TopN: ast.Search.Rerank.TopN,
}
}
if ast.Search.Citation != nil {
clone.Search.Citation = &searchTypes.CitationConfig{
Format: ast.Search.Citation.Format,
AutoInjectPrompt: ast.Search.Citation.AutoInjectPrompt,
CustomPrompt: ast.Search.Citation.CustomPrompt,
}
}
if ast.Search.Weights != nil {
clone.Search.Weights = &searchTypes.WeightsConfig{
User: ast.Search.Weights.User,
Hook: ast.Search.Weights.Hook,
Auto: ast.Search.Weights.Auto,
}
}
if ast.Search.Options != nil {
clone.Search.Options = &searchTypes.OptionsConfig{
SkipThreshold: ast.Search.Options.SkipThreshold,
}
}
}
return clone
}
@ -512,5 +615,29 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
ast.Workflow = workflow
}
// Uses
if v, has := data["uses"]; has {
uses, err := store.ToUses(v)
if err != nil {
return err
}
ast.Uses = uses
}
// Search
if v, has := data["search"]; has {
search, err := store.ToSearchConfig(v)
if err != nil {
return err
}
ast.Search = search
}
return ast.Validate()
}
// GetMergedSearchConfig returns the search config for this assistant
// Note: The config is already merged with global config during loading (loadMap)
func (ast *Assistant) GetMergedSearchConfig() *searchTypes.Config {
return ast.Search
}

36
agent/assistant/build.go
View file

@ -532,40 +532,10 @@ func (ast *Assistant) applyCreateResponseOptions(options *context.CompletionOpti
// getUses get the Uses configuration with priority: assistant.Uses > global settings
// Note: createResponse.Uses (applied in applyCreateResponseOptions) has even higher priority
// Final priority order: createResponse.Uses > assistant.Uses > global settings
// getUses returns the Uses config for this assistant
// Note: The config is already merged with global config during loading (loadMap)
func (ast *Assistant) getUses() *context.Uses {
// Priority 1: Assistant-specific Uses configuration
if ast.Uses != nil {
// Create a merged Uses by starting with global, then override with assistant-specific
merged := &context.Uses{}
// Start with global settings
if globalUses != nil {
merged.Vision = globalUses.Vision
merged.Audio = globalUses.Audio
merged.Search = globalUses.Search
merged.Fetch = globalUses.Fetch
}
// Override with assistant-specific settings (only if not empty)
if ast.Uses.Vision != "" {
merged.Vision = ast.Uses.Vision
}
if ast.Uses.Audio != "" {
merged.Audio = ast.Uses.Audio
}
if ast.Uses.Search != "" {
merged.Search = ast.Uses.Search
}
if ast.Uses.Fetch != "" {
merged.Fetch = ast.Uses.Fetch
}
return merged
}
// Priority 2: Global settings only
return globalUses
return ast.Uses
}
// applyMCPTools adds MCP tools to completion options and returns samples prompt

250
agent/assistant/load.go
View file

@ -13,6 +13,7 @@ import (
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
searchTypes "github.com/yaoapp/yao/agent/search/types"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/openai"
"gopkg.in/yaml.v3"
@ -22,12 +23,12 @@ import (
var loaded = NewCache(200) // 200 is the default capacity
var storage store.Store = nil
var storeSetting *store.Setting = nil // store setting from agent.yml
var search interface{} = nil
var modelCapabilities map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{}
var defaultConnector string = "" // default connector
var globalUses *context.Uses = nil // global uses configuration from agent.yml
var globalPrompts []store.Prompt = nil // global prompts from agent/prompts.yml
var globalKBSetting *store.KBSetting = nil // global KB setting from agent/kb.yml
var defaultConnector string = "" // default connector
var globalUses *context.Uses = nil // global uses configuration from agent.yml
var globalPrompts []store.Prompt = nil // global prompts from agent/prompts.yml
var globalKBSetting *store.KBSetting = nil // global KB setting from agent/kb.yml
var globalSearchConfig *searchTypes.Config = nil // global search config from agent/search.yml
// LoadBuiltIn load the built-in assistants
func LoadBuiltIn() error {
@ -183,6 +184,16 @@ func GetGlobalKBSetting() *store.KBSetting {
return globalKBSetting
}
// SetGlobalSearchConfig set the global search config from agent/search.yml
func SetGlobalSearchConfig(config *searchTypes.Config) {
globalSearchConfig = config
}
// GetGlobalSearchConfig returns the global search config
func GetGlobalSearchConfig() *searchTypes.Config {
return globalSearchConfig
}
// SetCache set the cache
func SetCache(capacity int) {
ClearCache()
@ -585,19 +596,24 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
}
}
// Search options
// Search configuration (from package.yao search block)
// This contains search options like web.max_results, kb.threshold, citation.format, etc.
// Merge hierarchy: global config < assistant config
if v, ok := data["search"].(map[string]interface{}); ok {
assistant.Search = &SearchOption{}
var assistantSearch searchTypes.Config
raw, err := jsoniter.Marshal(v)
if err != nil {
return nil, err
}
// Unmarshal the raw data
err = jsoniter.Unmarshal(raw, assistant.Search)
err = jsoniter.Unmarshal(raw, &assistantSearch)
if err != nil {
return nil, err
}
// Merge with global search config
assistant.Search = mergeSearchConfig(globalSearchConfig, &assistantSearch)
} else if globalSearchConfig != nil {
// No assistant-specific config, use global
assistant.Search = globalSearchConfig
}
// prompts
@ -681,12 +697,14 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
}
// uses (wrapper configurations for vision, audio, etc.)
// Merge hierarchy: global uses < assistant uses
if uses, has := data["uses"]; has {
var assistantUses *context.Uses
switch v := uses.(type) {
case *context.Uses:
assistant.Uses = v
assistantUses = v
case context.Uses:
assistant.Uses = &v
assistantUses = &v
default:
raw, err := jsoniter.Marshal(v)
if err != nil {
@ -697,8 +715,13 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
if err != nil {
return nil, err
}
assistant.Uses = &usesConfig
assistantUses = &usesConfig
}
// Merge with global uses
assistant.Uses = mergeUses(globalUses, assistantUses)
} else if globalUses != nil {
// No assistant-specific uses, use global
assistant.Uses = globalUses
}
// Load scripts (hook script and other scripts)
@ -761,3 +784,204 @@ func (ast *Assistant) initialize() error {
return nil
}
// mergeUses merges two Uses configs (base < override)
func mergeUses(base, override *context.Uses) *context.Uses {
if base == nil {
return override
}
if override == nil {
return base
}
result := *base // Copy base
// Override with non-empty values
if override.Vision != "" {
result.Vision = override.Vision
}
if override.Audio != "" {
result.Audio = override.Audio
}
if override.Search != "" {
result.Search = override.Search
}
if override.Fetch != "" {
result.Fetch = override.Fetch
}
if override.Web != "" {
result.Web = override.Web
}
if override.Keyword != "" {
result.Keyword = override.Keyword
}
if override.QueryDSL != "" {
result.QueryDSL = override.QueryDSL
}
if override.Rerank != "" {
result.Rerank = override.Rerank
}
return &result
}
// mergeSearchConfig merges two search configs (base < override)
func mergeSearchConfig(base, override *searchTypes.Config) *searchTypes.Config {
if base == nil {
return override
}
if override == nil {
return base
}
result := *base // Copy base
// Merge Web config
if override.Web != nil {
if result.Web == nil {
result.Web = override.Web
} else {
merged := *result.Web
if override.Web.Provider != "" {
merged.Provider = override.Web.Provider
}
if override.Web.APIKeyEnv != "" {
merged.APIKeyEnv = override.Web.APIKeyEnv
}
if override.Web.MaxResults > 0 {
merged.MaxResults = override.Web.MaxResults
}
result.Web = &merged
}
}
// Merge KB config
if override.KB != nil {
if result.KB == nil {
result.KB = override.KB
} else {
merged := *result.KB
if len(override.KB.Collections) > 0 {
merged.Collections = override.KB.Collections
}
if override.KB.Threshold > 0 {
merged.Threshold = override.KB.Threshold
}
if override.KB.Graph {
merged.Graph = override.KB.Graph
}
result.KB = &merged
}
}
// Merge DB config
if override.DB != nil {
if result.DB == nil {
result.DB = override.DB
} else {
merged := *result.DB
if len(override.DB.Models) > 0 {
merged.Models = override.DB.Models
}
if override.DB.MaxResults > 0 {
merged.MaxResults = override.DB.MaxResults
}
result.DB = &merged
}
}
// Merge Keyword config
if override.Keyword != nil {
if result.Keyword == nil {
result.Keyword = override.Keyword
} else {
merged := *result.Keyword
if override.Keyword.MaxKeywords > 0 {
merged.MaxKeywords = override.Keyword.MaxKeywords
}
if override.Keyword.Language != "" {
merged.Language = override.Keyword.Language
}
result.Keyword = &merged
}
}
// Merge QueryDSL config
if override.QueryDSL != nil {
if result.QueryDSL == nil {
result.QueryDSL = override.QueryDSL
} else {
merged := *result.QueryDSL
if override.QueryDSL.Strict {
merged.Strict = override.QueryDSL.Strict
}
result.QueryDSL = &merged
}
}
// Merge Rerank config
if override.Rerank != nil {
if result.Rerank == nil {
result.Rerank = override.Rerank
} else {
merged := *result.Rerank
if override.Rerank.TopN > 0 {
merged.TopN = override.Rerank.TopN
}
result.Rerank = &merged
}
}
// Merge Citation config
if override.Citation != nil {
if result.Citation == nil {
result.Citation = override.Citation
} else {
merged := *result.Citation
if override.Citation.Format != "" {
merged.Format = override.Citation.Format
}
// AutoInjectPrompt is a bool, so we check if it's explicitly set
// by checking if the whole Citation block was provided
merged.AutoInjectPrompt = override.Citation.AutoInjectPrompt
if override.Citation.CustomPrompt != "" {
merged.CustomPrompt = override.Citation.CustomPrompt
}
result.Citation = &merged
}
}
// Merge Weights config
if override.Weights != nil {
if result.Weights == nil {
result.Weights = override.Weights
} else {
merged := *result.Weights
if override.Weights.User > 0 {
merged.User = override.Weights.User
}
if override.Weights.Hook > 0 {
merged.Hook = override.Weights.Hook
}
if override.Weights.Auto > 0 {
merged.Auto = override.Weights.Auto
}
result.Weights = &merged
}
}
// Merge Options config
if override.Options != nil {
if result.Options == nil {
result.Options = override.Options
} else {
merged := *result.Options
if override.Options.SkipThreshold > 0 {
merged.SkipThreshold = override.Options.SkipThreshold
}
result.Options = &merged
}
}
return &result
}

366
agent/assistant/load_merge_test.go Normal file
View file

@ -0,0 +1,366 @@
package assistant_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/testutils"
)
// TestLoadPathMerge tests loading the merge test assistant
// This verifies that global config is properly merged with assistant-specific config
func TestLoadPathMerge(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/merge")
require.NoError(t, err)
require.NotNil(t, ast)
assert.Equal(t, "tests.merge", ast.ID)
assert.Equal(t, "Merge Config Test Assistant", ast.Name)
// Uses configuration - should merge global with assistant-specific
// Global (from agent/agent.yml):
// vision: "workers.system.vision"
// search: "workers.system.search"
// fetch: "workers.system.fetch"
// audio: (not set)
// querydsl: (not set)
// rerank: (not set)
// Assistant:
// web: "mcp:custom-web"
// keyword: "mcp:custom-keyword"
// Result: assistant values override, global values inherited
assert.NotNil(t, ast.Uses)
// Assistant overrides
assert.Equal(t, "mcp:custom-web", ast.Uses.Web) // overridden by assistant
assert.Equal(t, "mcp:custom-keyword", ast.Uses.Keyword) // overridden by assistant
// Inherited from global (agent/agent.yml)
assert.Equal(t, "workers.system.vision", ast.Uses.Vision) // inherited from global
assert.Equal(t, "workers.system.search", ast.Uses.Search) // inherited from global
assert.Equal(t, "workers.system.fetch", ast.Uses.Fetch) // inherited from global
// Not set in either global or assistant (should be empty)
assert.Empty(t, ast.Uses.Audio) // not set anywhere
assert.Empty(t, ast.Uses.QueryDSL) // not set anywhere
assert.Empty(t, ast.Uses.Rerank) // not set anywhere
// Search configuration - should merge global with assistant-specific
// Global (from agent/search.yml):
// web.provider=tavily, web.max_results=10
// kb.threshold=0.7, kb.graph=false
// db.max_results=20
// keyword.max_keywords=10, keyword.language=auto
// rerank.top_n=10
// citation.format=#ref:{id}, citation.auto_inject_prompt=true
// weights: user=1.0, hook=0.8, auto=0.6
// options.skip_threshold=5
// Assistant:
// web.provider=custom-provider, web.max_results=25
// kb.collections=[merge-test-kb], kb.threshold=0.85
assert.NotNil(t, ast.Search)
// Web config - assistant overrides global
assert.NotNil(t, ast.Search.Web)
assert.Equal(t, "custom-provider", ast.Search.Web.Provider) // overridden
assert.Equal(t, 25, ast.Search.Web.MaxResults) // overridden
// KB config - assistant overrides global
assert.NotNil(t, ast.Search.KB)
assert.Equal(t, []string{"merge-test-kb"}, ast.Search.KB.Collections) // overridden
assert.Equal(t, 0.85, ast.Search.KB.Threshold) // overridden
assert.False(t, ast.Search.KB.Graph) // inherited from global
// DB config - should inherit from global (assistant doesn't define it)
assert.NotNil(t, ast.Search.DB)
assert.Equal(t, 20, ast.Search.DB.MaxResults) // inherited from global
// Keyword config - should inherit from global
assert.NotNil(t, ast.Search.Keyword)
assert.Equal(t, 10, ast.Search.Keyword.MaxKeywords) // inherited from global
assert.Equal(t, "auto", ast.Search.Keyword.Language) // inherited from global
// Rerank config - should inherit from global
assert.NotNil(t, ast.Search.Rerank)
assert.Equal(t, 10, ast.Search.Rerank.TopN) // inherited from global
// Citation config - should inherit from global
assert.NotNil(t, ast.Search.Citation)
assert.Equal(t, "#ref:{id}", ast.Search.Citation.Format) // inherited from global
assert.True(t, ast.Search.Citation.AutoInjectPrompt) // inherited from global
// Weights config - should inherit from global
assert.NotNil(t, ast.Search.Weights)
assert.Equal(t, 1.0, ast.Search.Weights.User) // inherited from global
assert.Equal(t, 0.8, ast.Search.Weights.Hook) // inherited from global
assert.Equal(t, 0.6, ast.Search.Weights.Auto) // inherited from global
// Options config - should inherit from global
assert.NotNil(t, ast.Search.Options)
assert.Equal(t, 5, ast.Search.Options.SkipThreshold) // inherited from global
}
// TestLoadPathMergeOverride tests loading the merge-override test assistant
// This verifies that assistant config completely overrides global config
func TestLoadPathMergeOverride(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/merge-override")
require.NoError(t, err)
require.NotNil(t, ast)
assert.Equal(t, "tests.merge-override", ast.ID)
assert.Equal(t, "Merge Override Test Assistant", ast.Name)
// Uses configuration - all fields should be overridden by assistant
assert.NotNil(t, ast.Uses)
assert.Equal(t, "mcp:custom-vision", ast.Uses.Vision)
assert.Equal(t, "mcp:custom-audio", ast.Uses.Audio)
assert.Equal(t, "mcp:custom-search", ast.Uses.Search)
assert.Equal(t, "mcp:custom-fetch", ast.Uses.Fetch)
assert.Equal(t, "mcp:custom-web", ast.Uses.Web)
assert.Equal(t, "mcp:custom-keyword", ast.Uses.Keyword)
assert.Equal(t, "mcp:custom-querydsl", ast.Uses.QueryDSL)
assert.Equal(t, "mcp:custom-rerank", ast.Uses.Rerank)
// Search configuration - all fields should be overridden by assistant
assert.NotNil(t, ast.Search)
// Web config - all overridden
assert.NotNil(t, ast.Search.Web)
assert.Equal(t, "override-provider", ast.Search.Web.Provider)
assert.Equal(t, "$ENV.OVERRIDE_API_KEY", ast.Search.Web.APIKeyEnv)
assert.Equal(t, 100, ast.Search.Web.MaxResults)
// KB config - all overridden
assert.NotNil(t, ast.Search.KB)
assert.Equal(t, []string{"override-kb"}, ast.Search.KB.Collections)
assert.Equal(t, 0.99, ast.Search.KB.Threshold)
assert.True(t, ast.Search.KB.Graph)
// DB config - all overridden
assert.NotNil(t, ast.Search.DB)
assert.Equal(t, []string{"override-model"}, ast.Search.DB.Models)
assert.Equal(t, 200, ast.Search.DB.MaxResults)
// Keyword config - all overridden
assert.NotNil(t, ast.Search.Keyword)
assert.Equal(t, 20, ast.Search.Keyword.MaxKeywords)
assert.Equal(t, "zh", ast.Search.Keyword.Language)
// QueryDSL config - overridden
assert.NotNil(t, ast.Search.QueryDSL)
assert.True(t, ast.Search.QueryDSL.Strict)
// Rerank config - overridden
assert.NotNil(t, ast.Search.Rerank)
assert.Equal(t, 20, ast.Search.Rerank.TopN)
// Citation config - all overridden
assert.NotNil(t, ast.Search.Citation)
assert.Equal(t, "[override:{id}]", ast.Search.Citation.Format)
assert.False(t, ast.Search.Citation.AutoInjectPrompt)
assert.Equal(t, "Override citation prompt", ast.Search.Citation.CustomPrompt)
// Weights config - all overridden
assert.NotNil(t, ast.Search.Weights)
assert.Equal(t, 2.0, ast.Search.Weights.User)
assert.Equal(t, 1.5, ast.Search.Weights.Hook)
assert.Equal(t, 1.0, ast.Search.Weights.Auto)
// Options config - overridden
assert.NotNil(t, ast.Search.Options)
assert.Equal(t, 10, ast.Search.Options.SkipThreshold)
}
// TestLoadPathMergeEmpty tests loading the merge-empty test assistant
// This verifies that assistant with no uses/search config inherits all from global
func TestLoadPathMergeEmpty(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/merge-empty")
require.NoError(t, err)
require.NotNil(t, ast)
assert.Equal(t, "tests.merge-empty", ast.ID)
assert.Equal(t, "Merge Empty Test Assistant", ast.Name)
// Uses configuration - all inherited from global (agent/agent.yml)
assert.NotNil(t, ast.Uses)
assert.Equal(t, "workers.system.vision", ast.Uses.Vision) // from global
assert.Equal(t, "workers.system.search", ast.Uses.Search) // from global
assert.Equal(t, "workers.system.fetch", ast.Uses.Fetch) // from global
assert.Empty(t, ast.Uses.Audio) // not set in global
assert.Empty(t, ast.Uses.Web) // not set in global
assert.Empty(t, ast.Uses.Keyword) // not set in global
assert.Empty(t, ast.Uses.QueryDSL) // not set in global
assert.Empty(t, ast.Uses.Rerank) // not set in global
// Search configuration - all inherited from global (agent/search.yml)
assert.NotNil(t, ast.Search)
// Web config - from global
assert.NotNil(t, ast.Search.Web)
assert.Equal(t, "tavily", ast.Search.Web.Provider)
assert.Equal(t, 10, ast.Search.Web.MaxResults)
// KB config - from global
assert.NotNil(t, ast.Search.KB)
assert.Equal(t, 0.7, ast.Search.KB.Threshold)
assert.False(t, ast.Search.KB.Graph)
// DB config - from global
assert.NotNil(t, ast.Search.DB)
assert.Equal(t, 20, ast.Search.DB.MaxResults)
// Keyword config - from global
assert.NotNil(t, ast.Search.Keyword)
assert.Equal(t, 10, ast.Search.Keyword.MaxKeywords)
assert.Equal(t, "auto", ast.Search.Keyword.Language)
// Rerank config - from global
assert.NotNil(t, ast.Search.Rerank)
assert.Equal(t, 10, ast.Search.Rerank.TopN)
// Citation config - from global
assert.NotNil(t, ast.Search.Citation)
assert.Equal(t, "#ref:{id}", ast.Search.Citation.Format)
assert.True(t, ast.Search.Citation.AutoInjectPrompt)
// Weights config - from global
assert.NotNil(t, ast.Search.Weights)
assert.Equal(t, 1.0, ast.Search.Weights.User)
assert.Equal(t, 0.8, ast.Search.Weights.Hook)
assert.Equal(t, 0.6, ast.Search.Weights.Auto)
// Options config - from global
assert.NotNil(t, ast.Search.Options)
assert.Equal(t, 5, ast.Search.Options.SkipThreshold)
}
// TestLoadPathUsesAndSearchMerge tests loading fullfields assistant
// This verifies that uses and search configs are properly loaded and merged
func TestLoadPathUsesAndSearchMerge(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, ast)
// Uses configuration - assistant-specific values
assert.NotNil(t, ast.Uses)
assert.Equal(t, "agent", ast.Uses.Vision)
assert.Equal(t, "mcp:audio-server", ast.Uses.Audio)
assert.Equal(t, "agent", ast.Uses.Fetch)
assert.Equal(t, "builtin", ast.Uses.Web)
assert.Equal(t, "builtin", ast.Uses.Keyword)
assert.Equal(t, "builtin", ast.Uses.QueryDSL)
assert.Equal(t, "builtin", ast.Uses.Rerank)
// Search configuration - assistant-specific values
assert.NotNil(t, ast.Search)
// Web config - from assistant
assert.NotNil(t, ast.Search.Web)
assert.Equal(t, "tavily", ast.Search.Web.Provider)
assert.Equal(t, 15, ast.Search.Web.MaxResults)
// KB config - from assistant
assert.NotNil(t, ast.Search.KB)
assert.Equal(t, []string{"docs", "faq"}, ast.Search.KB.Collections)
assert.Equal(t, 0.8, ast.Search.KB.Threshold)
assert.True(t, ast.Search.KB.Graph)
// DB config - from assistant
assert.NotNil(t, ast.Search.DB)
assert.Equal(t, []string{"user", "product"}, ast.Search.DB.Models)
assert.Equal(t, 50, ast.Search.DB.MaxResults)
// Citation config - from assistant
assert.NotNil(t, ast.Search.Citation)
assert.Equal(t, "#ref:{id}", ast.Search.Citation.Format)
assert.True(t, ast.Search.Citation.AutoInjectPrompt)
// Weights config - from assistant
assert.NotNil(t, ast.Search.Weights)
assert.Equal(t, 1.0, ast.Search.Weights.User)
assert.Equal(t, 0.9, ast.Search.Weights.Hook)
assert.Equal(t, 0.7, ast.Search.Weights.Auto)
}
// TestLoadPathSearchAssistant tests loading the dedicated search test assistant
func TestLoadPathSearchAssistant(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/search")
require.NoError(t, err)
require.NotNil(t, ast)
assert.Equal(t, "tests.search", ast.ID)
assert.Equal(t, "Search Config Test Assistant", ast.Name)
// Uses configuration
assert.NotNil(t, ast.Uses)
assert.Equal(t, "builtin", ast.Uses.Web)
assert.Equal(t, "builtin", ast.Uses.Keyword)
assert.Equal(t, "builtin", ast.Uses.QueryDSL)
assert.Equal(t, "builtin", ast.Uses.Rerank)
// Search configuration
assert.NotNil(t, ast.Search)
// Web config
assert.NotNil(t, ast.Search.Web)
assert.Equal(t, "serper", ast.Search.Web.Provider)
assert.Equal(t, "$ENV.SERPER_API_KEY", ast.Search.Web.APIKeyEnv)
assert.Equal(t, 20, ast.Search.Web.MaxResults)
// KB config
assert.NotNil(t, ast.Search.KB)
assert.Equal(t, []string{"knowledge-base", "documents"}, ast.Search.KB.Collections)
assert.Equal(t, 0.75, ast.Search.KB.Threshold)
assert.False(t, ast.Search.KB.Graph)
// DB config
assert.NotNil(t, ast.Search.DB)
assert.Equal(t, []string{"article", "comment"}, ast.Search.DB.Models)
assert.Equal(t, 30, ast.Search.DB.MaxResults)
// Keyword config
assert.NotNil(t, ast.Search.Keyword)
assert.Equal(t, 8, ast.Search.Keyword.MaxKeywords)
assert.Equal(t, "auto", ast.Search.Keyword.Language)
// QueryDSL config
assert.NotNil(t, ast.Search.QueryDSL)
assert.True(t, ast.Search.QueryDSL.Strict)
// Rerank config
assert.NotNil(t, ast.Search.Rerank)
assert.Equal(t, 5, ast.Search.Rerank.TopN)
// Citation config
assert.NotNil(t, ast.Search.Citation)
assert.Equal(t, "#cite:{id}", ast.Search.Citation.Format)
assert.False(t, ast.Search.Citation.AutoInjectPrompt)
assert.Equal(t, "Please cite sources using #cite:{id} format.", ast.Search.Citation.CustomPrompt)
// Weights config
assert.NotNil(t, ast.Search.Weights)
assert.Equal(t, 1.0, ast.Search.Weights.User)
assert.Equal(t, 0.85, ast.Search.Weights.Hook)
assert.Equal(t, 0.65, ast.Search.Weights.Auto)
// Options config
assert.NotNil(t, ast.Search.Options)
assert.Equal(t, 3, ast.Search.Options.SkipThreshold)
}

244
agent/assistant/load_store_test.go
View file

@ -9,6 +9,7 @@ import (
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/context"
searchTypes "github.com/yaoapp/yao/agent/search/types"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/agent/testutils"
"github.com/yaoapp/yao/openapi/oauth/types"
@ -933,3 +934,246 @@ function Create(ctx: any, messages: any[]): any {
assert.False(t, *res.DisableGlobalPrompts)
})
}
// TestLoadStoreWithSearchConfig tests loading assistant with search configuration from database
func TestLoadStoreWithSearchConfig(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-with-search"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test With Search Config",
Type: "assistant",
Connector: "gpt-4o",
Uses: &context.Uses{
Vision: "agent",
Audio: "mcp:audio-server",
Fetch: "agent",
Web: "builtin",
Keyword: "builtin",
QueryDSL: "builtin",
Rerank: "builtin",
},
Search: &searchTypes.Config{
Web: &searchTypes.WebConfig{
Provider: "tavily",
MaxResults: 15,
},
KB: &searchTypes.KBConfig{
Collections: []string{"docs", "faq"},
Threshold: 0.8,
Graph: true,
},
DB: &searchTypes.DBConfig{
Models: []string{"user", "product"},
MaxResults: 50,
},
Keyword: &searchTypes.KeywordConfig{
MaxKeywords: 8,
Language: "auto",
},
QueryDSL: &searchTypes.QueryDSLConfig{
Strict: true,
},
Rerank: &searchTypes.RerankConfig{
TopN: 5,
},
Citation: &searchTypes.CitationConfig{
Format: "#cite:{id}",
AutoInjectPrompt: false,
CustomPrompt: "Please cite sources.",
},
Weights: &searchTypes.WeightsConfig{
User: 1.0,
Hook: 0.85,
Auto: 0.65,
},
Options: &searchTypes.OptionsConfig{
SkipThreshold: 3,
},
},
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
// Verify Uses
require.NotNil(t, loaded.Uses)
assert.Equal(t, "agent", loaded.Uses.Vision)
assert.Equal(t, "mcp:audio-server", loaded.Uses.Audio)
assert.Equal(t, "agent", loaded.Uses.Fetch)
assert.Equal(t, "builtin", loaded.Uses.Web)
assert.Equal(t, "builtin", loaded.Uses.Keyword)
assert.Equal(t, "builtin", loaded.Uses.QueryDSL)
assert.Equal(t, "builtin", loaded.Uses.Rerank)
// Verify Search config
require.NotNil(t, loaded.Search)
// Web config
require.NotNil(t, loaded.Search.Web)
assert.Equal(t, "tavily", loaded.Search.Web.Provider)
assert.Equal(t, 15, loaded.Search.Web.MaxResults)
// KB config
require.NotNil(t, loaded.Search.KB)
assert.Equal(t, []string{"docs", "faq"}, loaded.Search.KB.Collections)
assert.Equal(t, 0.8, loaded.Search.KB.Threshold)
assert.True(t, loaded.Search.KB.Graph)
// DB config
require.NotNil(t, loaded.Search.DB)
assert.Equal(t, []string{"user", "product"}, loaded.Search.DB.Models)
assert.Equal(t, 50, loaded.Search.DB.MaxResults)
// Keyword config
require.NotNil(t, loaded.Search.Keyword)
assert.Equal(t, 8, loaded.Search.Keyword.MaxKeywords)
assert.Equal(t, "auto", loaded.Search.Keyword.Language)
// QueryDSL config
require.NotNil(t, loaded.Search.QueryDSL)
assert.True(t, loaded.Search.QueryDSL.Strict)
// Rerank config
require.NotNil(t, loaded.Search.Rerank)
assert.Equal(t, 5, loaded.Search.Rerank.TopN)
// Citation config
require.NotNil(t, loaded.Search.Citation)
assert.Equal(t, "#cite:{id}", loaded.Search.Citation.Format)
assert.False(t, loaded.Search.Citation.AutoInjectPrompt)
assert.Equal(t, "Please cite sources.", loaded.Search.Citation.CustomPrompt)
// Weights config
require.NotNil(t, loaded.Search.Weights)
assert.Equal(t, 1.0, loaded.Search.Weights.User)
assert.Equal(t, 0.85, loaded.Search.Weights.Hook)
assert.Equal(t, 0.65, loaded.Search.Weights.Auto)
// Options config
require.NotNil(t, loaded.Search.Options)
assert.Equal(t, 3, loaded.Search.Options.SkipThreshold)
}
// TestLoadStoreWithPartialSearchConfig tests loading assistant with partial search configuration
func TestLoadStoreWithPartialSearchConfig(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-partial-search"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Partial Search Config",
Type: "assistant",
Connector: "gpt-4o",
Search: &searchTypes.Config{
Web: &searchTypes.WebConfig{
Provider: "serper",
},
// Only web config, others are nil
},
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
// Verify Search config
require.NotNil(t, loaded.Search)
// Web config should be set
require.NotNil(t, loaded.Search.Web)
assert.Equal(t, "serper", loaded.Search.Web.Provider)
// Other configs should be nil
assert.Nil(t, loaded.Search.KB)
assert.Nil(t, loaded.Search.DB)
assert.Nil(t, loaded.Search.Keyword)
assert.Nil(t, loaded.Search.QueryDSL)
assert.Nil(t, loaded.Search.Rerank)
assert.Nil(t, loaded.Search.Citation)
assert.Nil(t, loaded.Search.Weights)
assert.Nil(t, loaded.Search.Options)
}
// TestLoadStoreWithoutSearchConfig tests loading assistant without search configuration
func TestLoadStoreWithoutSearchConfig(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-no-search"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test No Search Config",
Type: "assistant",
Connector: "gpt-4o",
// No Search config
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
// Search config should be nil
assert.Nil(t, loaded.Search)
}

301
agent/assistant/search.go Normal file
View file

@ -0,0 +1,301 @@
package assistant
import (
"strings"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search"
"github.com/yaoapp/yao/agent/search/nlp/keyword"
searchTypes "github.com/yaoapp/yao/agent/search/types"
)
// shouldAutoSearch determines if auto search should be executed
// Returns false if:
// - uses.search is "disabled"
// - assistant has no search configuration
func (ast *Assistant) shouldAutoSearch(ctx *context.Context, createResponse *context.HookCreateResponse) bool {
// Get merged uses configuration
uses := ast.getMergedSearchUses(createResponse)
// Check if search is explicitly disabled
if uses != nil && uses.Search == "disabled" {
ctx.Logger.Info("Auto search disabled by uses.search=disabled")
return false
}
// Check if assistant has search configuration
if ast.Search == nil && (uses == nil || uses.Search == "") {
return false
}
// Check if search is enabled (builtin, agent, mcp, or empty means builtin)
return true
}
// getMergedSearchUses returns the merged uses configuration for search
// Priority: createResponse > assistant
func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResponse) *context.Uses {
// Start with assistant uses
var uses *context.Uses
if ast.Uses != nil {
uses = &context.Uses{
Search: ast.Uses.Search,
Web: ast.Uses.Web,
Keyword: ast.Uses.Keyword,
QueryDSL: ast.Uses.QueryDSL,
Rerank: ast.Uses.Rerank,
}
}
// Override with createResponse.Uses if provided (highest priority)
if createResponse != nil && createResponse.Uses != nil {
if uses == nil {
uses = &context.Uses{}
}
if createResponse.Uses.Search != "" {
uses.Search = createResponse.Uses.Search
}
if createResponse.Uses.Web != "" {
uses.Web = createResponse.Uses.Web
}
if createResponse.Uses.Keyword != "" {
uses.Keyword = createResponse.Uses.Keyword
}
if createResponse.Uses.QueryDSL != "" {
uses.QueryDSL = createResponse.Uses.QueryDSL
}
if createResponse.Uses.Rerank != "" {
uses.Rerank = createResponse.Uses.Rerank
}
}
return uses
}
// executeAutoSearch executes auto search based on configuration
// Returns ReferenceContext with results and formatted context
// opts is optional, used to check Skip.Keyword
func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, opts ...*context.Options) *searchTypes.ReferenceContext {
ctx.Logger.Phase("Search")
defer ctx.Logger.PhaseComplete("Search")
// Get merged uses configuration
uses := ast.getMergedSearchUses(createResponse)
// Convert to search.Uses
searchUses := &search.Uses{}
if uses != nil {
searchUses.Search = uses.Search
searchUses.Web = uses.Web
searchUses.Keyword = uses.Keyword
searchUses.QueryDSL = uses.QueryDSL
searchUses.Rerank = uses.Rerank
}
// Get merged search config
searchConfig := ast.GetMergedSearchConfig()
// Create searcher
searcher := search.New(searchConfig, searchUses)
// Extract query from messages
query := extractQueryFromMessages(messages)
if query == "" {
ctx.Logger.Info("No query found in messages, skipping auto search")
return nil
}
// Check if keyword extraction should be skipped
skipKeyword := false
if len(opts) > 0 && opts[0] != nil && opts[0].Skip != nil {
skipKeyword = opts[0].Skip.Keyword
}
// Extract keywords for web search if:
// 1. uses.keyword is configured (not empty)
// 2. Skip.Keyword is not true
// 3. Web search is enabled
webSearchEnabled := searchConfig != nil && searchConfig.Web != nil
if webSearchEnabled && !skipKeyword && searchUses.Keyword != "" {
extractor := keyword.NewExtractor(searchUses.Keyword, searchConfig.Keyword)
keywords, err := extractor.Extract(ctx, query, nil)
if err != nil {
ctx.Logger.Warn("Keyword extraction failed, using original query: %v", err)
} else if len(keywords) > 0 {
// Use extracted keywords as the search query for web search
optimizedQuery := strings.Join(keywords, " ")
ctx.Logger.Info("Extracted keywords for web search: %s -> %s", truncateString(query, 30), optimizedQuery)
query = optimizedQuery
}
}
// Build search requests based on configuration
requests := ast.buildSearchRequests(query, searchConfig)
if len(requests) == 0 {
ctx.Logger.Info("No search requests to execute")
return nil
}
// Execute searches in parallel
ctx.Logger.Info("Executing %d search requests for query: %s", len(requests), truncateString(query, 50))
results, err := searcher.All(ctx, requests)
if err != nil {
// Log error but don't fail - search errors shouldn't block the main flow
ctx.Logger.Error("Auto search failed: %v", err)
return nil
}
// Build reference context (includes references, XML, and prompt)
var citationConfig *searchTypes.CitationConfig
if searchConfig != nil {
citationConfig = searchConfig.Citation
}
refCtx := search.BuildReferenceContext(results, citationConfig)
if len(refCtx.References) == 0 {
ctx.Logger.Info("No search results found")
return nil
}
ctx.Logger.Info("Auto search completed: %d references", len(refCtx.References))
return refCtx
}
// buildSearchRequests builds search requests based on assistant configuration
func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Config) []*searchTypes.Request {
var requests []*searchTypes.Request
// Web search - check if web search is configured
if config != nil && config.Web != nil {
requests = append(requests, &searchTypes.Request{
Type: searchTypes.SearchTypeWeb,
Query: query,
Source: searchTypes.SourceAuto,
Limit: config.Web.MaxResults,
})
}
// KB search - check if KB is configured
if ast.KB != nil && len(ast.KB.Collections) > 0 {
limit := 10
threshold := 0.7
if config != nil && config.KB != nil {
if config.KB.Threshold > 0 {
threshold = config.KB.Threshold
}
}
requests = append(requests, &searchTypes.Request{
Type: searchTypes.SearchTypeKB,
Query: query,
Source: searchTypes.SourceAuto,
Limit: limit,
Collections: ast.KB.Collections,
Threshold: threshold,
Graph: config != nil && config.KB != nil && config.KB.Graph,
})
}
// DB search - check if DB is configured
if ast.DB != nil && len(ast.DB.Models) > 0 {
limit := 20
if config != nil && config.DB != nil && config.DB.MaxResults > 0 {
limit = config.DB.MaxResults
}
requests = append(requests, &searchTypes.Request{
Type: searchTypes.SearchTypeDB,
Query: query,
Source: searchTypes.SourceAuto,
Limit: limit,
Models: ast.DB.Models,
})
}
return requests
}
// injectSearchContext injects search results into messages
// Adds search context as a system message after existing system messages
func (ast *Assistant) injectSearchContext(messages []context.Message, refCtx *searchTypes.ReferenceContext) []context.Message {
if refCtx == nil || len(refCtx.References) == 0 {
return messages
}
// Build the search context message
var contentParts []string
// Add citation prompt
if refCtx.Prompt != "" {
contentParts = append(contentParts, refCtx.Prompt)
}
// Add XML context
if refCtx.XML != "" {
contentParts = append(contentParts, refCtx.XML)
}
if len(contentParts) == 0 {
return messages
}
// Create system message with search context
searchMessage := context.Message{
Role: "system",
Content: strings.Join(contentParts, "\n\n"),
}
// Find the position to insert the search message
// Insert after any existing system messages but before user messages
insertIndex := 0
for i, msg := range messages {
if msg.Role == "system" {
insertIndex = i + 1
} else {
break
}
}
// Insert the search message
result := make([]context.Message, 0, len(messages)+1)
result = append(result, messages[:insertIndex]...)
result = append(result, searchMessage)
result = append(result, messages[insertIndex:]...)
return result
}
// extractQueryFromMessages extracts the search query from messages
// Uses the last user message as the query
func extractQueryFromMessages(messages []context.Message) string {
// Find the last user message
for i := len(messages) - 1; i >= 0; i-- {
if messages[i].Role == "user" {
content := messages[i].Content
// Handle string content
if str, ok := content.(string); ok {
return str
}
// Handle content parts (array of objects)
if parts, ok := content.([]interface{}); ok {
for _, part := range parts {
if partMap, ok := part.(map[string]interface{}); ok {
if partMap["type"] == "text" {
if text, ok := partMap["text"].(string); ok {
return text
}
}
}
}
}
}
}
return ""
}
// truncateString truncates a string to maxLen characters
func truncateString(s string, maxLen int) string {
if len(s) <= maxLen {
return s
}
return s[:maxLen] + "..."
}

84
agent/assistant/search_auto_disabled_test.go Normal file
View file

@ -0,0 +1,84 @@
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"
)
// newSearchAutoDisabledTestContext creates a test context
func newSearchAutoDisabledTestContext(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 TestSearchAutoDisabled(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/search-auto-disabled")
require.NoError(t, err)
require.NotNil(t, ast)
t.Run("ShouldHaveSearchConfig", func(t *testing.T) {
// Search config is set but uses.search is disabled
assert.NotNil(t, ast.Search, "search config should be set")
assert.NotNil(t, ast.Search.Web, "web search config should be set")
})
t.Run("ShouldHaveDisabledUses", func(t *testing.T) {
assert.NotNil(t, ast.Uses, "uses config should be set")
assert.Equal(t, "disabled", ast.Uses.Search, "uses.search should be disabled")
})
t.Run("StreamShouldNotExecuteSearch", func(t *testing.T) {
// Get agent via assistant.Get (required for Stream)
agent, err := assistant.Get("tests.search-auto-disabled")
require.NoError(t, err)
require.NotNil(t, agent)
// Create context
ctx := newSearchAutoDisabledTestContext("test-search-auto-disabled", "tests.search-auto-disabled")
// Create messages
messages := []context.Message{
{
Role: "user",
Content: "Hello, how are you?",
},
}
// Execute stream - should NOT trigger search because uses.search is "disabled"
response, err := agent.Stream(ctx, messages)
require.NoError(t, err)
require.NotNil(t, response)
resp := response.(*context.Response)
assert.NotNil(t, resp.Completion, "should have completion")
t.Logf("✓ Stream executed without search (disabled)")
})
}

125
agent/assistant/search_auto_full_test.go Normal file
View file

@ -0,0 +1,125 @@
package assistant_test
import (
stdContext "context"
"strings"
"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"
)
// newSearchAutoFullTestContext creates a test context
func newSearchAutoFullTestContext(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 TestSearchAutoFull(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/search-auto-full")
require.NoError(t, err)
require.NotNil(t, ast)
t.Run("ShouldHaveWebSearchConfig", func(t *testing.T) {
assert.NotNil(t, ast.Search, "search config should be set")
assert.NotNil(t, ast.Search.Web, "web search config should be set")
assert.Equal(t, "tavily", ast.Search.Web.Provider)
assert.Equal(t, 3, ast.Search.Web.MaxResults)
})
t.Run("ShouldHaveKBSearchConfig", func(t *testing.T) {
assert.NotNil(t, ast.Search.KB, "kb search config should be set")
assert.Equal(t, 0.7, ast.Search.KB.Threshold)
assert.False(t, ast.Search.KB.Graph)
})
t.Run("ShouldHaveDBSearchConfig", func(t *testing.T) {
assert.NotNil(t, ast.Search.DB, "db search config should be set")
assert.Equal(t, 10, ast.Search.DB.MaxResults)
})
t.Run("ShouldHaveKBCollections", func(t *testing.T) {
assert.NotNil(t, ast.KB, "kb config should be set")
assert.Contains(t, ast.KB.Collections, "test-collection")
})
t.Run("ShouldHaveDBModels", func(t *testing.T) {
assert.NotNil(t, ast.DB, "db config should be set")
assert.Contains(t, ast.DB.Models, "user")
assert.Contains(t, ast.DB.Models, "article")
})
t.Run("ShouldHaveCitationConfig", func(t *testing.T) {
assert.NotNil(t, ast.Search.Citation, "citation config should be set")
assert.Equal(t, "xml", ast.Search.Citation.Format)
assert.True(t, ast.Search.Citation.AutoInjectPrompt)
})
t.Run("ShouldHaveUsesConfig", func(t *testing.T) {
assert.NotNil(t, ast.Uses, "uses config should be set")
assert.Equal(t, "builtin", ast.Uses.Search)
assert.Equal(t, "builtin", ast.Uses.Web)
})
t.Run("StreamShouldExecuteMultipleSearchTypes", func(t *testing.T) {
// Get agent via assistant.Get (required for Stream)
agent, err := assistant.Get("tests.search-auto-full")
require.NoError(t, err)
require.NotNil(t, agent)
// Create context
ctx := newSearchAutoFullTestContext("test-search-auto-full", "tests.search-auto-full")
// Create messages with a search query
messages := []context.Message{
{
Role: "user",
Content: "Find information about machine learning",
},
}
// Execute stream - should trigger Web + KB + DB searches
response, err := agent.Stream(ctx, messages)
// Assert no error (if API key is configured)
if err != nil {
// If error contains "API key", it's expected in CI without keys
if strings.Contains(err.Error(), "API key") || strings.Contains(err.Error(), "api_key") {
t.Logf("Expected error without API key: %v", err)
return
}
// Other errors should fail
require.NoError(t, err)
}
require.NotNil(t, response)
resp := response.(*context.Response)
assert.NotNil(t, resp.Completion, "should have completion")
t.Logf("✓ Stream executed with full search config (Web + KB + DB)")
})
}

108
agent/assistant/search_auto_hook_disable_test.go Normal file
View file

@ -0,0 +1,108 @@
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)
resp := response.(*context.Response)
assert.NotNil(t, resp.Completion, "should have completion")
t.Logf("✓ Stream executed with hook disabling search")
})
}

186
agent/assistant/search_auto_keyword_test.go Normal file
View file

@ -0,0 +1,186 @@
package assistant_test
import (
stdContext "context"
"strings"
"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"
)
// newKeywordTestContext creates a test context for keyword extraction tests
func newKeywordTestContext(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 TestSearchAutoKeyword(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/search-auto-keyword")
require.NoError(t, err)
require.NotNil(t, ast)
t.Run("ShouldHaveKeywordConfig", func(t *testing.T) {
assert.NotNil(t, ast.Search, "search config should be set")
assert.NotNil(t, ast.Search.Keyword, "keyword config should be set")
assert.Equal(t, 5, ast.Search.Keyword.MaxKeywords)
assert.Equal(t, "auto", ast.Search.Keyword.Language)
})
t.Run("ShouldHaveKeywordInUses", func(t *testing.T) {
assert.NotNil(t, ast.Uses, "uses config should be set")
assert.Equal(t, "builtin", ast.Uses.Keyword)
})
t.Run("StreamWithKeywordExtraction", func(t *testing.T) {
// Get agent via assistant.Get (required for Stream)
agent, err := assistant.Get("tests.search-auto-keyword")
require.NoError(t, err)
require.NotNil(t, agent)
// Create context
ctx := newKeywordTestContext("test-search-keyword", "tests.search-auto-keyword")
// Create messages with a verbose query that should benefit from keyword extraction
messages := []context.Message{
{
Role: "user",
Content: "I want to find the best wireless headphones under 100 dollars for programming and music",
},
}
// Execute stream without Skip.Keyword (keyword extraction should happen)
response, err := agent.Stream(ctx, messages)
// Assert no error (if API key is configured)
if err != nil {
if strings.Contains(err.Error(), "API key") || strings.Contains(err.Error(), "api_key") {
t.Logf("Expected error without API key: %v", err)
return
}
require.NoError(t, err)
}
require.NotNil(t, response)
resp := response.(*context.Response)
assert.NotNil(t, resp.Completion, "should have completion")
t.Logf("✓ Stream with keyword extraction executed successfully")
})
t.Run("StreamWithSkipKeyword", func(t *testing.T) {
// Get agent via assistant.Get (required for Stream)
agent, err := assistant.Get("tests.search-auto-keyword")
require.NoError(t, err)
require.NotNil(t, agent)
// Create context
ctx := newKeywordTestContext("test-search-skip-keyword", "tests.search-auto-keyword")
// Create messages
messages := []context.Message{
{
Role: "user",
Content: "I want to find the best wireless headphones under 100 dollars",
},
}
// Execute stream with Skip.Keyword = true (keyword extraction should be skipped)
opts := &context.Options{
Skip: &context.Skip{
Keyword: true,
},
}
response, err := agent.Stream(ctx, messages, opts)
// Assert no error (if API key is configured)
if err != nil {
if strings.Contains(err.Error(), "API key") || strings.Contains(err.Error(), "api_key") {
t.Logf("Expected error without API key: %v", err)
return
}
require.NoError(t, err)
}
require.NotNil(t, response)
resp := response.(*context.Response)
assert.NotNil(t, resp.Completion, "should have completion")
t.Logf("✓ Stream with Skip.Keyword executed successfully")
})
}
func TestSearchAutoKeywordNotConfigured(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Use the search-auto-web assistant which does NOT have uses.keyword configured
ast, err := assistant.LoadPath("/assistants/tests/search-auto-web")
require.NoError(t, err)
require.NotNil(t, ast)
t.Run("ShouldNotHaveKeywordInUses", func(t *testing.T) {
// uses.keyword should be empty (not configured)
if ast.Uses != nil {
assert.Empty(t, ast.Uses.Keyword, "uses.keyword should be empty")
}
})
t.Run("StreamShouldSkipKeywordExtraction", func(t *testing.T) {
// Get agent via assistant.Get (required for Stream)
agent, err := assistant.Get("tests.search-auto-web")
require.NoError(t, err)
require.NotNil(t, agent)
// Create context
ctx := newKeywordTestContext("test-no-keyword", "tests.search-auto-web")
// Create messages
messages := []context.Message{
{
Role: "user",
Content: "What is the latest news about AI?",
},
}
// Execute stream - keyword extraction should NOT happen because uses.keyword is not set
response, err := agent.Stream(ctx, messages)
// Assert no error (if API key is configured)
if err != nil {
if strings.Contains(err.Error(), "API key") || strings.Contains(err.Error(), "api_key") {
t.Logf("Expected error without API key: %v", err)
return
}
require.NoError(t, err)
}
require.NotNil(t, response)
resp := response.(*context.Response)
assert.NotNil(t, resp.Completion, "should have completion")
t.Logf("✓ Stream without keyword config executed successfully")
})
}

103
agent/assistant/search_auto_web_test.go Normal file
View file

@ -0,0 +1,103 @@
package assistant_test
import (
stdContext "context"
"strings"
"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"
)
// newSearchAutoTestContext creates a test context for search auto tests
func newSearchAutoTestContext(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 TestSearchAutoWeb(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.LoadPath("/assistants/tests/search-auto-web")
require.NoError(t, err)
require.NotNil(t, ast)
t.Run("ShouldHaveSearchConfig", func(t *testing.T) {
assert.NotNil(t, ast.Search, "search config should be set")
assert.NotNil(t, ast.Search.Web, "web search config should be set")
assert.Equal(t, "tavily", ast.Search.Web.Provider)
assert.Equal(t, 3, ast.Search.Web.MaxResults)
})
t.Run("ShouldHaveUsesConfig", func(t *testing.T) {
assert.NotNil(t, ast.Uses, "uses config should be set")
assert.Equal(t, "builtin", ast.Uses.Search)
assert.Equal(t, "builtin", ast.Uses.Web)
})
t.Run("ShouldHaveCitationConfig", func(t *testing.T) {
assert.NotNil(t, ast.Search.Citation, "citation config should be set")
assert.Equal(t, "xml", ast.Search.Citation.Format)
assert.True(t, ast.Search.Citation.AutoInjectPrompt)
})
t.Run("StreamShouldExecuteAutoSearch", func(t *testing.T) {
// Get agent via assistant.Get (required for Stream)
agent, err := assistant.Get("tests.search-auto-web")
require.NoError(t, err)
require.NotNil(t, agent)
// Create context
ctx := newSearchAutoTestContext("test-search-auto-web", "tests.search-auto-web")
// Create messages with a search query
messages := []context.Message{
{
Role: "user",
Content: "What is the latest news about artificial intelligence?",
},
}
// Execute stream
response, err := agent.Stream(ctx, messages)
// Assert no error (if API key is configured)
if err != nil {
// If error contains "API key", it's expected in CI without keys
if strings.Contains(err.Error(), "API key") || strings.Contains(err.Error(), "api_key") {
t.Logf("Expected error without API key: %v", err)
return
}
// Other errors should fail
require.NoError(t, err)
}
require.NotNil(t, response)
resp := response.(*context.Response)
assert.NotNil(t, resp.Completion, "should have completion")
t.Logf("✓ Stream executed successfully with auto search")
})
}

13
agent/assistant/types.go
View file

@ -21,12 +21,6 @@ type API interface {
GetPlaceholder(locale string) *store.Placeholder
}
// SearchOption the search option
type SearchOption struct {
WebSearch *bool `json:"web_search,omitempty" yaml:"web_search,omitempty"` // Whether to search the web
Knowledge *bool `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Whether to search the knowledge
}
// Script the script scripts except hook script
type Script struct {
*v8.Script
@ -35,16 +29,13 @@ type Script struct {
// Assistant the assistant
type Assistant struct {
store.AssistantModel
Search *SearchOption `json:"search,omitempty" yaml:"search,omitempty"` // Whether this assistant supports search
HookScript *hook.Script `json:"-" yaml:"-"` // Hook Script (index.ts)
Scripts map[string]*Script `json:"-" yaml:"-"` // Other scripts
HookScript *hook.Script `json:"-" yaml:"-"` // Hook Script (index.ts)
Scripts map[string]*Script `json:"-" yaml:"-"` // Other scripts
// Internal
// ===============================
openai *api.OpenAI // OpenAI API
search bool // Whether this assistant supports search
vision bool // Whether this assistant supports vision
// toolCalls bool // Whether this assistant supports tool_calls
}
// MCPTool represents a simplified MCP tool for building LLM requests

17
agent/caller/caller.go Normal file
View file

@ -0,0 +1,17 @@
// Package caller provides a shared interface for calling agents
// This package is used by both content and search packages to avoid circular dependencies
package caller
import (
agentContext "github.com/yaoapp/yao/agent/context"
)
// AgentCaller interface for calling agents (to avoid circular dependency)
// Used by content handlers (vision, audio, etc.) and search handlers (agent mode)
type AgentCaller interface {
Stream(ctx *agentContext.Context, messages []agentContext.Message, options ...*agentContext.Options) (interface{}, error)
}
// AgentGetterFunc is a function type that gets an agent by ID
// This should be set by the assistant package during initialization
var AgentGetterFunc func(agentID string) (AgentCaller, error)

13
agent/content/tools.go
View file

@ -9,29 +9,22 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/mcp"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/caller"
agentContext "github.com/yaoapp/yao/agent/context"
)
// AgentCaller interface for calling agents (to avoid circular dependency)
type AgentCaller interface {
Stream(ctx *agentContext.Context, messages []agentContext.Message, options ...*agentContext.Options) (interface{}, error)
}
// AgentGetterFunc is a function type that gets an agent by ID
var AgentGetterFunc func(agentID string) (AgentCaller, error)
// fileInfoMutex protects concurrent access to files_info list in Space
var fileInfoMutex sync.Mutex
// CallAgent calls an agent to process content (vision, audio, etc.)
// This is a generic function that can be used by any handler
func CallAgent(ctx *agentContext.Context, agentID string, message agentContext.Message) (string, error) {
if AgentGetterFunc == nil {
if caller.AgentGetterFunc == nil {
return "", fmt.Errorf("AgentGetterFunc not initialized")
}
// Load the agent by ID using the injected function
agent, err := AgentGetterFunc(agentID)
agent, err := caller.AgentGetterFunc(agentID)
if err != nil {
return "", fmt.Errorf("failed to load agent %s: %w", agentID, err)
}

3
agent/context/jsapi.go
View file

@ -62,6 +62,9 @@ func (ctx *Context) NewObject(v8ctx *v8go.Context) (*v8go.Value, error) {
// Set mcp object
jsObject.Set("mcp", ctx.newMCPObject(v8ctx.Isolate()))
// Set search object
jsObject.Set("search", ctx.newSearchObject(v8ctx.Isolate()))
// Note: Space object will be set after instance creation (requires v8ctx)
// Create instance

362
agent/context/jsapi_search.go Normal file
View file

@ -0,0 +1,362 @@
package context
import (
"github.com/yaoapp/gou/runtime/v8/bridge"
"rogchap.com/v8go"
)
// SearchAPI defines the search JSAPI interface for ctx.search.*
// This interface is defined here to avoid circular dependency between context and search packages.
// The actual implementation is in agent/search/jsapi.go
type SearchAPI interface {
// Web executes web search
// Returns *types.Result or error information
Web(query string, opts map[string]interface{}) interface{}
// KB executes knowledge base search
// Returns *types.Result or error information
KB(query string, opts map[string]interface{}) interface{}
// DB executes database search
// Returns *types.Result or error information
DB(query string, opts map[string]interface{}) interface{}
// Parallel search methods - inspired by JavaScript Promise
// All waits for all searches to complete (like Promise.all)
All(requests []interface{}) []interface{}
// Any returns when any search succeeds with results (like Promise.any)
Any(requests []interface{}) []interface{}
// Race returns when any search completes (like Promise.race)
Race(requests []interface{}) []interface{}
}
// SearchAPIFactory is a function type that creates a SearchAPI for a context
// This is set by the search package during initialization
var SearchAPIFactory func(ctx *Context) SearchAPI
// Search returns the search API for this context
// Returns nil if SearchAPIFactory is not set
func (ctx *Context) Search() SearchAPI {
if SearchAPIFactory == nil {
return nil
}
return SearchAPIFactory(ctx)
}
// newSearchObject creates a new search object with all search methods
// This is called from jsapi.go NewObject() to mount ctx.search
func (ctx *Context) newSearchObject(iso *v8go.Isolate) *v8go.ObjectTemplate {
searchObj := v8go.NewObjectTemplate(iso)
// Single search methods
searchObj.Set("Web", ctx.searchWebMethod(iso))
searchObj.Set("KB", ctx.searchKBMethod(iso))
searchObj.Set("DB", ctx.searchDBMethod(iso))
// Parallel search methods - inspired by JavaScript Promise
searchObj.Set("All", ctx.searchAllMethod(iso))
searchObj.Set("Any", ctx.searchAnyMethod(iso))
searchObj.Set("Race", ctx.searchRaceMethod(iso))
return searchObj
}
// searchWebMethod implements ctx.search.Web(query, options?)
// Options:
// - limit: number - max results (default: 10)
// - sites: string[] - restrict to specific sites
// - time_range: string - "day", "week", "month", "year"
// - rerank: { top_n: number } - rerank options
func (ctx *Context) searchWebMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
v8ctx := info.Context()
args := info.Args()
// Validate arguments
if len(args) < 1 {
return bridge.JsException(v8ctx, "Web requires query parameter")
}
// Get query string
if !args[0].IsString() {
return bridge.JsException(v8ctx, "query must be a string")
}
query := args[0].String()
// Parse options (optional)
var opts map[string]interface{}
if len(args) >= 2 && !args[1].IsUndefined() && !args[1].IsNull() {
goVal, err := bridge.GoValue(args[1], v8ctx)
if err != nil {
return bridge.JsException(v8ctx, "invalid options: "+err.Error())
}
if optsMap, ok := goVal.(map[string]interface{}); ok {
opts = optsMap
}
}
// Get search API
searchAPI := ctx.Search()
if searchAPI == nil {
return bridge.JsException(v8ctx, "search API not available")
}
// Execute search
result := searchAPI.Web(query, opts)
// Convert result to JS value
jsVal, err := bridge.JsValue(v8ctx, result)
if err != nil {
return bridge.JsException(v8ctx, "failed to convert result: "+err.Error())
}
return jsVal
})
}
// searchKBMethod implements ctx.search.KB(query, options?)
// Options:
// - collections: string[] - collection IDs
// - threshold: number - similarity threshold (0-1)
// - limit: number - max results
// - graph: boolean - enable graph association
// - rerank: { top_n: number } - rerank options
func (ctx *Context) searchKBMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
v8ctx := info.Context()
args := info.Args()
// Validate arguments
if len(args) < 1 {
return bridge.JsException(v8ctx, "KB requires query parameter")
}
// Get query string
if !args[0].IsString() {
return bridge.JsException(v8ctx, "query must be a string")
}
query := args[0].String()
// Parse options (optional)
var opts map[string]interface{}
if len(args) >= 2 && !args[1].IsUndefined() && !args[1].IsNull() {
goVal, err := bridge.GoValue(args[1], v8ctx)
if err != nil {
return bridge.JsException(v8ctx, "invalid options: "+err.Error())
}
if optsMap, ok := goVal.(map[string]interface{}); ok {
opts = optsMap
}
}
// Get search API
searchAPI := ctx.Search()
if searchAPI == nil {
return bridge.JsException(v8ctx, "search API not available")
}
// Execute search
result := searchAPI.KB(query, opts)
// Convert result to JS value
jsVal, err := bridge.JsValue(v8ctx, result)
if err != nil {
return bridge.JsException(v8ctx, "failed to convert result: "+err.Error())
}
return jsVal
})
}
// searchDBMethod implements ctx.search.DB(query, options?)
// Options:
// - models: string[] - model IDs
// - wheres: Where[] - pre-defined filters (GOU QueryDSL Where format)
// - orders: Order[] - sort orders (GOU QueryDSL Order format)
// - select: string[] - fields to return
// - limit: number - max results
// - rerank: { top_n: number } - rerank options
func (ctx *Context) searchDBMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
v8ctx := info.Context()
args := info.Args()
// Validate arguments
if len(args) < 1 {
return bridge.JsException(v8ctx, "DB requires query parameter")
}
// Get query string
if !args[0].IsString() {
return bridge.JsException(v8ctx, "query must be a string")
}
query := args[0].String()
// Parse options (optional)
var opts map[string]interface{}
if len(args) >= 2 && !args[1].IsUndefined() && !args[1].IsNull() {
goVal, err := bridge.GoValue(args[1], v8ctx)
if err != nil {
return bridge.JsException(v8ctx, "invalid options: "+err.Error())
}
if optsMap, ok := goVal.(map[string]interface{}); ok {
opts = optsMap
}
}
// Get search API
searchAPI := ctx.Search()
if searchAPI == nil {
return bridge.JsException(v8ctx, "search API not available")
}
// Execute search
result := searchAPI.DB(query, opts)
// Convert result to JS value
jsVal, err := bridge.JsValue(v8ctx, result)
if err != nil {
return bridge.JsException(v8ctx, "failed to convert result: "+err.Error())
}
return jsVal
})
}
// searchAllMethod implements ctx.search.All(requests)
// Waits for all searches to complete (like Promise.all)
// Each request should have:
// - type: string - "web", "kb", or "db"
// - query: string - search query
// - ... other type-specific options
func (ctx *Context) searchAllMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
v8ctx := info.Context()
args := info.Args()
// Validate arguments
if len(args) < 1 {
return bridge.JsException(v8ctx, "All requires requests parameter")
}
// Parse requests array
goVal, err := bridge.GoValue(args[0], v8ctx)
if err != nil {
return bridge.JsException(v8ctx, "invalid requests: "+err.Error())
}
requestsArray, ok := goVal.([]interface{})
if !ok {
return bridge.JsException(v8ctx, "requests must be an array")
}
// Get search API
searchAPI := ctx.Search()
if searchAPI == nil {
return bridge.JsException(v8ctx, "search API not available")
}
// Execute parallel search
results := searchAPI.All(requestsArray)
// Convert results to JS value
jsVal, err := bridge.JsValue(v8ctx, results)
if err != nil {
return bridge.JsException(v8ctx, "failed to convert results: "+err.Error())
}
return jsVal
})
}
// searchAnyMethod implements ctx.search.Any(requests)
// Returns when any search succeeds with results (like Promise.any)
// Each request should have:
// - type: string - "web", "kb", or "db"
// - query: string - search query
// - ... other type-specific options
func (ctx *Context) searchAnyMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
v8ctx := info.Context()
args := info.Args()
// Validate arguments
if len(args) < 1 {
return bridge.JsException(v8ctx, "Any requires requests parameter")
}
// Parse requests array
goVal, err := bridge.GoValue(args[0], v8ctx)
if err != nil {
return bridge.JsException(v8ctx, "invalid requests: "+err.Error())
}
requestsArray, ok := goVal.([]interface{})
if !ok {
return bridge.JsException(v8ctx, "requests must be an array")
}
// Get search API
searchAPI := ctx.Search()
if searchAPI == nil {
return bridge.JsException(v8ctx, "search API not available")
}
// Execute parallel search
results := searchAPI.Any(requestsArray)
// Convert results to JS value
jsVal, err := bridge.JsValue(v8ctx, results)
if err != nil {
return bridge.JsException(v8ctx, "failed to convert results: "+err.Error())
}
return jsVal
})
}
// searchRaceMethod implements ctx.search.Race(requests)
// Returns when any search completes (like Promise.race)
// Each request should have:
// - type: string - "web", "kb", or "db"
// - query: string - search query
// - ... other type-specific options
func (ctx *Context) searchRaceMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
v8ctx := info.Context()
args := info.Args()
// Validate arguments
if len(args) < 1 {
return bridge.JsException(v8ctx, "Race requires requests parameter")
}
// Parse requests array
goVal, err := bridge.GoValue(args[0], v8ctx)
if err != nil {
return bridge.JsException(v8ctx, "invalid requests: "+err.Error())
}
requestsArray, ok := goVal.([]interface{})
if !ok {
return bridge.JsException(v8ctx, "requests must be an array")
}
// Get search API
searchAPI := ctx.Search()
if searchAPI == nil {
return bridge.JsException(v8ctx, "search API not available")
}
// Execute parallel search
results := searchAPI.Race(requestsArray)
// Convert results to JS value
jsVal, err := bridge.JsValue(v8ctx, results)
if err != nil {
return bridge.JsException(v8ctx, "failed to convert results: "+err.Error())
}
return jsVal
})
}

342
agent/context/jsapi_search_test.go Normal file
View file

@ -0,0 +1,342 @@
package context_test
import (
stdContext "context"
"encoding/json"
"strings"
"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/search/types"
"github.com/yaoapp/yao/agent/testutils"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
)
// Note: SearchAPIFactory is set by assistant.init() with proper config getter
// We import assistant package to ensure init() runs before tests
// newSearchTestContext creates a Context for search JSAPI testing
func newSearchTestContext(chatID, assistantID string) *context.Context {
authorized := &oauthTypes.AuthorizedInfo{
Subject: "test-user",
ClientID: "test-client-id",
Scope: "openid profile email",
SessionID: "test-session-id",
UserID: "test-user-123",
}
ctx := context.New(stdContext.Background(), authorized, chatID)
ctx.AssistantID = assistantID
ctx.Locale = "en-us"
ctx.Referer = context.RefererAPI
ctx.Accept = context.AcceptWebCUI
ctx.Metadata = make(map[string]interface{})
return ctx
}
// getResponseContent extracts the content from the first assistant message
func getResponseContent(res *context.HookCreateResponse) string {
if res == nil || len(res.Messages) == 0 {
return ""
}
for _, msg := range res.Messages {
if msg.Role == "assistant" {
if content, ok := msg.Content.(string); ok {
return content
}
}
}
return ""
}
// TestSearchJSAPI_Web tests ctx.search.Web() via Create Hook
func TestSearchJSAPI_Web(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the search-jsapi test assistant
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err, "Failed to get tests.search-jsapi assistant")
require.NotNil(t, agent.HookScript, "The tests.search-jsapi assistant has no script")
ctx := newSearchTestContext("chat-search-web", "tests.search-jsapi")
// Call Create hook with test:web command
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "test:web Yao App Engine"}})
require.NoError(t, err, "Create hook failed")
require.NotNil(t, res, "Expected non-nil response")
// Get response content from messages
content := getResponseContent(res)
require.NotEmpty(t, content, "Expected response content")
// Parse the JSON response
var result types.Result
err = json.Unmarshal([]byte(content), &result)
require.NoError(t, err, "Response should be valid JSON: %s", content)
// Verify result
assert.Equal(t, types.SearchTypeWeb, result.Type, "type should be web")
assert.Equal(t, "Yao App Engine", result.Query, "query should match")
assert.Empty(t, result.Error, "should not have error: %s", result.Error)
assert.Greater(t, len(result.Items), 0, "should have items")
t.Logf("Web search returned %d items", len(result.Items))
for i, item := range result.Items {
if i < 3 {
t.Logf(" [%s] %s - %s", item.CitationID, item.Title, item.URL)
}
}
}
// TestSearchJSAPI_WebWithSites tests ctx.search.Web() with site restriction
func TestSearchJSAPI_WebWithSites(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err)
require.NotNil(t, agent.HookScript)
ctx := newSearchTestContext("chat-search-web-sites", "tests.search-jsapi")
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "test:web_sites Yao App Engine"}})
require.NoError(t, err)
require.NotNil(t, res)
content := getResponseContent(res)
require.NotEmpty(t, content, "Expected response content")
var result types.Result
err = json.Unmarshal([]byte(content), &result)
require.NoError(t, err, "Response should be valid JSON: %s", content)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Empty(t, result.Error, "should not have error: %s", result.Error)
assert.Greater(t, len(result.Items), 0, "should have items")
// Verify all results are from allowed sites
allowedSites := []string{"github.com", "yaoapps.com"}
for _, item := range result.Items {
isAllowed := false
for _, site := range allowedSites {
if strings.Contains(item.URL, site) {
isAllowed = true
break
}
}
assert.True(t, isAllowed, "URL %s should be from allowed sites", item.URL)
}
t.Logf("Site-restricted search returned %d items", len(result.Items))
}
// TestSearchJSAPI_KB tests ctx.search.KB() via Create Hook (skeleton)
func TestSearchJSAPI_KB(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err)
require.NotNil(t, agent.HookScript)
ctx := newSearchTestContext("chat-search-kb", "tests.search-jsapi")
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "test:kb test query"}})
require.NoError(t, err)
require.NotNil(t, res)
content := getResponseContent(res)
require.NotEmpty(t, content, "Expected response content")
var result types.Result
err = json.Unmarshal([]byte(content), &result)
require.NoError(t, err, "Response should be valid JSON: %s", content)
assert.Equal(t, types.SearchTypeKB, result.Type, "type should be kb")
assert.Equal(t, "test query", result.Query, "query should match")
assert.Equal(t, types.SourceHook, result.Source, "source should be hook")
}
// TestSearchJSAPI_DB tests ctx.search.DB() via Create Hook (skeleton)
func TestSearchJSAPI_DB(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err)
require.NotNil(t, agent.HookScript)
ctx := newSearchTestContext("chat-search-db", "tests.search-jsapi")
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "test:db test query"}})
require.NoError(t, err)
require.NotNil(t, res)
content := getResponseContent(res)
require.NotEmpty(t, content, "Expected response content")
var result types.Result
err = json.Unmarshal([]byte(content), &result)
require.NoError(t, err, "Response should be valid JSON: %s", content)
assert.Equal(t, types.SearchTypeDB, result.Type, "type should be db")
assert.Equal(t, "test query", result.Query, "query should match")
assert.Equal(t, types.SourceHook, result.Source, "source should be hook")
}
// TestSearchJSAPI_All tests ctx.search.All() via Create Hook
func TestSearchJSAPI_All(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err)
require.NotNil(t, agent.HookScript)
ctx := newSearchTestContext("chat-search-all", "tests.search-jsapi")
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "test:all"}})
require.NoError(t, err)
require.NotNil(t, res)
content := getResponseContent(res)
require.NotEmpty(t, content, "Expected response content")
// Parse as array of results
var results []*types.Result
err = json.Unmarshal([]byte(content), &results)
require.NoError(t, err, "Response should be valid JSON array: %s", content)
assert.Len(t, results, 2, "should have 2 results")
// Both should succeed
successCount := 0
totalItems := 0
for _, r := range results {
if r != nil && r.Error == "" {
successCount++
totalItems += len(r.Items)
}
}
assert.Equal(t, 2, successCount, "both searches should succeed")
assert.Greater(t, totalItems, 0, "should have items")
t.Logf("All search: %d results, %d total items", len(results), totalItems)
}
// TestSearchJSAPI_Any tests ctx.search.Any() via Create Hook
func TestSearchJSAPI_Any(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err)
require.NotNil(t, agent.HookScript)
ctx := newSearchTestContext("chat-search-any", "tests.search-jsapi")
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "test:any"}})
require.NoError(t, err)
require.NotNil(t, res)
content := getResponseContent(res)
require.NotEmpty(t, content, "Expected response content")
var results []*types.Result
err = json.Unmarshal([]byte(content), &results)
require.NoError(t, err, "Response should be valid JSON array: %s", content)
assert.Len(t, results, 2, "should have 2 result slots")
// At least one should have results
hasSuccess := false
for _, r := range results {
if r != nil && len(r.Items) > 0 && r.Error == "" {
hasSuccess = true
break
}
}
assert.True(t, hasSuccess, "at least one search should succeed")
t.Logf("Any search completed")
}
// TestSearchJSAPI_Race tests ctx.search.Race() via Create Hook
func TestSearchJSAPI_Race(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err)
require.NotNil(t, agent.HookScript)
ctx := newSearchTestContext("chat-search-race", "tests.search-jsapi")
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "test:race"}})
require.NoError(t, err)
require.NotNil(t, res)
content := getResponseContent(res)
require.NotEmpty(t, content, "Expected response content")
var results []*types.Result
err = json.Unmarshal([]byte(content), &results)
require.NoError(t, err, "Response should be valid JSON array: %s", content)
assert.Len(t, results, 2, "should have 2 result slots")
// At least one should have completed
hasResult := false
for _, r := range results {
if r != nil {
hasResult = true
break
}
}
assert.True(t, hasResult, "at least one search should complete")
t.Logf("Race search completed")
}
// TestSearchJSAPI_InvalidCommand tests invalid test command
func TestSearchJSAPI_InvalidCommand(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err)
require.NotNil(t, agent.HookScript)
ctx := newSearchTestContext("chat-search-invalid", "tests.search-jsapi")
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "invalid command"}})
require.NoError(t, err)
require.NotNil(t, res)
content := getResponseContent(res)
assert.Contains(t, content, "Invalid test command", "should return error message")
}
// TestSearchJSAPI_UnknownMethod tests unknown test method
func TestSearchJSAPI_UnknownMethod(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.search-jsapi")
require.NoError(t, err)
require.NotNil(t, agent.HookScript)
ctx := newSearchTestContext("chat-search-unknown", "tests.search-jsapi")
res, _, err := agent.HookScript.Create(ctx, []context.Message{{Role: "user", Content: "test:unknown"}})
require.NoError(t, err)
require.NotNil(t, res)
content := getResponseContent(res)
assert.Contains(t, content, "Unknown test method", "should return error message")
}

1
agent/context/types.go
View file

@ -195,6 +195,7 @@ type Skip struct {
History bool `json:"history"` // Skip saving chat history (for internal calls like title/prompt generation)
Trace bool `json:"trace"` // Skip trace logging
Output bool `json:"output"` // Skip output to client (for internal A2A calls that only need response data)
Keyword bool `json:"keyword"` // Skip keyword extraction for web search (use raw query directly)
}
// MessageMetadata stores metadata for sent messages

8
agent/context/types_llm.go
View file

@ -10,8 +10,14 @@ import (
type Uses struct {
Vision string `json:"vision,omitempty"` // Vision processing tool. Format: "agent" or "mcp:server_id"
Audio string `json:"audio,omitempty"` // Audio processing tool. Format: "agent" or "mcp:server_id"
Search string `json:"search,omitempty"` // Search tool. Format: "agent" or "mcp:server_id"
Search string `json:"search,omitempty"` // Search tool. Format: "builtin", "disabled", "<assistant-id>", "mcp:<server>.<tool>"
Fetch string `json:"fetch,omitempty"` // Fetch/retrieval tool. Format: "agent" or "mcp:server_id"
// Search-related processing tools (NLP)
Web string `json:"web,omitempty"` // Web search handler: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
Keyword string `json:"keyword,omitempty"` // Keyword extraction: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
QueryDSL string `json:"querydsl,omitempty"` // QueryDSL generation: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
Rerank string `json:"rerank,omitempty"` // Result reranking: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
}
// VisionFormat specifies the vision input format

183
agent/load.go
View file

@ -10,6 +10,8 @@ import (
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
searchDefaults "github.com/yaoapp/yao/agent/search/defaults"
searchTypes "github.com/yaoapp/yao/agent/search/types"
storeMongo "github.com/yaoapp/yao/agent/store/mongo"
storeRedis "github.com/yaoapp/yao/agent/store/redis"
store "github.com/yaoapp/yao/agent/store/types"
@ -92,6 +94,12 @@ func Load(cfg config.Config) error {
return err
}
// Initialize Search Configuration
err = initSearchConfig()
if err != nil {
return err
}
// Initialize Assistant
err = initAssistant()
if err != nil {
@ -222,6 +230,10 @@ func initAssistant() error {
assistant.SetGlobalKBSetting(agentDSL.KB)
}
if agentDSL.Search != nil {
assistant.SetGlobalSearchConfig(agentDSL.Search)
}
// Load Built-in Assistants
err := assistant.LoadBuiltIn()
if err != nil {
@ -261,6 +273,177 @@ func initKBConfig() error {
return nil
}
// initSearchConfig initialize the search configuration from agent/search.yml
func initSearchConfig() error {
// Start with system defaults
agentDSL.Search = searchDefaults.SystemDefaults
path := filepath.Join("agent", "search.yml")
if exists, _ := application.App.Exists(path); !exists {
return nil // Search config is optional, use defaults
}
// Read the search configuration
bytes, err := application.App.Read(path)
if err != nil {
return err
}
var searchConfig searchTypes.Config
err = application.Parse("search.yml", bytes, &searchConfig)
if err != nil {
return err
}
// Merge with defaults
agentDSL.Search = mergeSearchConfig(searchDefaults.SystemDefaults, &searchConfig)
return nil
}
// mergeSearchConfig merges two search configs (base < override)
func mergeSearchConfig(base, override *searchTypes.Config) *searchTypes.Config {
if base == nil {
return override
}
if override == nil {
return base
}
result := *base // Copy base
// Merge Web config
if override.Web != nil {
if result.Web == nil {
result.Web = override.Web
} else {
if override.Web.Provider != "" {
result.Web.Provider = override.Web.Provider
}
if override.Web.APIKeyEnv != "" {
result.Web.APIKeyEnv = override.Web.APIKeyEnv
}
if override.Web.MaxResults > 0 {
result.Web.MaxResults = override.Web.MaxResults
}
}
}
// Merge KB config
if override.KB != nil {
if result.KB == nil {
result.KB = override.KB
} else {
if len(override.KB.Collections) > 0 {
result.KB.Collections = override.KB.Collections
}
if override.KB.Threshold > 0 {
result.KB.Threshold = override.KB.Threshold
}
if override.KB.Graph {
result.KB.Graph = override.KB.Graph
}
}
}
// Merge DB config
if override.DB != nil {
if result.DB == nil {
result.DB = override.DB
} else {
if len(override.DB.Models) > 0 {
result.DB.Models = override.DB.Models
}
if override.DB.MaxResults > 0 {
result.DB.MaxResults = override.DB.MaxResults
}
}
}
// Merge Keyword config
if override.Keyword != nil {
if result.Keyword == nil {
result.Keyword = override.Keyword
} else {
if override.Keyword.MaxKeywords > 0 {
result.Keyword.MaxKeywords = override.Keyword.MaxKeywords
}
if override.Keyword.Language != "" {
result.Keyword.Language = override.Keyword.Language
}
}
}
// Merge QueryDSL config
if override.QueryDSL != nil {
result.QueryDSL = override.QueryDSL
}
// Merge Rerank config
if override.Rerank != nil {
if result.Rerank == nil {
result.Rerank = override.Rerank
} else {
if override.Rerank.TopN > 0 {
result.Rerank.TopN = override.Rerank.TopN
}
}
}
// Merge Citation config
if override.Citation != nil {
if result.Citation == nil {
result.Citation = override.Citation
} else {
if override.Citation.Format != "" {
result.Citation.Format = override.Citation.Format
}
// AutoInjectPrompt is a bool, need to check if explicitly set
result.Citation.AutoInjectPrompt = override.Citation.AutoInjectPrompt
if override.Citation.CustomPrompt != "" {
result.Citation.CustomPrompt = override.Citation.CustomPrompt
}
}
}
// Merge Weights config
if override.Weights != nil {
if result.Weights == nil {
result.Weights = override.Weights
} else {
if override.Weights.User > 0 {
result.Weights.User = override.Weights.User
}
if override.Weights.Hook > 0 {
result.Weights.Hook = override.Weights.Hook
}
if override.Weights.Auto > 0 {
result.Weights.Auto = override.Weights.Auto
}
}
}
// Merge Options config
if override.Options != nil {
if result.Options == nil {
result.Options = override.Options
} else {
if override.Options.SkipThreshold > 0 {
result.Options.SkipThreshold = override.Options.SkipThreshold
}
}
}
return &result
}
// GetSearchConfig returns the global search configuration
func GetSearchConfig() *searchTypes.Config {
if agentDSL == nil {
return searchDefaults.SystemDefaults
}
return agentDSL.Search
}
// defaultAssistant get the default assistant
func defaultAssistant() (*assistant.Assistant, error) {
if agentDSL.Uses == nil || agentDSL.Uses.Default == "" {

43
agent/load_test.go
View file

@ -94,6 +94,49 @@ func TestLoad(t *testing.T) {
assert.Equal(t, "__yao.utf8", agent.KB.Chat.DocumentDefaults.Converter.ProviderID)
assert.Equal(t, "standard-text", agent.KB.Chat.DocumentDefaults.Converter.OptionID)
})
t.Run("LoadSearchConfig", func(t *testing.T) {
// Search configuration should be loaded from agent/search.yml
assert.NotNil(t, agent.Search)
// Verify web config
assert.NotNil(t, agent.Search.Web)
assert.Equal(t, "tavily", agent.Search.Web.Provider)
assert.Equal(t, 10, agent.Search.Web.MaxResults)
// Verify KB config
assert.NotNil(t, agent.Search.KB)
assert.Equal(t, 0.7, agent.Search.KB.Threshold)
assert.False(t, agent.Search.KB.Graph)
// Verify DB config
assert.NotNil(t, agent.Search.DB)
assert.Equal(t, 20, agent.Search.DB.MaxResults)
// Verify keyword config
assert.NotNil(t, agent.Search.Keyword)
assert.Equal(t, 10, agent.Search.Keyword.MaxKeywords)
assert.Equal(t, "auto", agent.Search.Keyword.Language)
// Verify rerank config
assert.NotNil(t, agent.Search.Rerank)
assert.Equal(t, 10, agent.Search.Rerank.TopN)
// Verify citation config
assert.NotNil(t, agent.Search.Citation)
assert.Equal(t, "#ref:{id}", agent.Search.Citation.Format)
assert.True(t, agent.Search.Citation.AutoInjectPrompt)
// Verify weights config
assert.NotNil(t, agent.Search.Weights)
assert.Equal(t, 1.0, agent.Search.Weights.User)
assert.Equal(t, 0.8, agent.Search.Weights.Hook)
assert.Equal(t, 0.6, agent.Search.Weights.Auto)
// Verify options config
assert.NotNil(t, agent.Search.Options)
assert.Equal(t, 5, agent.Search.Options.SkipThreshold)
})
}
func TestGetGlobalPrompts(t *testing.T) {

927
agent/search/DESIGN.md
View file

File diff suppressed because it is too large Load diff

27
agent/search/citation.go Normal file
View file

@ -0,0 +1,27 @@
package search
import (
"fmt"
"sync/atomic"
)
// CitationGenerator generates unique citation IDs
type CitationGenerator struct {
counter uint64
}
// NewCitationGenerator creates a new citation generator
func NewCitationGenerator() *CitationGenerator {
return &CitationGenerator{}
}
// Next generates the next citation ID
func (g *CitationGenerator) Next() string {
n := atomic.AddUint64(&g.counter, 1)
return fmt.Sprintf("ref_%03d", n)
}
// Reset resets the counter (for testing)
func (g *CitationGenerator) Reset() {
atomic.StoreUint64(&g.counter, 0)
}

88
agent/search/citation_test.go Normal file
View file

@ -0,0 +1,88 @@
package search
import (
"sync"
"testing"
"github.com/stretchr/testify/assert"
)
func TestCitationGenerator_Next(t *testing.T) {
gen := NewCitationGenerator()
// First ID should be ref_001
id1 := gen.Next()
assert.Equal(t, "ref_001", id1)
// Second ID should be ref_002
id2 := gen.Next()
assert.Equal(t, "ref_002", id2)
// Third ID should be ref_003
id3 := gen.Next()
assert.Equal(t, "ref_003", id3)
}
func TestCitationGenerator_Reset(t *testing.T) {
gen := NewCitationGenerator()
// Generate some IDs
gen.Next()
gen.Next()
gen.Next()
// Reset
gen.Reset()
// Next ID should be ref_001 again
id := gen.Next()
assert.Equal(t, "ref_001", id)
}
func TestCitationGenerator_Format(t *testing.T) {
gen := NewCitationGenerator()
// Generate 999 IDs to test padding
for i := 0; i < 999; i++ {
gen.Next()
}
// 1000th ID should be ref_1000 (no padding limit)
id := gen.Next()
assert.Equal(t, "ref_1000", id)
}
func TestCitationGenerator_Concurrent(t *testing.T) {
gen := NewCitationGenerator()
// Run 100 goroutines, each generating 10 IDs
var wg sync.WaitGroup
ids := make(chan string, 1000)
for i := 0; i < 100; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 10; j++ {
ids <- gen.Next()
}
}()
}
wg.Wait()
close(ids)
// Collect all IDs
idSet := make(map[string]bool)
for id := range ids {
idSet[id] = true
}
// All 1000 IDs should be unique
assert.Equal(t, 1000, len(idSet))
}
func TestNewCitationGenerator(t *testing.T) {
gen := NewCitationGenerator()
assert.NotNil(t, gen)
}

63
agent/search/defaults/defaults.go Normal file
View file

@ -0,0 +1,63 @@
package defaults
import "github.com/yaoapp/yao/agent/search/types"
// SystemDefaults provides hardcoded default values
// Used by agent/load.go for merging with agent/search.yao
var SystemDefaults = &types.Config{
// Web search defaults
Web: &types.WebConfig{
Provider: "tavily",
MaxResults: 10,
},
// KB search defaults
KB: &types.KBConfig{
Threshold: 0.7,
Graph: false,
},
// DB search defaults
DB: &types.DBConfig{
MaxResults: 20,
},
// Keyword extraction options (uses.keyword)
Keyword: &types.KeywordConfig{
MaxKeywords: 10,
Language: "auto",
},
// QueryDSL generation options (uses.querydsl)
QueryDSL: &types.QueryDSLConfig{
Strict: false,
},
// Rerank options (uses.rerank)
Rerank: &types.RerankConfig{
TopN: 10,
},
// Citation
Citation: &types.CitationConfig{
Format: "#ref:{id}",
AutoInjectPrompt: true,
},
// Source weights
Weights: &types.WeightsConfig{
User: 1.0,
Hook: 0.8,
Auto: 0.6,
},
// Behavior options
Options: &types.OptionsConfig{
SkipThreshold: 5,
},
}
// GetWeight returns the weight for a source type using default config
func GetWeight(source types.SourceType) float64 {
return SystemDefaults.GetWeight(source)
}

93
agent/search/handlers/db/handler.go Normal file
View file

@ -0,0 +1,93 @@
package db
import (
"time"
"github.com/yaoapp/yao/agent/search/types"
)
// Handler implements DB search
type Handler struct {
usesQueryDSL string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
config *types.DBConfig // DB search configuration
}
// NewHandler creates a new DB search handler
func NewHandler(usesQueryDSL string, cfg *types.DBConfig) *Handler {
return &Handler{usesQueryDSL: usesQueryDSL, config: cfg}
}
// Type returns the search type this handler supports
func (h *Handler) Type() types.SearchType {
return types.SearchTypeDB
}
// Search converts NL to QueryDSL and executes
// TODO: Implement actual QueryDSL generation and model query logic
func (h *Handler) Search(req *types.Request) (*types.Result, error) {
start := time.Now()
// Validate request
if req.Query == "" {
return &types.Result{
Type: types.SearchTypeDB,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(start).Milliseconds(),
Error: "query is required",
}, nil
}
// Get models from request or config
models := req.Models
if len(models) == 0 && h.config != nil {
models = h.config.Models
}
// If no models specified, return empty result
if len(models) == 0 {
return &types.Result{
Type: types.SearchTypeDB,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(start).Milliseconds(),
}, nil
}
// Get max results
maxResults := req.Limit
if maxResults == 0 && h.config != nil && h.config.MaxResults > 0 {
maxResults = h.config.MaxResults
}
if maxResults == 0 {
maxResults = 20 // default
}
// TODO: Implement actual DB search
// 1. Get model schemas for specified models
// 2. Generate QueryDSL from natural language query using uses.querydsl mode:
// - "builtin": template-based generation
// - "<assistant-id>": delegate to LLM assistant
// - "mcp:<server>.<tool>": call external MCP tool
// 3. Execute QueryDSL on each model
// 4. Format results and return
// For now, return empty result (skeleton)
result := &types.Result{
Type: types.SearchTypeDB,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(start).Milliseconds(),
}
// Store maxResults for later use
_ = maxResults
return result, nil
}

215
agent/search/handlers/db/handler_test.go Normal file
View file

@ -0,0 +1,215 @@
package db
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/types"
)
func TestNewHandler(t *testing.T) {
t.Run("with nil config", func(t *testing.T) {
h := NewHandler("builtin", nil)
assert.NotNil(t, h)
assert.Equal(t, "builtin", h.usesQueryDSL)
assert.Nil(t, h.config)
})
t.Run("with config", func(t *testing.T) {
cfg := &types.DBConfig{
Models: []string{"product", "order"},
MaxResults: 50,
}
h := NewHandler("workers.nlp.querydsl", cfg)
assert.NotNil(t, h)
assert.Equal(t, "workers.nlp.querydsl", h.usesQueryDSL)
assert.Equal(t, cfg, h.config)
})
t.Run("with mcp mode", func(t *testing.T) {
h := NewHandler("mcp:nlp.generate_querydsl", nil)
assert.NotNil(t, h)
assert.Equal(t, "mcp:nlp.generate_querydsl", h.usesQueryDSL)
})
}
func TestHandler_Type(t *testing.T) {
h := NewHandler("builtin", nil)
assert.Equal(t, types.SearchTypeDB, h.Type())
}
func TestHandler_Search(t *testing.T) {
tests := []struct {
name string
usesQueryDSL string
config *types.DBConfig
req *types.Request
expectError string
expectItems int
}{
{
name: "empty query",
usesQueryDSL: "builtin",
config: nil,
req: &types.Request{
Type: types.SearchTypeDB,
Query: "",
},
expectError: "query is required",
expectItems: 0,
},
{
name: "no models in request or config",
usesQueryDSL: "builtin",
config: nil,
req: &types.Request{
Type: types.SearchTypeDB,
Query: "find products under $100",
},
expectError: "",
expectItems: 0,
},
{
name: "models from config",
usesQueryDSL: "builtin",
config: &types.DBConfig{
Models: []string{"product"},
MaxResults: 20,
},
req: &types.Request{
Type: types.SearchTypeDB,
Query: "find products under $100",
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "models from request",
usesQueryDSL: "builtin",
config: nil,
req: &types.Request{
Type: types.SearchTypeDB,
Query: "find products under $100",
Models: []string{"product", "order"},
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "with limit",
usesQueryDSL: "builtin",
config: &types.DBConfig{
Models: []string{"product"},
},
req: &types.Request{
Type: types.SearchTypeDB,
Query: "find products",
Models: []string{"product"},
Limit: 5,
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "with wheres",
usesQueryDSL: "builtin",
config: &types.DBConfig{
Models: []string{"product"},
},
req: &types.Request{
Type: types.SearchTypeDB,
Query: "find products",
Models: []string{"product"},
// Wheres would be set here in real usage
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "agent mode",
usesQueryDSL: "workers.nlp.querydsl",
config: &types.DBConfig{
Models: []string{"product"},
},
req: &types.Request{
Type: types.SearchTypeDB,
Query: "find products",
Models: []string{"product"},
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "mcp mode",
usesQueryDSL: "mcp:nlp.generate_querydsl",
config: &types.DBConfig{
Models: []string{"product"},
},
req: &types.Request{
Type: types.SearchTypeDB,
Query: "find products",
Models: []string{"product"},
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
h := NewHandler(tt.usesQueryDSL, tt.config)
result, err := h.Search(tt.req)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.Equal(t, types.SearchTypeDB, result.Type)
assert.Equal(t, tt.req.Query, result.Query)
assert.Equal(t, tt.expectItems, len(result.Items))
if tt.expectError != "" {
assert.Equal(t, tt.expectError, result.Error)
} else {
assert.Empty(t, result.Error)
}
// Duration should be set
assert.GreaterOrEqual(t, result.Duration, int64(0))
})
}
}
func TestHandler_Search_SourcePreserved(t *testing.T) {
h := NewHandler("builtin", &types.DBConfig{Models: []string{"product"}})
sources := []types.SourceType{types.SourceUser, types.SourceHook, types.SourceAuto}
for _, source := range sources {
req := &types.Request{
Type: types.SearchTypeDB,
Query: "test",
Source: source,
Models: []string{"product"},
}
result, err := h.Search(req)
assert.NoError(t, err)
assert.Equal(t, source, result.Source)
}
}
func TestHandler_Search_MaxResultsFromConfig(t *testing.T) {
cfg := &types.DBConfig{
Models: []string{"product"},
MaxResults: 50,
}
h := NewHandler("builtin", cfg)
req := &types.Request{
Type: types.SearchTypeDB,
Query: "test",
Models: []string{"product"},
// No limit in request, should use config's MaxResults
}
result, err := h.Search(req)
assert.NoError(t, err)
assert.NotNil(t, result)
// Skeleton doesn't actually use maxResults yet, but the test ensures the handler runs
}

95
agent/search/handlers/kb/handler.go Normal file
View file

@ -0,0 +1,95 @@
package kb
import (
"time"
"github.com/yaoapp/yao/agent/search/types"
)
// Handler implements KB search
type Handler struct {
config *types.KBConfig // KB search configuration
}
// NewHandler creates a new KB search handler
func NewHandler(cfg *types.KBConfig) *Handler {
return &Handler{config: cfg}
}
// Type returns the search type this handler supports
func (h *Handler) Type() types.SearchType {
return types.SearchTypeKB
}
// Search executes vector search and optional graph association
// TODO: Implement actual vector search and graph association logic
func (h *Handler) Search(req *types.Request) (*types.Result, error) {
start := time.Now()
// Validate request
if req.Query == "" {
return &types.Result{
Type: types.SearchTypeKB,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(start).Milliseconds(),
Error: "query is required",
}, nil
}
// Get collections from request or config
collections := req.Collections
if len(collections) == 0 && h.config != nil {
collections = h.config.Collections
}
// If no collections specified, return empty result
if len(collections) == 0 {
return &types.Result{
Type: types.SearchTypeKB,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(start).Milliseconds(),
}, nil
}
// Get threshold from request or config
threshold := req.Threshold
if threshold == 0 && h.config != nil && h.config.Threshold > 0 {
threshold = h.config.Threshold
}
if threshold == 0 {
threshold = 0.7 // default
}
// Get limit
limit := req.Limit
if limit == 0 {
limit = 10 // default
}
// TODO: Implement actual vector search
// 1. Generate embedding for query using collection's embedding config
// 2. Search each collection with vector similarity
// 3. If req.Graph is true, perform graph association
// 4. Merge and return results
// For now, return empty result (skeleton)
result := &types.Result{
Type: types.SearchTypeKB,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(start).Milliseconds(),
}
// Store threshold in result metadata for debugging
_ = threshold
return result, nil
}

170
agent/search/handlers/kb/handler_test.go Normal file
View file

@ -0,0 +1,170 @@
package kb
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/types"
)
func TestNewHandler(t *testing.T) {
t.Run("with nil config", func(t *testing.T) {
h := NewHandler(nil)
assert.NotNil(t, h)
assert.Nil(t, h.config)
})
t.Run("with config", func(t *testing.T) {
cfg := &types.KBConfig{
Collections: []string{"docs", "faq"},
Threshold: 0.8,
Graph: true,
}
h := NewHandler(cfg)
assert.NotNil(t, h)
assert.Equal(t, cfg, h.config)
})
}
func TestHandler_Type(t *testing.T) {
h := NewHandler(nil)
assert.Equal(t, types.SearchTypeKB, h.Type())
}
func TestHandler_Search(t *testing.T) {
tests := []struct {
name string
config *types.KBConfig
req *types.Request
expectError string
expectItems int
}{
{
name: "empty query",
config: nil,
req: &types.Request{
Type: types.SearchTypeKB,
Query: "",
},
expectError: "query is required",
expectItems: 0,
},
{
name: "no collections in request or config",
config: nil,
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
},
expectError: "",
expectItems: 0,
},
{
name: "collections from config",
config: &types.KBConfig{
Collections: []string{"docs"},
Threshold: 0.7,
},
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "collections from request",
config: nil,
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs", "faq"},
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "with threshold from request",
config: &types.KBConfig{
Collections: []string{"docs"},
Threshold: 0.7,
},
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Threshold: 0.9,
Collections: []string{"docs"},
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "with graph enabled",
config: &types.KBConfig{
Collections: []string{"docs"},
Graph: true,
},
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs"},
Graph: true,
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "with limit",
config: &types.KBConfig{
Collections: []string{"docs"},
},
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs"},
Limit: 5,
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
h := NewHandler(tt.config)
result, err := h.Search(tt.req)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.Equal(t, types.SearchTypeKB, result.Type)
assert.Equal(t, tt.req.Query, result.Query)
assert.Equal(t, tt.expectItems, len(result.Items))
if tt.expectError != "" {
assert.Equal(t, tt.expectError, result.Error)
} else {
assert.Empty(t, result.Error)
}
// Duration should be set
assert.GreaterOrEqual(t, result.Duration, int64(0))
})
}
}
func TestHandler_Search_SourcePreserved(t *testing.T) {
h := NewHandler(&types.KBConfig{Collections: []string{"docs"}})
sources := []types.SourceType{types.SourceUser, types.SourceHook, types.SourceAuto}
for _, source := range sources {
req := &types.Request{
Type: types.SearchTypeKB,
Query: "test",
Source: source,
Collections: []string{"docs"},
}
result, err := h.Search(req)
assert.NoError(t, err)
assert.Equal(t, source, result.Source)
}
}

232
agent/search/handlers/web/agent.go Normal file
View file

@ -0,0 +1,232 @@
package web
import (
"encoding/json"
"fmt"
"time"
"github.com/yaoapp/yao/agent/caller"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// AgentProvider implements web search using another agent (AI Search)
type AgentProvider struct {
agentID string // Agent/Assistant ID (e.g., "workers.search.web")
}
// NewAgentProvider creates a new Agent provider
func NewAgentProvider(agentID string) *AgentProvider {
return &AgentProvider{
agentID: agentID,
}
}
// Search executes web search via agent delegation
// The agent can understand intent, generate optimized queries, and return structured results
func (p *AgentProvider) Search(ctx *agentContext.Context, req *types.Request) (*types.Result, error) {
startTime := time.Now()
// Check if context is provided
if ctx == nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: "Agent mode requires context",
}, nil
}
// Check if AgentGetterFunc is initialized
if caller.AgentGetterFunc == nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: "AgentGetterFunc not initialized",
}, nil
}
// Get the agent
agent, err := caller.AgentGetterFunc(p.agentID)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: fmt.Sprintf("Agent '%s' not found: %v", p.agentID, err),
}, nil
}
// Build message for the agent
// Include search parameters in the message content
searchParams := map[string]interface{}{
"query": req.Query,
"type": "web",
"source": string(req.Source),
}
if req.Limit > 0 {
searchParams["limit"] = req.Limit
}
if len(req.Sites) > 0 {
searchParams["sites"] = req.Sites
}
if req.TimeRange != "" {
searchParams["time_range"] = req.TimeRange
}
// Convert to JSON for the message
paramsJSON, err := json.Marshal(searchParams)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: fmt.Sprintf("Failed to serialize search params: %v", err),
}, nil
}
// Create message for the agent
message := agentContext.Message{
Role: "user",
Content: string(paramsJSON),
}
// Call the agent with skip options (no history, no output)
opts := &agentContext.Options{
Skip: &agentContext.Skip{
History: true,
Output: true,
},
}
response, err := agent.Stream(ctx, []agentContext.Message{message}, opts)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: fmt.Sprintf("Agent call failed: %v", err),
}, nil
}
// Parse the agent response
items, total, parseErr := p.parseAgentResponse(response, req.Source)
if parseErr != "" {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: parseErr,
}, nil
}
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: items,
Total: total,
Duration: time.Since(startTime).Milliseconds(),
}, nil
}
// parseAgentResponse parses the agent response into search result items
// The agent should return a JSON structure with search results
func (p *AgentProvider) parseAgentResponse(response interface{}, source types.SourceType) ([]*types.ResultItem, int, string) {
if response == nil {
return nil, 0, "Agent returned nil response"
}
// Try to extract data from response
var data map[string]interface{}
// Handle different response types
switch v := response.(type) {
case map[string]interface{}:
data = v
case string:
// Try to parse as JSON
if err := json.Unmarshal([]byte(v), &data); err != nil {
return nil, 0, fmt.Sprintf("Failed to parse agent response as JSON: %v", err)
}
default:
// Try to marshal and unmarshal
jsonBytes, err := json.Marshal(response)
if err != nil {
return nil, 0, fmt.Sprintf("Failed to serialize agent response: %v", err)
}
if err := json.Unmarshal(jsonBytes, &data); err != nil {
return nil, 0, fmt.Sprintf("Failed to parse agent response: %v", err)
}
}
// Check for "next" field (custom hook data)
if next, hasNext := data["next"]; hasNext && next != nil {
if nextMap, ok := next.(map[string]interface{}); ok {
data = nextMap
} else if nextStr, ok := next.(string); ok {
// Try to parse as JSON
if err := json.Unmarshal([]byte(nextStr), &data); err != nil {
return nil, 0, fmt.Sprintf("Failed to parse next hook data: %v", err)
}
}
}
// Extract items from data
items := []*types.ResultItem{}
total := 0
if itemsData, ok := data["items"].([]interface{}); ok {
for _, itemData := range itemsData {
if item, ok := itemData.(map[string]interface{}); ok {
resultItem := &types.ResultItem{
Type: types.SearchTypeWeb,
Source: source,
}
if title, ok := item["title"].(string); ok {
resultItem.Title = title
}
if content, ok := item["content"].(string); ok {
resultItem.Content = content
}
if url, ok := item["url"].(string); ok {
resultItem.URL = url
}
if score, ok := item["score"].(float64); ok {
resultItem.Score = score
}
items = append(items, resultItem)
}
}
}
if totalVal, ok := data["total"].(float64); ok {
total = int(totalVal)
} else {
total = len(items)
}
return items, total, ""
}

284
agent/search/handlers/web/agent_test.go Normal file
View file

@ -0,0 +1,284 @@
package web_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/handlers/web"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
)
// TestAgentProviderWithAssistantConfig tests AgentProvider using web-agent-caller assistant config
func TestAgentProviderWithAssistantConfig(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-agent-caller test assistant to get its config
ast, err := assistant.LoadPath("/assistants/tests/web-agent-caller")
require.NoError(t, err)
require.NotNil(t, ast)
require.NotNil(t, ast.Uses)
// Verify assistant config
assert.Equal(t, "tests.web-agent-caller", ast.ID)
assert.Equal(t, "tests.web-agent", ast.Uses.Web)
// Create AgentProvider from uses.web
provider := web.NewAgentProvider(ast.Uses.Web)
require.NotNil(t, provider)
// Create a mock context
ctx := createTestContext(t)
// Execute search
req := &types.Request{
Query: "Yao App Engine",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 5,
}
result, err := provider.Search(ctx, req)
require.NoError(t, err)
require.NotNil(t, result)
// Verify result structure
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "Yao App Engine", result.Query)
assert.Equal(t, types.SourceAuto, result.Source)
// Agent should return mock results from Next hook
if result.Error == "" {
assert.Greater(t, result.Total, 0)
assert.NotEmpty(t, result.Items)
assert.Greater(t, result.Duration, int64(0))
// Verify result item structure
for _, item := range result.Items {
assert.Equal(t, types.SearchTypeWeb, item.Type)
assert.Equal(t, types.SourceAuto, item.Source)
assert.NotEmpty(t, item.Title)
assert.NotEmpty(t, item.URL)
}
t.Logf("Agent search returned %d results in %dms", result.Total, result.Duration)
} else {
t.Logf("Agent search returned error: %s", result.Error)
}
}
// TestAgentProviderWithSiteRestriction tests AgentProvider with domain restriction
func TestAgentProviderWithSiteRestriction(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create AgentProvider
provider := web.NewAgentProvider("tests.web-agent")
// Create a mock context
ctx := createTestContext(t)
// Execute search with site restriction
req := &types.Request{
Query: "documentation",
Type: types.SearchTypeWeb,
Source: types.SourceHook,
Sites: []string{"github.com"},
Limit: 3,
}
result, err := provider.Search(ctx, req)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, types.SourceHook, result.Source)
if result.Error == "" {
// All results should be from github.com (mock data respects sites)
for _, item := range result.Items {
assert.Contains(t, item.URL, "github.com", "Result URL should be from github.com")
}
t.Logf("Site-restricted agent search returned %d results", result.Total)
} else {
t.Logf("Agent search returned error: %s", result.Error)
}
}
// TestAgentProviderWithTimeRange tests AgentProvider with time range filter
func TestAgentProviderWithTimeRange(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create AgentProvider
provider := web.NewAgentProvider("tests.web-agent")
// Create a mock context
ctx := createTestContext(t)
// Execute search with time range
req := &types.Request{
Query: "artificial intelligence news",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
TimeRange: "week",
Limit: 5,
}
result, err := provider.Search(ctx, req)
require.NoError(t, err)
require.NotNil(t, result)
if result.Error == "" {
t.Logf("Time-ranged agent search (last week) returned %d results in %dms", result.Total, result.Duration)
} else {
t.Logf("Agent search returned error: %s", result.Error)
}
}
// TestAgentProviderNotFound tests AgentProvider when agent is not found
func TestAgentProviderNotFound(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create AgentProvider with non-existent agent
provider := web.NewAgentProvider("nonexistent.agent")
// Create a mock context
ctx := createTestContext(t)
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
}
result, err := provider.Search(ctx, req)
// Should not return error, but result should have error message
require.NoError(t, err)
require.NotNil(t, result)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "not found")
}
// TestAgentProviderWithoutContext tests AgentProvider without context
func TestAgentProviderWithoutContext(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create AgentProvider
provider := web.NewAgentProvider("tests.web-agent")
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
}
// Call without context (nil)
result, err := provider.Search(nil, req)
// Should still work - agent provider handles nil context
require.NoError(t, err)
require.NotNil(t, result)
// May have error if context is required for agent call
t.Logf("Agent search without context: error=%s, total=%d", result.Error, result.Total)
}
// TestWebHandlerAgentMode tests the web handler in agent mode
func TestWebHandlerAgentMode(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create handler with agent mode
handler := web.NewHandler("tests.web-agent", nil)
require.NotNil(t, handler)
// Verify type
assert.Equal(t, types.SearchTypeWeb, handler.Type())
// Create a mock context
ctx := createTestContext(t)
// Execute search with context
req := &types.Request{
Query: "Yao framework",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 5,
}
result, err := handler.SearchWithContext(ctx, req)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "Yao framework", result.Query)
if result.Error == "" {
t.Logf("Handler agent mode returned %d results", result.Total)
} else {
t.Logf("Handler agent mode returned error: %s", result.Error)
}
}
// TestWebHandlerAgentModeWithoutContext tests the web handler in agent mode without context
func TestWebHandlerAgentModeWithoutContext(t *testing.T) {
// Create handler with agent mode
handler := web.NewHandler("tests.web-agent", nil)
require.NotNil(t, handler)
req := &types.Request{
Query: "test",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
}
// Call Search() without context (uses SearchWithContext with nil)
result, err := handler.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "requires context")
}
// createTestContext creates a test context for agent calls
func createTestContext(t *testing.T) *agentContext.Context {
authorized := &oauthTypes.AuthorizedInfo{
UserID: "test-user",
TenantID: "test-tenant",
}
ctx := agentContext.New(nil, authorized, "test-chat-id")
ctx.AssistantID = "tests.web-agent-caller"
return ctx
}

110
agent/search/handlers/web/handler.go Normal file
View file

@ -0,0 +1,110 @@
package web
import (
"fmt"
"strings"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// Handler implements web search
type Handler struct {
usesWeb string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
config *types.WebConfig // Web search configuration
}
// NewHandler creates a new web search handler
func NewHandler(usesWeb string, cfg *types.WebConfig) *Handler {
return &Handler{usesWeb: usesWeb, config: cfg}
}
// Type returns the search type this handler supports
func (h *Handler) Type() types.SearchType {
return types.SearchTypeWeb
}
// Search executes web search based on uses.web mode
// ctx is optional and only required for agent mode
func (h *Handler) Search(req *types.Request) (*types.Result, error) {
return h.SearchWithContext(nil, req)
}
// SearchWithContext executes web search with optional agent context
// ctx is required for agent mode, optional for builtin and MCP modes
func (h *Handler) SearchWithContext(ctx *agentContext.Context, req *types.Request) (*types.Result, error) {
switch {
case h.usesWeb == "builtin" || h.usesWeb == "":
return h.builtinSearch(req)
case strings.HasPrefix(h.usesWeb, "mcp:"):
return h.mcpSearch(req)
default:
// Agent mode: delegate to assistant for AI-powered search
if ctx == nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Error: "Agent mode requires context",
}, nil
}
return h.agentSearch(ctx, req)
}
}
// builtinSearch uses Tavily/Serper/SerpAPI directly
func (h *Handler) builtinSearch(req *types.Request) (*types.Result, error) {
// Determine provider from config
providerName := "tavily" // default
if h.config != nil && h.config.Provider != "" {
providerName = h.config.Provider
}
switch providerName {
case "tavily":
return NewTavilyProvider(h.config).Search(req)
case "serper":
// Serper (serper.dev) - POST request with X-API-KEY header
return NewSerperProvider(h.config).Search(req)
case "serpapi":
// SerpAPI (serpapi.com) - GET request with api_key parameter
return NewSerpAPIProvider(h.config).Search(req)
default:
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Error: fmt.Sprintf("Unknown provider: %s (supported: tavily, serper, serpapi)", providerName),
}, nil
}
}
// agentSearch delegates to an assistant for AI-powered search
func (h *Handler) agentSearch(ctx *agentContext.Context, req *types.Request) (*types.Result, error) {
provider := NewAgentProvider(h.usesWeb)
return provider.Search(ctx, req)
}
// mcpSearch calls external MCP tool
func (h *Handler) mcpSearch(req *types.Request) (*types.Result, error) {
// Parse "mcp:server.tool"
mcpRef := strings.TrimPrefix(h.usesWeb, "mcp:")
provider, err := NewMCPProvider(mcpRef)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Error: fmt.Sprintf("Invalid MCP format: %v", err),
}, nil
}
return provider.Search(req)
}

203
agent/search/handlers/web/mcp.go Normal file
View file

@ -0,0 +1,203 @@
package web
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
"github.com/yaoapp/gou/mcp"
gouMCPTypes "github.com/yaoapp/gou/mcp/types"
"github.com/yaoapp/yao/agent/search/types"
)
// MCPProvider implements web search using MCP tool
type MCPProvider struct {
serverID string // MCP server ID (e.g., "search")
toolName string // MCP tool name (e.g., "web_search")
}
// NewMCPProvider creates a new MCP provider from "mcp:server.tool" format
func NewMCPProvider(mcpRef string) (*MCPProvider, error) {
// Parse "server.tool" format
parts := strings.SplitN(mcpRef, ".", 2)
if len(parts) != 2 {
return nil, fmt.Errorf("invalid MCP format, expected 'server.tool', got '%s'", mcpRef)
}
return &MCPProvider{
serverID: parts[0],
toolName: parts[1],
}, nil
}
// Search executes web search via MCP tool
func (p *MCPProvider) Search(req *types.Request) (*types.Result, error) {
startTime := time.Now()
// Select MCP client
client, err := mcp.Select(p.serverID)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: fmt.Sprintf("MCP client '%s' not found: %v", p.serverID, err),
}, nil
}
// Build MCP tool arguments
args := map[string]interface{}{
"query": req.Query,
}
if req.Limit > 0 {
args["limit"] = req.Limit
}
if len(req.Sites) > 0 {
args["sites"] = req.Sites
}
if req.TimeRange != "" {
args["time_range"] = req.TimeRange
}
// Call MCP tool
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
result, err := client.CallTool(ctx, p.toolName, args)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: fmt.Sprintf("MCP tool call failed: %v", err),
}, nil
}
// Parse MCP result
items, total, parseErr := p.parseResult(result, req.Source)
if parseErr != "" {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: parseErr,
}, nil
}
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: items,
Total: total,
Duration: time.Since(startTime).Milliseconds(),
}, nil
}
// parseResult parses MCP tool result into search result items
func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse, source types.SourceType) ([]*types.ResultItem, int, string) {
if result == nil {
return nil, 0, "MCP returned nil result"
}
// Check for errors in result
if result.IsError {
errMsg := "MCP tool returned error"
if len(result.Content) > 0 && result.Content[0].Text != "" {
errMsg = result.Content[0].Text
}
return nil, 0, errMsg
}
// Parse content - expect JSON data
if len(result.Content) == 0 {
return []*types.ResultItem{}, 0, ""
}
// Try to extract data from content
var data map[string]interface{}
for _, content := range result.Content {
// Check text content type
if content.Type == gouMCPTypes.ToolContentTypeText && content.Text != "" {
// Try to parse as JSON
if parsed, ok := parseJSON(content.Text); ok {
data = parsed
break
}
}
}
if data == nil {
return []*types.ResultItem{}, 0, ""
}
// Extract items from data
items := []*types.ResultItem{}
total := 0
if itemsData, ok := data["items"].([]interface{}); ok {
for _, itemData := range itemsData {
if item, ok := itemData.(map[string]interface{}); ok {
resultItem := &types.ResultItem{
Type: types.SearchTypeWeb,
Source: source,
}
if title, ok := item["title"].(string); ok {
resultItem.Title = title
}
if content, ok := item["content"].(string); ok {
resultItem.Content = content
}
if url, ok := item["url"].(string); ok {
resultItem.URL = url
}
if score, ok := item["score"].(float64); ok {
resultItem.Score = score
}
items = append(items, resultItem)
}
}
}
if totalVal, ok := data["total"].(float64); ok {
total = int(totalVal)
} else {
total = len(items)
}
return items, total, ""
}
// parseJSON attempts to parse a string as JSON
func parseJSON(s string) (map[string]interface{}, bool) {
// Simple JSON detection - if it starts with { and ends with }
s = strings.TrimSpace(s)
if !strings.HasPrefix(s, "{") || !strings.HasSuffix(s, "}") {
return nil, false
}
// Use encoding/json for parsing
var result map[string]interface{}
if err := json.Unmarshal([]byte(s), &result); err != nil {
return nil, false
}
return result, true
}

241
agent/search/handlers/web/mcp_test.go Normal file
View file

@ -0,0 +1,241 @@
package web_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/search/handlers/web"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
)
// TestMCPProviderWithAssistantConfig tests MCPProvider using web-mcp assistant config
func TestMCPProviderWithAssistantConfig(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-mcp test assistant to get its config
ast, err := assistant.LoadPath("/assistants/tests/web-mcp")
require.NoError(t, err)
require.NotNil(t, ast)
require.NotNil(t, ast.Uses)
// Verify assistant config
assert.Equal(t, "tests.web-mcp", ast.ID)
assert.Equal(t, "mcp:search.web_search", ast.Uses.Web)
// Create MCPProvider from uses.web
mcpRef := ast.Uses.Web[4:] // Remove "mcp:" prefix
provider, err := web.NewMCPProvider(mcpRef)
require.NoError(t, err)
require.NotNil(t, provider)
// Execute search
req := &types.Request{
Query: "Yao App Engine",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// Verify result structure
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "Yao App Engine", result.Query)
assert.Equal(t, types.SourceAuto, result.Source)
// MCP should return mock results
if result.Error == "" {
assert.Greater(t, result.Total, 0)
assert.NotEmpty(t, result.Items)
assert.Greater(t, result.Duration, int64(0))
// Verify result item structure
for _, item := range result.Items {
assert.Equal(t, types.SearchTypeWeb, item.Type)
assert.Equal(t, types.SourceAuto, item.Source)
assert.NotEmpty(t, item.Title)
assert.NotEmpty(t, item.URL)
}
t.Logf("MCP search returned %d results in %dms", result.Total, result.Duration)
} else {
t.Logf("MCP search returned error (expected if MCP not loaded): %s", result.Error)
}
}
// TestMCPProviderWithSiteRestriction tests MCPProvider with domain restriction
func TestMCPProviderWithSiteRestriction(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create MCPProvider
provider, err := web.NewMCPProvider("search.web_search")
require.NoError(t, err)
// Execute search with site restriction
req := &types.Request{
Query: "documentation",
Type: types.SearchTypeWeb,
Source: types.SourceHook,
Sites: []string{"github.com"},
Limit: 3,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, types.SourceHook, result.Source)
if result.Error == "" {
t.Logf("Site-restricted MCP search returned %d results", result.Total)
} else {
t.Logf("MCP search returned error (expected if MCP not loaded): %s", result.Error)
}
}
// TestMCPProviderWithTimeRange tests MCPProvider with time range filter
func TestMCPProviderWithTimeRange(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create MCPProvider
provider, err := web.NewMCPProvider("search.web_search")
require.NoError(t, err)
// Execute search with time range
req := &types.Request{
Query: "artificial intelligence news",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
TimeRange: "week",
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
if result.Error == "" {
t.Logf("Time-ranged MCP search (last week) returned %d results in %dms", result.Total, result.Duration)
} else {
t.Logf("MCP search returned error (expected if MCP not loaded): %s", result.Error)
}
}
// TestMCPProviderInvalidFormat tests MCPProvider with invalid format
func TestMCPProviderInvalidFormat(t *testing.T) {
// Test invalid format without dot
_, err := web.NewMCPProvider("invalid")
require.Error(t, err)
assert.Contains(t, err.Error(), "invalid MCP format")
// Test empty string
_, err = web.NewMCPProvider("")
require.Error(t, err)
assert.Contains(t, err.Error(), "invalid MCP format")
}
// TestMCPProviderNotFound tests MCPProvider when MCP server is not found
func TestMCPProviderNotFound(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create MCPProvider with non-existent server
provider, err := web.NewMCPProvider("nonexistent.web_search")
require.NoError(t, err)
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
}
result, err := provider.Search(req)
// Should not return error, but result should have error message
require.NoError(t, err)
require.NotNil(t, result)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "not found")
}
// TestWebHandlerMCPMode tests the web handler in MCP mode
func TestWebHandlerMCPMode(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Create handler with MCP mode
handler := web.NewHandler("mcp:search.web_search", nil)
require.NotNil(t, handler)
// Verify type
assert.Equal(t, types.SearchTypeWeb, handler.Type())
// Execute search
req := &types.Request{
Query: "Yao framework",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 5,
}
result, err := handler.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "Yao framework", result.Query)
if result.Error == "" {
t.Logf("Handler MCP mode returned %d results", result.Total)
} else {
t.Logf("Handler MCP mode returned error (expected if MCP not loaded): %s", result.Error)
}
}
// TestWebHandlerInvalidMCPFormat tests the web handler with invalid MCP format
func TestWebHandlerInvalidMCPFormat(t *testing.T) {
// Create handler with invalid MCP format
handler := web.NewHandler("mcp:invalid", nil)
require.NotNil(t, handler)
req := &types.Request{
Query: "test",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
}
result, err := handler.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "Invalid MCP format")
}

302
agent/search/handlers/web/serpapi.go Normal file
View file

@ -0,0 +1,302 @@
package web
import (
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"os"
"time"
"github.com/yaoapp/yao/agent/search/types"
)
const (
serpAPIURL = "https://serpapi.com/search.json"
serpAPITimeout = 30 * time.Second
)
// SerpAPIProvider implements web search using SerpAPI (supports multiple search engines)
type SerpAPIProvider struct {
apiKey string
maxResults int
engine string // Search engine: "google", "bing", "baidu", "yandex", "duckduckgo", etc.
}
// NewSerpAPIProvider creates a new SerpAPI provider
func NewSerpAPIProvider(cfg *types.WebConfig) *SerpAPIProvider {
apiKey := ""
if cfg != nil && cfg.APIKeyEnv != "" {
// Support both "$ENV.VAR_NAME" and "VAR_NAME" formats
envName := cfg.APIKeyEnv
if len(envName) > 5 && envName[:5] == "$ENV." {
envName = envName[5:]
}
apiKey = os.Getenv(envName)
}
maxResults := 10
if cfg != nil && cfg.MaxResults > 0 {
maxResults = cfg.MaxResults
}
engine := "google" // Default to Google
if cfg != nil && cfg.Engine != "" {
engine = cfg.Engine
}
return &SerpAPIProvider{
apiKey: apiKey,
maxResults: maxResults,
engine: engine,
}
}
// serpAPIResponse represents the response from SerpAPI
type serpAPIResponse struct {
SearchMetadata serpAPIMetadata `json:"search_metadata"`
SearchParameters serpAPIParams `json:"search_parameters"`
SearchInformation serpAPIInfo `json:"search_information"`
OrganicResults []serpAPIResult `json:"organic_results"`
AnswerBox *serpAPIAnswerBox `json:"answer_box,omitempty"`
KnowledgeGraph *serpAPIKnowledge `json:"knowledge_graph,omitempty"`
RelatedSearches []serpAPIRelated `json:"related_searches,omitempty"`
RelatedQuestions []serpAPIQuestion `json:"related_questions,omitempty"`
}
// serpAPIMetadata contains metadata from response
type serpAPIMetadata struct {
ID string `json:"id"`
Status string `json:"status"`
CreatedAt string `json:"created_at"`
ProcessedAt string `json:"processed_at"`
TotalTimeTaken float64 `json:"total_time_taken"`
}
// serpAPIParams contains search parameters from response
type serpAPIParams struct {
Engine string `json:"engine"`
Q string `json:"q"`
Location string `json:"location_used"`
GoogleDomain string `json:"google_domain"`
HL string `json:"hl"`
GL string `json:"gl"`
Device string `json:"device"`
}
// serpAPIInfo contains search information
type serpAPIInfo struct {
QueryDisplayed string `json:"query_displayed"`
TotalResults int64 `json:"total_results"`
TimeTakenDisplayed float64 `json:"time_taken_displayed"`
OrganicResultsState string `json:"organic_results_state"`
}
// serpAPIResult represents a single organic search result
type serpAPIResult struct {
Position int `json:"position"`
Title string `json:"title"`
Link string `json:"link"`
RedirectLink string `json:"redirect_link,omitempty"`
DisplayedLink string `json:"displayed_link"`
Snippet string `json:"snippet"`
Date string `json:"date,omitempty"`
CachedPageLink string `json:"cached_page_link,omitempty"`
}
// serpAPIAnswerBox represents the answer box (featured snippet)
type serpAPIAnswerBox struct {
Type string `json:"type,omitempty"`
Title string `json:"title,omitempty"`
Snippet string `json:"snippet,omitempty"`
Link string `json:"link,omitempty"`
}
// serpAPIKnowledge represents knowledge graph data
type serpAPIKnowledge struct {
Title string `json:"title,omitempty"`
Type string `json:"type,omitempty"`
Description string `json:"description,omitempty"`
}
// serpAPIRelated represents related searches
type serpAPIRelated struct {
Query string `json:"query"`
Link string `json:"link"`
}
// serpAPIQuestion represents related questions (People Also Ask)
type serpAPIQuestion struct {
Question string `json:"question"`
Snippet string `json:"snippet,omitempty"`
Title string `json:"title,omitempty"`
Link string `json:"link,omitempty"`
}
// Search executes a web search using SerpAPI
func (p *SerpAPIProvider) Search(req *types.Request) (*types.Result, error) {
startTime := time.Now()
// Validate API key
if p.apiKey == "" {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Error: "SerpAPI API key not configured",
}, nil
}
// Determine max results
maxResults := p.maxResults
if req.Limit > 0 {
maxResults = req.Limit
}
// Build query parameters
params := url.Values{}
params.Set("engine", p.engine)
params.Set("api_key", p.apiKey)
params.Set("num", fmt.Sprintf("%d", maxResults))
// Build search query with site restrictions if specified
query := req.Query
if len(req.Sites) > 0 {
siteQuery := ""
for i, site := range req.Sites {
if i > 0 {
siteQuery += " OR "
}
siteQuery += "site:" + site
}
query = "(" + siteQuery + ") " + req.Query
}
params.Set("q", query)
// Add time range if specified (tbs parameter)
if req.TimeRange != "" {
tbs := convertSerpAPITimeRange(req.TimeRange)
if tbs != "" {
params.Set("tbs", tbs)
}
}
// Execute API call
serpResp, err := p.callAPI(params)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: fmt.Sprintf("SerpAPI error: %v", err),
}, nil
}
// Convert results
items := make([]*types.ResultItem, 0, len(serpResp.OrganicResults))
// Add answer box as first result if available
if serpResp.AnswerBox != nil && serpResp.AnswerBox.Snippet != "" {
items = append(items, &types.ResultItem{
Type: types.SearchTypeWeb,
Title: serpResp.AnswerBox.Title,
Content: serpResp.AnswerBox.Snippet,
URL: serpResp.AnswerBox.Link,
Score: 1.0, // Featured snippet gets highest score
Source: req.Source,
Metadata: map[string]interface{}{
"type": "answer_box",
},
})
}
// Add organic results
for _, r := range serpResp.OrganicResults {
// Calculate score based on position (1st = 0.95, 2nd = 0.90, etc.)
score := 1.0 - float64(r.Position)*0.05
if score < 0.1 {
score = 0.1
}
items = append(items, &types.ResultItem{
Type: types.SearchTypeWeb,
Title: r.Title,
Content: r.Snippet,
URL: r.Link,
Score: score,
Source: req.Source,
})
}
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: items,
Total: len(items),
Duration: time.Since(startTime).Milliseconds(),
}, nil
}
// callAPI makes the HTTP GET request to SerpAPI
func (p *SerpAPIProvider) callAPI(params url.Values) (*serpAPIResponse, error) {
// Build URL with query parameters
reqURL := serpAPIURL + "?" + params.Encode()
// Create HTTP request
httpReq, err := http.NewRequest(http.MethodGet, reqURL, nil)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
// Execute request
client := &http.Client{Timeout: serpAPITimeout}
resp, err := client.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("request failed: %w", err)
}
defer resp.Body.Close()
// Read response body
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
// Check status code
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(respBody))
}
// Parse response
var serpResp serpAPIResponse
if err := json.Unmarshal(respBody, &serpResp); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
return &serpResp, nil
}
// convertSerpAPITimeRange converts time range to SerpAPI tbs format
func convertSerpAPITimeRange(timeRange string) string {
switch timeRange {
case "hour":
return "qdr:h"
case "day":
return "qdr:d"
case "week":
return "qdr:w"
case "month":
return "qdr:m"
case "year":
return "qdr:y"
default:
return ""
}
}

394
agent/search/handlers/web/serpapi_test.go Normal file
View file

@ -0,0 +1,394 @@
package web_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/search/handlers/web"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
)
// TestSerpAPIProviderWithAssistantConfig tests SerpAPIProvider using web-serpapi assistant config
func TestSerpAPIProviderWithAssistantConfig(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serpapi test assistant to get its config
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
require.NoError(t, err)
require.NotNil(t, ast)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Verify assistant config
assert.Equal(t, "tests.web-serpapi", ast.ID)
assert.Equal(t, "serpapi", ast.Search.Web.Provider)
assert.Equal(t, "$ENV.SERPAPI_API_KEY", ast.Search.Web.APIKeyEnv)
assert.Equal(t, 10, ast.Search.Web.MaxResults)
// Create SerpAPIProvider with assistant's web config
provider := web.NewSerpAPIProvider(ast.Search.Web)
require.NotNil(t, provider)
// Execute search
req := &types.Request{
Query: "Yao App Engine",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// Verify result structure
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "Yao App Engine", result.Query)
assert.Equal(t, types.SourceAuto, result.Source)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
// Verify we got results
assert.Greater(t, result.Total, 0)
assert.NotEmpty(t, result.Items)
assert.Greater(t, result.Duration, int64(0))
// Verify result item structure
for _, item := range result.Items {
assert.Equal(t, types.SearchTypeWeb, item.Type)
assert.Equal(t, types.SourceAuto, item.Source)
assert.NotEmpty(t, item.Title)
assert.NotEmpty(t, item.URL)
assert.Greater(t, item.Score, 0.0)
}
t.Logf("Search returned %d results in %dms", result.Total, result.Duration)
}
// TestSerpAPIProviderWithSiteRestriction tests SerpAPIProvider with domain restriction
func TestSerpAPIProviderWithSiteRestriction(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serpapi test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create SerpAPIProvider
provider := web.NewSerpAPIProvider(ast.Search.Web)
// Execute search with site restriction
req := &types.Request{
Query: "documentation",
Type: types.SearchTypeWeb,
Source: types.SourceHook,
Sites: []string{"github.com"},
Limit: 3,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, types.SourceHook, result.Source)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
require.NotEmpty(t, result.Items, "Search should return results")
// All results should be from github.com
for _, item := range result.Items {
assert.Contains(t, item.URL, "github.com", "Result URL should be from github.com")
}
t.Logf("Site-restricted search returned %d results from github.com", result.Total)
}
// TestSerpAPIProviderWithMultipleSites tests SerpAPIProvider with multiple domain restrictions
func TestSerpAPIProviderWithMultipleSites(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serpapi test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create SerpAPIProvider
provider := web.NewSerpAPIProvider(ast.Search.Web)
// Execute search with multiple site restrictions
req := &types.Request{
Query: "golang tutorial",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Sites: []string{"github.com", "golang.org"},
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
require.NotEmpty(t, result.Items, "Search should return results")
// Results should be from either github.com or golang.org
for _, item := range result.Items {
isValidSite := false
for _, site := range req.Sites {
if containsSite(item.URL, site) {
isValidSite = true
break
}
}
assert.True(t, isValidSite, "Result URL should be from github.com or golang.org: %s", item.URL)
}
t.Logf("Multi-site search returned %d results", result.Total)
}
// TestSerpAPIProviderWithTimeRange tests SerpAPIProvider with time range filter
func TestSerpAPIProviderWithTimeRange(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serpapi test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create SerpAPIProvider
provider := web.NewSerpAPIProvider(ast.Search.Web)
// Execute search with time range
req := &types.Request{
Query: "artificial intelligence news",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
TimeRange: "week", // Last week
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
t.Logf("Time-ranged search (last week) returned %d results in %dms", result.Total, result.Duration)
}
// TestSerpAPIProviderWithoutAPIKey tests graceful degradation when API key is missing
func TestSerpAPIProviderWithoutAPIKey(t *testing.T) {
// Create provider with nil config (no API key)
provider := web.NewSerpAPIProvider(nil)
require.NotNil(t, provider)
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
}
result, err := provider.Search(req)
// Should not return error, but result should have error message
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "test query", result.Query)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "API key")
assert.Empty(t, result.Items)
assert.Equal(t, 0, result.Total)
}
// TestSerpAPIProviderWithEmptyConfig tests provider with empty config
func TestSerpAPIProviderWithEmptyConfig(t *testing.T) {
// Create provider with empty config
cfg := &types.WebConfig{}
provider := web.NewSerpAPIProvider(cfg)
require.NotNil(t, provider)
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceUser,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "API key")
}
// TestSerpAPIProviderMaxResults tests that max_results from config is respected
func TestSerpAPIProviderMaxResults(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serpapi test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create SerpAPIProvider
provider := web.NewSerpAPIProvider(ast.Search.Web)
// Execute search without limit (should use config's max_results)
req := &types.Request{
Query: "machine learning",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
// No Limit set, should use config's max_results (10)
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
// Should respect max_results from config (+1 for possible answer box)
assert.LessOrEqual(t, result.Total, ast.Search.Web.MaxResults+1)
t.Logf("Search without limit returned %d results (max: %d)", result.Total, ast.Search.Web.MaxResults)
// Execute search with explicit limit
req2 := &types.Request{
Query: "machine learning",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 3, // Override config's max_results
}
result2, err := provider.Search(req2)
require.NoError(t, err)
require.NotNil(t, result2)
// API key must be valid - search should succeed
require.Empty(t, result2.Error, "Search should succeed with valid API key, got error: %s", result2.Error)
// Should respect request's limit (+1 for possible answer box)
assert.LessOrEqual(t, result2.Total, 4)
t.Logf("Search with limit=3 returned %d results", result2.Total)
}
// TestSerpAPIProviderWithBingEngine tests SerpAPIProvider with Bing search engine
func TestSerpAPIProviderWithBingEngine(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serpapi test assistant to get base config
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create config with Bing engine
bingConfig := &types.WebConfig{
Provider: "serpapi",
APIKeyEnv: ast.Search.Web.APIKeyEnv,
MaxResults: 5,
Engine: "bing",
}
// Create SerpAPIProvider with Bing engine
provider := web.NewSerpAPIProvider(bingConfig)
require.NotNil(t, provider)
// Execute search
req := &types.Request{
Query: "Golang programming",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// Verify result structure
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "Golang programming", result.Query)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Bing search should succeed with valid API key, got error: %s", result.Error)
// Verify we got results
assert.Greater(t, result.Total, 0)
assert.NotEmpty(t, result.Items)
t.Logf("Bing search returned %d results in %dms", result.Total, result.Duration)
}
// TestSerpAPIProviderEngineDefault tests that default engine is Google
func TestSerpAPIProviderEngineDefault(t *testing.T) {
// Create provider with config that has no engine specified
cfg := &types.WebConfig{
Provider: "serpapi",
APIKeyEnv: "SERPAPI_API_KEY",
MaxResults: 10,
// Engine not set - should default to "google"
}
provider := web.NewSerpAPIProvider(cfg)
require.NotNil(t, provider)
// We can't directly check the engine field since it's private,
// but we verify the provider is created successfully
// The actual engine usage is tested in integration tests
}
// containsSite checks if url contains the site domain
func containsSite(url, site string) bool {
return len(url) >= len(site) && containsHelper(url, site)
}
// containsHelper is a helper function for string containment check
func containsHelper(s, substr string) bool {
for i := 0; i <= len(s)-len(substr); i++ {
if s[i:i+len(substr)] == substr {
return true
}
}
return false
}

280
agent/search/handlers/web/serper.go Normal file
View file

@ -0,0 +1,280 @@
package web
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"time"
"github.com/yaoapp/yao/agent/search/types"
)
const (
serperAPIURL = "https://google.serper.dev/search"
serperAPITimeout = 30 * time.Second
)
// SerperProvider implements web search using Serper API (serper.dev)
type SerperProvider struct {
apiKey string
maxResults int
}
// NewSerperProvider creates a new Serper provider
func NewSerperProvider(cfg *types.WebConfig) *SerperProvider {
apiKey := ""
if cfg != nil && cfg.APIKeyEnv != "" {
// Support both "$ENV.VAR_NAME" and "VAR_NAME" formats
envName := cfg.APIKeyEnv
if len(envName) > 5 && envName[:5] == "$ENV." {
envName = envName[5:]
}
apiKey = os.Getenv(envName)
}
maxResults := 10
if cfg != nil && cfg.MaxResults > 0 {
maxResults = cfg.MaxResults
}
return &SerperProvider{
apiKey: apiKey,
maxResults: maxResults,
}
}
// serperRequest represents the request body for Serper API
type serperRequest struct {
Q string `json:"q"` // Search query
Num int `json:"num,omitempty"` // Number of results (default: 10, max: 100)
GL string `json:"gl,omitempty"` // Country code (e.g., "us", "cn")
HL string `json:"hl,omitempty"` // Language code (e.g., "en", "zh-cn")
TBS string `json:"tbs,omitempty"` // Time-based search (qdr:h, qdr:d, qdr:w, qdr:m, qdr:y)
Page int `json:"page,omitempty"` // Page number (default: 1)
AutoCor bool `json:"autocorrect"` // Auto-correct spelling
}
// serperResponse represents the response from Serper API
type serperResponse struct {
SearchParameters serperSearchParams `json:"searchParameters"`
Organic []serperResult `json:"organic"`
AnswerBox *serperAnswerBox `json:"answerBox,omitempty"`
KnowledgeGraph *serperKnowledge `json:"knowledgeGraph,omitempty"`
RelatedSearches []serperRelated `json:"relatedSearches,omitempty"`
}
// serperSearchParams contains search parameters from response
type serperSearchParams struct {
Q string `json:"q"`
Type string `json:"type"`
GL string `json:"gl"`
HL string `json:"hl"`
Num int `json:"num"`
}
// serperResult represents a single organic search result
type serperResult struct {
Title string `json:"title"`
Link string `json:"link"`
Snippet string `json:"snippet"`
Position int `json:"position"`
Date string `json:"date,omitempty"`
}
// serperAnswerBox represents the answer box (featured snippet)
type serperAnswerBox struct {
Title string `json:"title,omitempty"`
Snippet string `json:"snippet,omitempty"`
Link string `json:"link,omitempty"`
}
// serperKnowledge represents knowledge graph data
type serperKnowledge struct {
Title string `json:"title,omitempty"`
Type string `json:"type,omitempty"`
Description string `json:"description,omitempty"`
}
// serperRelated represents related searches
type serperRelated struct {
Query string `json:"query"`
}
// Search executes a web search using Serper API
func (p *SerperProvider) Search(req *types.Request) (*types.Result, error) {
startTime := time.Now()
// Validate API key
if p.apiKey == "" {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Error: "Serper API key not configured",
}, nil
}
// Determine max results
maxResults := p.maxResults
if req.Limit > 0 {
maxResults = req.Limit
}
// Build search query with site restrictions if specified
query := req.Query
if len(req.Sites) > 0 {
// Serper uses "site:domain" syntax in query
if len(req.Sites) == 1 {
query = "site:" + req.Sites[0] + " " + req.Query
} else {
// Multiple sites: (site:domain1 OR site:domain2) query
siteQuery := ""
for i, site := range req.Sites {
if i > 0 {
siteQuery += " OR "
}
siteQuery += "site:" + site
}
query = "(" + siteQuery + ") " + req.Query
}
}
// Build request body
serperReq := serperRequest{
Q: query,
Num: maxResults,
AutoCor: true,
}
// Add time range if specified
if req.TimeRange != "" {
serperReq.TBS = convertSerperTimeRange(req.TimeRange)
}
// Execute API call
serperResp, err := p.callAPI(&serperReq)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: fmt.Sprintf("Serper API error: %v", err),
}, nil
}
// Convert results
items := make([]*types.ResultItem, 0, len(serperResp.Organic))
// Add answer box as first result if available
if serperResp.AnswerBox != nil && serperResp.AnswerBox.Snippet != "" {
items = append(items, &types.ResultItem{
Type: types.SearchTypeWeb,
Title: serperResp.AnswerBox.Title,
Content: serperResp.AnswerBox.Snippet,
URL: serperResp.AnswerBox.Link,
Score: 1.0, // Featured snippet gets highest score
Source: req.Source,
Metadata: map[string]interface{}{
"type": "answer_box",
},
})
}
// Add organic results
for _, r := range serperResp.Organic {
// Calculate score based on position (1st = 0.95, 2nd = 0.90, etc.)
score := 1.0 - float64(r.Position)*0.05
if score < 0.1 {
score = 0.1
}
items = append(items, &types.ResultItem{
Type: types.SearchTypeWeb,
Title: r.Title,
Content: r.Snippet,
URL: r.Link,
Score: score,
Source: req.Source,
})
}
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: items,
Total: len(items),
Duration: time.Since(startTime).Milliseconds(),
}, nil
}
// callAPI makes the HTTP POST request to Serper API
func (p *SerperProvider) callAPI(req *serperRequest) (*serperResponse, error) {
// Serialize request body
body, err := json.Marshal(req)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
// Create HTTP request
httpReq, err := http.NewRequest(http.MethodPost, serperAPIURL, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("X-API-KEY", p.apiKey)
// Execute request
client := &http.Client{Timeout: serperAPITimeout}
resp, err := client.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("request failed: %w", err)
}
defer resp.Body.Close()
// Read response body
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
// Check status code
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(respBody))
}
// Parse response
var serperResp serperResponse
if err := json.Unmarshal(respBody, &serperResp); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
return &serperResp, nil
}
// convertSerperTimeRange converts time range to Serper tbs format
func convertSerperTimeRange(timeRange string) string {
switch timeRange {
case "hour":
return "qdr:h"
case "day":
return "qdr:d"
case "week":
return "qdr:w"
case "month":
return "qdr:m"
case "year":
return "qdr:y"
default:
return ""
}
}

327
agent/search/handlers/web/serper_test.go Normal file
View file

@ -0,0 +1,327 @@
package web_test
import (
"os"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/search/handlers/web"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
)
// skipIfNoSerperKey skips the test if SERPER_API_KEY is not set
// Note: Serper (serper.dev) requires registration at https://serper.dev
func skipIfNoSerperKey(t *testing.T) {
if os.Getenv("SERPER_API_KEY") == "" {
t.Skip("Skipping Serper test: SERPER_API_KEY not set. Register at https://serper.dev for free 2500 queries.")
}
}
// TestSerperProviderWithAssistantConfig tests SerperProvider using web-serper assistant config
func TestSerperProviderWithAssistantConfig(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
skipIfNoSerperKey(t)
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serper test assistant to get its config
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
require.NoError(t, err)
require.NotNil(t, ast)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Verify assistant config
assert.Equal(t, "tests.web-serper", ast.ID)
assert.Equal(t, "serper", ast.Search.Web.Provider)
assert.Equal(t, "$ENV.SERPER_API_KEY", ast.Search.Web.APIKeyEnv)
assert.Equal(t, 10, ast.Search.Web.MaxResults)
// Create SerperProvider with assistant's web config
provider := web.NewSerperProvider(ast.Search.Web)
require.NotNil(t, provider)
// Execute search
req := &types.Request{
Query: "Yao App Engine",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// Verify result structure
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "Yao App Engine", result.Query)
assert.Equal(t, types.SourceAuto, result.Source)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
// Verify we got results
assert.Greater(t, result.Total, 0)
assert.NotEmpty(t, result.Items)
assert.Greater(t, result.Duration, int64(0))
// Verify result item structure
for _, item := range result.Items {
assert.Equal(t, types.SearchTypeWeb, item.Type)
assert.Equal(t, types.SourceAuto, item.Source)
assert.NotEmpty(t, item.Title)
assert.NotEmpty(t, item.URL)
assert.Greater(t, item.Score, 0.0)
}
t.Logf("Search returned %d results in %dms", result.Total, result.Duration)
}
// TestSerperProviderWithSiteRestriction tests SerperProvider with domain restriction
func TestSerperProviderWithSiteRestriction(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
skipIfNoSerperKey(t)
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serper test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create SerperProvider
provider := web.NewSerperProvider(ast.Search.Web)
// Execute search with site restriction
req := &types.Request{
Query: "documentation",
Type: types.SearchTypeWeb,
Source: types.SourceHook,
Sites: []string{"github.com"},
Limit: 3,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, types.SourceHook, result.Source)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
require.NotEmpty(t, result.Items, "Search should return results")
// All results should be from github.com
for _, item := range result.Items {
assert.Contains(t, item.URL, "github.com", "Result URL should be from github.com")
}
t.Logf("Site-restricted search returned %d results from github.com", result.Total)
}
// TestSerperProviderWithMultipleSites tests SerperProvider with multiple domain restrictions
func TestSerperProviderWithMultipleSites(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
skipIfNoSerperKey(t)
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serper test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create SerperProvider
provider := web.NewSerperProvider(ast.Search.Web)
// Execute search with multiple site restrictions
req := &types.Request{
Query: "golang tutorial",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Sites: []string{"github.com", "golang.org"},
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
require.NotEmpty(t, result.Items, "Search should return results")
// Results should be from either github.com or golang.org
for _, item := range result.Items {
isValidSite := false
for _, site := range req.Sites {
if contains(item.URL, site) {
isValidSite = true
break
}
}
assert.True(t, isValidSite, "Result URL should be from github.com or golang.org: %s", item.URL)
}
t.Logf("Multi-site search returned %d results", result.Total)
}
// TestSerperProviderWithTimeRange tests SerperProvider with time range filter
func TestSerperProviderWithTimeRange(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
skipIfNoSerperKey(t)
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serper test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create SerperProvider
provider := web.NewSerperProvider(ast.Search.Web)
// Execute search with time range
req := &types.Request{
Query: "artificial intelligence news",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
TimeRange: "week", // Last week
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
t.Logf("Time-ranged search (last week) returned %d results in %dms", result.Total, result.Duration)
}
// TestSerperProviderWithoutAPIKey tests graceful degradation when API key is missing
func TestSerperProviderWithoutAPIKey(t *testing.T) {
// Create provider with nil config (no API key)
provider := web.NewSerperProvider(nil)
require.NotNil(t, provider)
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
}
result, err := provider.Search(req)
// Should not return error, but result should have error message
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "test query", result.Query)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "API key")
assert.Empty(t, result.Items)
assert.Equal(t, 0, result.Total)
}
// TestSerperProviderWithEmptyConfig tests provider with empty config
func TestSerperProviderWithEmptyConfig(t *testing.T) {
// Create provider with empty config
cfg := &types.WebConfig{}
provider := web.NewSerperProvider(cfg)
require.NotNil(t, provider)
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceUser,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "API key")
}
// TestSerperProviderMaxResults tests that max_results from config is respected
func TestSerperProviderMaxResults(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
skipIfNoSerperKey(t)
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serper test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create SerperProvider
provider := web.NewSerperProvider(ast.Search.Web)
// Execute search without limit (should use config's max_results)
req := &types.Request{
Query: "machine learning",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
// No Limit set, should use config's max_results (10)
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
// Should respect max_results from config (+1 for possible answer box)
assert.LessOrEqual(t, result.Total, ast.Search.Web.MaxResults+1)
t.Logf("Search without limit returned %d results (max: %d)", result.Total, ast.Search.Web.MaxResults)
// Execute search with explicit limit
req2 := &types.Request{
Query: "machine learning",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 3, // Override config's max_results
}
result2, err := provider.Search(req2)
require.NoError(t, err)
require.NotNil(t, result2)
// API key must be valid - search should succeed
require.Empty(t, result2.Error, "Search should succeed with valid API key, got error: %s", result2.Error)
// Should respect request's limit (+1 for possible answer box)
assert.LessOrEqual(t, result2.Total, 4)
t.Logf("Search with limit=3 returned %d results", result2.Total)
}
// contains checks if s contains substr (uses containsSite from serpapi_test.go)
func contains(s, substr string) bool {
return containsSite(s, substr)
}

192
agent/search/handlers/web/tavily.go Normal file
View file

@ -0,0 +1,192 @@
package web
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"time"
"github.com/yaoapp/yao/agent/search/types"
)
const (
tavilyAPIURL = "https://api.tavily.com/search"
tavilyAPITimeout = 30 * time.Second
)
// TavilyProvider implements web search using Tavily API
type TavilyProvider struct {
apiKey string
maxResults int
}
// NewTavilyProvider creates a new Tavily provider
func NewTavilyProvider(cfg *types.WebConfig) *TavilyProvider {
apiKey := ""
if cfg != nil && cfg.APIKeyEnv != "" {
// Support both "$ENV.VAR_NAME" and "VAR_NAME" formats
envName := cfg.APIKeyEnv
if len(envName) > 5 && envName[:5] == "$ENV." {
envName = envName[5:]
}
apiKey = os.Getenv(envName)
}
maxResults := 10
if cfg != nil && cfg.MaxResults > 0 {
maxResults = cfg.MaxResults
}
return &TavilyProvider{
apiKey: apiKey,
maxResults: maxResults,
}
}
// tavilyRequest represents the request body for Tavily API
type tavilyRequest struct {
APIKey string `json:"api_key"`
Query string `json:"query"`
SearchDepth string `json:"search_depth,omitempty"` // "basic" or "advanced"
IncludeAnswer bool `json:"include_answer,omitempty"` // Include AI-generated answer
IncludeRawContent bool `json:"include_raw_content,omitempty"` // Include raw HTML content
MaxResults int `json:"max_results,omitempty"` // Max number of results
IncludeDomains []string `json:"include_domains,omitempty"` // Limit to specific domains
ExcludeDomains []string `json:"exclude_domains,omitempty"` // Exclude specific domains
}
// tavilyResponse represents the response from Tavily API
type tavilyResponse struct {
Query string `json:"query"`
Answer string `json:"answer,omitempty"`
Results []tavilyResult `json:"results"`
}
// tavilyResult represents a single search result from Tavily
type tavilyResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
Score float64 `json:"score"`
RawContent string `json:"raw_content,omitempty"`
}
// Search executes a web search using Tavily API
func (p *TavilyProvider) Search(req *types.Request) (*types.Result, error) {
startTime := time.Now()
// Validate API key
if p.apiKey == "" {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Error: "Tavily API key not configured",
}, nil
}
// Determine max results
maxResults := p.maxResults
if req.Limit > 0 {
maxResults = req.Limit
}
// Build request body
tavilyReq := tavilyRequest{
APIKey: p.apiKey,
Query: req.Query,
SearchDepth: "basic",
IncludeAnswer: false,
MaxResults: maxResults,
}
// Add domain restrictions if specified
if len(req.Sites) > 0 {
tavilyReq.IncludeDomains = req.Sites
}
// Execute API call
tavilyResp, err := p.callAPI(&tavilyReq)
if err != nil {
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(startTime).Milliseconds(),
Error: fmt.Sprintf("Tavily API error: %v", err),
}, nil
}
// Convert results
items := make([]*types.ResultItem, 0, len(tavilyResp.Results))
for _, r := range tavilyResp.Results {
items = append(items, &types.ResultItem{
Type: types.SearchTypeWeb,
Title: r.Title,
Content: r.Content,
URL: r.URL,
Score: r.Score,
Source: req.Source,
})
}
return &types.Result{
Type: types.SearchTypeWeb,
Query: req.Query,
Source: req.Source,
Items: items,
Total: len(items),
Duration: time.Since(startTime).Milliseconds(),
}, nil
}
// callAPI makes the HTTP request to Tavily API
func (p *TavilyProvider) callAPI(req *tavilyRequest) (*tavilyResponse, error) {
// Serialize request body
body, err := json.Marshal(req)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
// Create HTTP request
httpReq, err := http.NewRequest(http.MethodPost, tavilyAPIURL, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
// Execute request
client := &http.Client{Timeout: tavilyAPITimeout}
resp, err := client.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("request failed: %w", err)
}
defer resp.Body.Close()
// Read response body
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
// Check status code
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(respBody))
}
// Parse response
var tavilyResp tavilyResponse
if err := json.Unmarshal(respBody, &tavilyResp); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
return &tavilyResp, nil
}

223
agent/search/handlers/web/tavily_test.go Normal file
View file

@ -0,0 +1,223 @@
package web_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/search/handlers/web"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
)
// TestTavilyProviderWithAssistantConfig tests TavilyProvider using web-tavily assistant config
func TestTavilyProviderWithAssistantConfig(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-tavily test assistant to get its config
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
require.NoError(t, err)
require.NotNil(t, ast)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Verify assistant config
assert.Equal(t, "tests.web-tavily", ast.ID)
assert.Equal(t, "tavily", ast.Search.Web.Provider)
assert.Equal(t, "$ENV.TAVILY_API_KEY", ast.Search.Web.APIKeyEnv)
assert.Equal(t, 10, ast.Search.Web.MaxResults)
// Create TavilyProvider with assistant's web config
provider := web.NewTavilyProvider(ast.Search.Web)
require.NotNil(t, provider)
// Execute search
req := &types.Request{
Query: "Yao App Engine",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 5,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// Verify result structure
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "Yao App Engine", result.Query)
assert.Equal(t, types.SourceAuto, result.Source)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
// Verify we got results
assert.Greater(t, result.Total, 0)
assert.NotEmpty(t, result.Items)
assert.Greater(t, result.Duration, int64(0))
// Verify result item structure
for _, item := range result.Items {
assert.Equal(t, types.SearchTypeWeb, item.Type)
assert.Equal(t, types.SourceAuto, item.Source)
assert.NotEmpty(t, item.Title)
assert.NotEmpty(t, item.URL)
// Content may be empty for some results
}
t.Logf("Search returned %d results in %dms", result.Total, result.Duration)
}
// TestTavilyProviderWithSiteRestriction tests TavilyProvider with domain restriction
func TestTavilyProviderWithSiteRestriction(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-tavily test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create TavilyProvider
provider := web.NewTavilyProvider(ast.Search.Web)
// Execute search with site restriction
req := &types.Request{
Query: "documentation",
Type: types.SearchTypeWeb,
Source: types.SourceHook,
Sites: []string{"github.com"},
Limit: 3,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, types.SourceHook, result.Source)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
require.NotEmpty(t, result.Items, "Search should return results")
// All results should be from github.com
for _, item := range result.Items {
assert.Contains(t, item.URL, "github.com", "Result URL should be from github.com")
}
t.Logf("Site-restricted search returned %d results from github.com", result.Total)
}
// TestTavilyProviderWithoutAPIKey tests graceful degradation when API key is missing
func TestTavilyProviderWithoutAPIKey(t *testing.T) {
// Create provider with nil config (no API key)
provider := web.NewTavilyProvider(nil)
require.NotNil(t, provider)
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
}
result, err := provider.Search(req)
// Should not return error, but result should have error message
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "test query", result.Query)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "API key")
assert.Empty(t, result.Items)
assert.Equal(t, 0, result.Total)
}
// TestTavilyProviderWithEmptyConfig tests provider with empty config
func TestTavilyProviderWithEmptyConfig(t *testing.T) {
// Create provider with empty config
cfg := &types.WebConfig{}
provider := web.NewTavilyProvider(cfg)
require.NotNil(t, provider)
req := &types.Request{
Query: "test query",
Type: types.SearchTypeWeb,
Source: types.SourceUser,
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
assert.NotEmpty(t, result.Error)
assert.Contains(t, result.Error, "API key")
}
// TestTavilyProviderMaxResults tests that max_results from config is respected
func TestTavilyProviderMaxResults(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test")
}
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-tavily test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Create TavilyProvider
provider := web.NewTavilyProvider(ast.Search.Web)
// Execute search without limit (should use config's max_results)
req := &types.Request{
Query: "artificial intelligence",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
// No Limit set, should use config's max_results (10)
}
result, err := provider.Search(req)
require.NoError(t, err)
require.NotNil(t, result)
// API key must be valid - search should succeed
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
// Should respect max_results from config
assert.LessOrEqual(t, result.Total, ast.Search.Web.MaxResults)
t.Logf("Search without limit returned %d results (max: %d)", result.Total, ast.Search.Web.MaxResults)
// Execute search with explicit limit
req2 := &types.Request{
Query: "artificial intelligence",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Limit: 3, // Override config's max_results
}
result2, err := provider.Search(req2)
require.NoError(t, err)
require.NotNil(t, result2)
// API key must be valid - search should succeed
require.Empty(t, result2.Error, "Search should succeed with valid API key, got error: %s", result2.Error)
// Should respect request's limit
assert.LessOrEqual(t, result2.Total, 3)
t.Logf("Search with limit=3 returned %d results", result2.Total)
}

14
agent/search/interfaces/handler.go Normal file
View file

@ -0,0 +1,14 @@
package interfaces
import (
"github.com/yaoapp/yao/agent/search/types"
)
// Handler defines the interface for search implementations
type Handler interface {
// Type returns the search type this handler supports
Type() types.SearchType
// Search executes the search and returns results
Search(req *types.Request) (*types.Result, error)
}

24
agent/search/interfaces/nlp.go Normal file
View file

@ -0,0 +1,24 @@
package interfaces
import (
"github.com/yaoapp/gou/model"
"github.com/yaoapp/gou/query/gou"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// KeywordExtractor extracts keywords for web search
type KeywordExtractor interface {
// Extract extracts search keywords from user message
// ctx is required for Agent and MCP modes, can be nil for builtin mode
Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error)
}
// QueryDSLGenerator generates QueryDSL for DB search
type QueryDSLGenerator interface {
// Generate converts natural language to QueryDSL
Generate(query string, models []*model.Model) (*gou.QueryDSL, error)
}
// Note: Embedding is handled by KB collection's own config (embedding provider + model),
// not defined here. See KB handler for details.

13
agent/search/interfaces/reranker.go Normal file
View file

@ -0,0 +1,13 @@
package interfaces
import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// Reranker reorders search results by relevance
type Reranker interface {
// Rerank reorders results based on query relevance
// ctx is required for Agent and MCP modes, can be nil for builtin mode
Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error)
}

23
agent/search/interfaces/searcher.go Normal file
View file

@ -0,0 +1,23 @@
package interfaces
import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// Searcher is the main interface exposed to external callers
type Searcher interface {
// Search executes a single search request
Search(ctx *context.Context, req *types.Request) (*types.Result, error)
// Parallel search methods - inspired by JavaScript Promise
// All waits for all searches to complete (like Promise.all)
All(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error)
// Any returns when any search succeeds with results (like Promise.any)
Any(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error)
// Race returns when any search completes (like Promise.race)
Race(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error)
// BuildReferences converts search results to unified Reference format for LLM
BuildReferences(results []*types.Result) []*types.Reference
}

229
agent/search/jsapi.go Normal file
View file

@ -0,0 +1,229 @@
package search
import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// JSAPI implements context.SearchAPI interface
// Provides ctx.search.Web(), ctx.search.KB(), ctx.search.DB(), ctx.search.All(), ctx.search.Any(), ctx.search.Race()
type JSAPI struct {
ctx *context.Context
searcher *Searcher
}
// NewJSAPI creates a new search JSAPI instance
func NewJSAPI(ctx *context.Context, config *types.Config, uses *Uses) *JSAPI {
return &JSAPI{
ctx: ctx,
searcher: New(config, uses),
}
}
// Web executes web search
// Options:
// - limit: int - max results (default: 10)
// - sites: []string - restrict to specific sites
// - time_range: string - "day", "week", "month", "year"
// - rerank: map[string]interface{} - rerank options
func (api *JSAPI) Web(query string, opts map[string]interface{}) interface{} {
req := api.buildRequest(types.SearchTypeWeb, query, opts)
result, _ := api.searcher.Search(api.ctx, req)
return result
}
// KB executes knowledge base search
// Options:
// - collections: []string - collection IDs
// - threshold: float64 - similarity threshold (0-1)
// - limit: int - max results
// - graph: bool - enable graph association
// - rerank: map[string]interface{} - rerank options
func (api *JSAPI) KB(query string, opts map[string]interface{}) interface{} {
req := api.buildRequest(types.SearchTypeKB, query, opts)
result, _ := api.searcher.Search(api.ctx, req)
return result
}
// DB executes database search
// Options:
// - models: []string - model IDs
// - wheres: []map[string]interface{} - pre-defined filters (GOU QueryDSL Where format)
// - orders: []map[string]interface{} - sort orders (GOU QueryDSL Order format)
// - select: []string - fields to return
// - limit: int - max results
// - rerank: map[string]interface{} - rerank options
func (api *JSAPI) DB(query string, opts map[string]interface{}) interface{} {
req := api.buildRequest(types.SearchTypeDB, query, opts)
result, _ := api.searcher.Search(api.ctx, req)
return result
}
// All executes all searches and waits for all to complete (like Promise.all)
// Each request should have:
// - type: string - "web", "kb", or "db"
// - query: string - search query
// - ... other type-specific options
func (api *JSAPI) All(requests []interface{}) []interface{} {
reqs := api.parseRequests(requests)
results, _ := api.searcher.All(api.ctx, reqs)
return api.convertResults(results)
}
// Any returns as soon as any search succeeds with results (like Promise.any)
// Each request should have:
// - type: string - "web", "kb", or "db"
// - query: string - search query
// - ... other type-specific options
func (api *JSAPI) Any(requests []interface{}) []interface{} {
reqs := api.parseRequests(requests)
results, _ := api.searcher.Any(api.ctx, reqs)
return api.convertResults(results)
}
// Race returns as soon as any search completes (like Promise.race)
// Each request should have:
// - type: string - "web", "kb", or "db"
// - query: string - search query
// - ... other type-specific options
func (api *JSAPI) Race(requests []interface{}) []interface{} {
reqs := api.parseRequests(requests)
results, _ := api.searcher.Race(api.ctx, reqs)
return api.convertResults(results)
}
// buildRequest builds a Request from query and options
func (api *JSAPI) buildRequest(searchType types.SearchType, query string, opts map[string]interface{}) *types.Request {
req := &types.Request{
Type: searchType,
Query: query,
Source: types.SourceHook, // JSAPI calls are from hooks
}
if opts == nil {
return req
}
// Common options
if limit, ok := opts["limit"].(float64); ok {
req.Limit = int(limit)
} else if limit, ok := opts["limit"].(int); ok {
req.Limit = limit
}
// Web-specific options
if searchType == types.SearchTypeWeb {
if sites, ok := opts["sites"].([]interface{}); ok {
req.Sites = toStringSlice(sites)
}
if timeRange, ok := opts["time_range"].(string); ok {
req.TimeRange = timeRange
}
}
// KB-specific options
if searchType == types.SearchTypeKB {
if collections, ok := opts["collections"].([]interface{}); ok {
req.Collections = toStringSlice(collections)
}
if threshold, ok := opts["threshold"].(float64); ok {
req.Threshold = threshold
}
if graph, ok := opts["graph"].(bool); ok {
req.Graph = graph
}
}
// DB-specific options
if searchType == types.SearchTypeDB {
if models, ok := opts["models"].([]interface{}); ok {
req.Models = toStringSlice(models)
}
if selectFields, ok := opts["select"].([]interface{}); ok {
req.Select = toStringSlice(selectFields)
}
// Note: wheres and orders are more complex, handled by QueryDSL generator
}
// Rerank options
if rerankOpts, ok := opts["rerank"].(map[string]interface{}); ok {
req.Rerank = &types.RerankOptions{}
if topN, ok := rerankOpts["top_n"].(float64); ok {
req.Rerank.TopN = int(topN)
} else if topN, ok := rerankOpts["top_n"].(int); ok {
req.Rerank.TopN = topN
}
}
return req
}
// parseRequests parses an array of request objects into typed Requests
func (api *JSAPI) parseRequests(requests []interface{}) []*types.Request {
reqs := make([]*types.Request, 0, len(requests))
for _, r := range requests {
reqMap, ok := r.(map[string]interface{})
if !ok {
continue
}
// Get type
typeStr, ok := reqMap["type"].(string)
if !ok {
continue
}
searchType := types.SearchType(typeStr)
// Get query
query, ok := reqMap["query"].(string)
if !ok {
continue
}
// Build request with remaining options
req := api.buildRequest(searchType, query, reqMap)
reqs = append(reqs, req)
}
return reqs
}
// convertResults converts typed Results to interface slice for JS
func (api *JSAPI) convertResults(results []*types.Result) []interface{} {
out := make([]interface{}, len(results))
for i, r := range results {
out[i] = r
}
return out
}
// toStringSlice converts []interface{} to []string
func toStringSlice(arr []interface{}) []string {
result := make([]string, 0, len(arr))
for _, v := range arr {
if s, ok := v.(string); ok {
result = append(result, s)
}
}
return result
}
// ConfigGetter is a function type that retrieves search config and uses for an assistant
type ConfigGetter func(assistantID string) (*types.Config, *Uses)
// configGetter is set by assistant package during initialization
var configGetter ConfigGetter
// SetJSAPIFactory sets the factory function for creating SearchAPI instances
// Called by assistant package during initialization
// getter: function to get search config and uses from assistant ID
func SetJSAPIFactory(getter ConfigGetter) {
configGetter = getter
context.SearchAPIFactory = func(ctx *context.Context) context.SearchAPI {
var config *types.Config
var uses *Uses
if configGetter != nil && ctx.AssistantID != "" {
config, uses = configGetter(ctx.AssistantID)
}
return NewJSAPI(ctx, config, uses)
}
}

328
agent/search/jsapi_test.go Normal file
View file

@ -0,0 +1,328 @@
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"
)
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) {
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) {
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) {
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) {
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) {
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) {
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) {
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.Context{}
// 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.Context{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)
}

175
agent/search/nlp/keyword/agent.go Normal file
View file

@ -0,0 +1,175 @@
package keyword
import (
"encoding/json"
"fmt"
"github.com/yaoapp/yao/agent/caller"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// AgentProvider delegates keyword extraction to an LLM-powered assistant
// The assistant can understand context and extract semantically relevant keywords
type AgentProvider struct {
agentID string // Assistant ID to delegate to
}
// NewAgentProvider creates a new agent-based keyword extractor
func NewAgentProvider(agentID string) *AgentProvider {
return &AgentProvider{
agentID: agentID,
}
}
// Extract extracts keywords by calling the target agent
// The agent receives the content and returns extracted keywords
func (p *AgentProvider) Extract(ctx *agentContext.Context, content string, opts *types.KeywordOptions) ([]string, error) {
if ctx == nil {
return nil, fmt.Errorf("context is required for agent keyword extraction")
}
// Check if AgentGetterFunc is initialized
if caller.AgentGetterFunc == nil {
return nil, fmt.Errorf("AgentGetterFunc not initialized")
}
// Get the agent
agent, err := caller.AgentGetterFunc(p.agentID)
if err != nil {
return nil, fmt.Errorf("failed to get agent %s: %w", p.agentID, err)
}
// Build the request message
requestData := map[string]interface{}{
"content": content,
"max_keywords": opts.MaxKeywords,
"language": opts.Language,
}
requestJSON, _ := json.Marshal(requestData)
// Create message for the agent
messages := []agentContext.Message{
{
Role: "user",
Content: string(requestJSON),
},
}
// Call the agent with skip options (no history, no output)
options := &agentContext.Options{
Skip: &agentContext.Skip{
History: true,
Output: true,
},
}
result, err := agent.Stream(ctx, messages, options)
if err != nil {
return nil, fmt.Errorf("agent call failed: %w", err)
}
// Debug: log the result type and value
// fmt.Printf("DEBUG Agent result type: %T, value: %+v\n", result, result)
// Parse the result
return p.parseResult(result)
}
// parseResult extracts keywords from the agent's response
// The agent should return data in NextHookResponse format: { data: { keywords: [...] } }
// The Stream() response wraps this in: { next: { data: { keywords: [...] } } }
func (p *AgentProvider) parseResult(result interface{}) ([]string, error) {
if result == nil {
return []string{}, nil
}
// Try to convert to map first (most common case)
var data map[string]interface{}
switch v := result.(type) {
case map[string]interface{}:
data = v
case string:
// Try to parse as JSON
if err := json.Unmarshal([]byte(v), &data); err != nil {
// Not a JSON object, try as array
var keywords []string
if err := json.Unmarshal([]byte(v), &keywords); err == nil {
return keywords, nil
}
// Return as single keyword
return []string{v}, nil
}
case []string:
return v, nil
case []interface{}:
keywords := make([]string, 0, len(v))
for _, item := range v {
if s, ok := item.(string); ok {
keywords = append(keywords, s)
}
}
return keywords, nil
default:
// Try to marshal and unmarshal
jsonBytes, err := json.Marshal(result)
if err != nil {
return []string{}, nil
}
if err := json.Unmarshal(jsonBytes, &data); err != nil {
return []string{}, nil
}
}
// Check for "next" field (custom hook data from NextHookResponse)
// Stream() returns: { next: { data: { keywords: [...] } } }
if next, hasNext := data["next"]; hasNext && next != nil {
if nextMap, ok := next.(map[string]interface{}); ok {
data = nextMap
} else if nextStr, ok := next.(string); ok {
if err := json.Unmarshal([]byte(nextStr), &data); err != nil {
return []string{}, nil
}
}
}
// Extract keywords from data
// Try common field names: "keywords", "data", "data.keywords"
if kw, ok := data["keywords"]; ok {
return p.extractKeywordsFromValue(kw)
}
if d, ok := data["data"]; ok {
if dm, ok := d.(map[string]interface{}); ok {
if kw, ok := dm["keywords"]; ok {
return p.extractKeywordsFromValue(kw)
}
}
return p.extractKeywordsFromValue(d)
}
return []string{}, nil
}
// extractKeywordsFromValue extracts string array from various types
func (p *AgentProvider) extractKeywordsFromValue(v interface{}) ([]string, error) {
switch kw := v.(type) {
case []string:
return kw, nil
case []interface{}:
keywords := make([]string, 0, len(kw))
for _, item := range kw {
if s, ok := item.(string); ok {
keywords = append(keywords, s)
}
}
return keywords, nil
case string:
var keywords []string
if err := json.Unmarshal([]byte(kw), &keywords); err == nil {
return keywords, nil
}
return []string{kw}, nil
}
return []string{}, nil
}

120
agent/search/nlp/keyword/agent_test.go Normal file
View file

@ -0,0 +1,120 @@
package keyword_test
import (
"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/search/nlp/keyword"
searchTypes "github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
)
func TestAgentProviderWithAssistantConfig(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the keyword-agent assistant that will provide keywords
ast, err := assistant.Get("tests.keyword-agent")
require.NoError(t, err)
require.NotNil(t, ast)
// Create test context
ctx := newTestContext(t)
// Create extractor with agent mode
extractor := keyword.NewExtractor("tests.keyword-agent", &searchTypes.KeywordConfig{
MaxKeywords: 5,
Language: "auto",
})
// Test extraction
content := "Machine learning and deep learning are subfields of artificial intelligence"
keywords, err := extractor.Extract(ctx, content, nil)
require.NoError(t, err)
assert.NotEmpty(t, keywords, "Agent should return keywords")
assert.LessOrEqual(t, len(keywords), 5, "Should respect max_keywords")
// Verify keywords are relevant
t.Logf("Extracted keywords: %v", keywords)
}
func TestAgentProviderWithCustomOptions(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newTestContext(t)
// Create extractor with agent mode
extractor := keyword.NewExtractor("tests.keyword-agent", &searchTypes.KeywordConfig{
MaxKeywords: 10,
})
// Test with runtime options override
content := "Python programming language for data science and web development"
keywords, err := extractor.Extract(ctx, content, &searchTypes.KeywordOptions{
MaxKeywords: 3, // Override to 3
})
require.NoError(t, err)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 3, "Should respect runtime max_keywords override")
t.Logf("Extracted keywords (max 3): %v", keywords)
}
func TestAgentProviderWithoutContext(t *testing.T) {
// Test that agent mode requires context
extractor := keyword.NewExtractor("tests.keyword-agent", nil)
_, err := extractor.Extract(nil, "test content", nil)
assert.Error(t, err, "Agent mode should require context")
assert.Contains(t, err.Error(), "context is required")
}
func TestAgentProviderAgentNotFound(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newTestContext(t)
// Create extractor with non-existent agent
extractor := keyword.NewExtractor("non-existent-agent", nil)
_, err := extractor.Extract(ctx, "test content", nil)
assert.Error(t, err, "Should error for non-existent agent")
assert.Contains(t, err.Error(), "failed to get agent")
}
// newTestContext creates a test context with required fields
func newTestContext(t *testing.T) *context.Context {
t.Helper()
authorized := &oauthTypes.AuthorizedInfo{
UserID: "test-user",
}
chatID := "test-chat-keyword"
ctx := context.New(t.Context(), authorized, chatID)
return ctx
}

243
agent/search/nlp/keyword/builtin.go Normal file
View file

@ -0,0 +1,243 @@
package keyword
import (
"regexp"
"sort"
"strings"
"unicode"
)
// BuiltinExtractor implements simple frequency-based keyword extraction
// This is a lightweight implementation with no external dependencies.
//
// Algorithm:
// 1. Tokenize text (split by whitespace and punctuation)
// 2. Normalize (lowercase, trim)
// 3. Filter stop words and short words
// 4. Count word frequency
// 5. Return top N words by frequency
//
// Limitations:
// - No semantic understanding
// - No phrase extraction (single words only)
// - Basic Chinese support (splits by punctuation, no proper segmentation)
//
// For better results, use Agent or MCP mode with LLM-based extraction.
type BuiltinExtractor struct {
stopWords map[string]bool
minLength int // minimum word length to consider
}
// Result represents an extracted keyword with its score
type Result struct {
Word string `json:"word"`
Score float64 `json:"score"` // frequency-based score (0-1)
}
// NewBuiltinExtractor creates a new builtin keyword extractor
func NewBuiltinExtractor() *BuiltinExtractor {
return &BuiltinExtractor{
stopWords: defaultStopWords,
minLength: 2,
}
}
// Extract extracts keywords from text using frequency-based algorithm
func (e *BuiltinExtractor) Extract(text string, limit int) []Result {
if text == "" || limit <= 0 {
return []Result{}
}
// Step 1: Tokenize
tokens := e.tokenize(text)
// Step 2 & 3: Normalize and filter
var words []string
for _, token := range tokens {
word := e.normalize(token)
if e.shouldKeep(word) {
words = append(words, word)
}
}
if len(words) == 0 {
return []Result{}
}
// Step 4: Count frequency
freq := make(map[string]int)
for _, word := range words {
freq[word]++
}
// Step 5: Sort by frequency and return top N
type wordFreq struct {
word string
freq int
}
var sorted []wordFreq
for word, count := range freq {
sorted = append(sorted, wordFreq{word, count})
}
sort.Slice(sorted, func(i, j int) bool {
return sorted[i].freq > sorted[j].freq
})
// Calculate max frequency for normalization
maxFreq := 1
if len(sorted) > 0 {
maxFreq = sorted[0].freq
}
// Build result with normalized scores
result := make([]Result, 0, limit)
for i := 0; i < len(sorted) && i < limit; i++ {
result = append(result, Result{
Word: sorted[i].word,
Score: float64(sorted[i].freq) / float64(maxFreq),
})
}
return result
}
// ExtractAsStrings is a convenience method that returns just the keyword strings
func (e *BuiltinExtractor) ExtractAsStrings(text string, limit int) []string {
results := e.Extract(text, limit)
words := make([]string, len(results))
for i, r := range results {
words[i] = r.Word
}
return words
}
// tokenize splits text into tokens
// Handles both English (space-separated) and Chinese (character-based with punctuation splits)
func (e *BuiltinExtractor) tokenize(text string) []string {
// Split by whitespace and common punctuation
splitter := regexp.MustCompile(`[\s\p{P}\p{S}]+`)
tokens := splitter.Split(text, -1)
// Further split mixed Chinese/English text
var result []string
for _, token := range tokens {
if token == "" {
continue
}
// Split Chinese characters as individual tokens (basic approach)
// For proper Chinese segmentation, use Agent/MCP mode
subTokens := e.splitMixedText(token)
result = append(result, subTokens...)
}
return result
}
// splitMixedText handles mixed Chinese/English text
// Chinese characters are grouped together, English words stay as-is
func (e *BuiltinExtractor) splitMixedText(text string) []string {
var result []string
var current strings.Builder
var lastType int // 0=none, 1=chinese, 2=other
for _, r := range text {
currentType := 0
if unicode.Is(unicode.Han, r) {
currentType = 1
} else if unicode.IsLetter(r) || unicode.IsDigit(r) {
currentType = 2
}
if currentType == 0 {
// Non-word character, flush current
if current.Len() > 0 {
result = append(result, current.String())
current.Reset()
}
lastType = 0
continue
}
if lastType != 0 && lastType != currentType {
// Type changed, flush current
if current.Len() > 0 {
result = append(result, current.String())
current.Reset()
}
}
current.WriteRune(r)
lastType = currentType
}
// Flush remaining
if current.Len() > 0 {
result = append(result, current.String())
}
return result
}
// normalize converts word to lowercase and trims whitespace
func (e *BuiltinExtractor) normalize(word string) string {
return strings.ToLower(strings.TrimSpace(word))
}
// shouldKeep checks if a word should be kept (not a stop word, meets length requirement)
func (e *BuiltinExtractor) shouldKeep(word string) bool {
if len(word) < e.minLength {
return false
}
if e.stopWords[word] {
return false
}
// Keep if it contains at least one letter or Chinese character
for _, r := range word {
if unicode.IsLetter(r) {
return true
}
}
return false
}
// defaultStopWords contains common stop words for English and Chinese
// This is a minimal set to keep the implementation lightweight.
// For comprehensive stop word filtering, use Agent/MCP mode.
var defaultStopWords = map[string]bool{
// English stop words (most common ~100)
"a": true, "an": true, "the": true, "and": true, "or": true, "but": true,
"is": true, "are": true, "was": true, "were": true, "be": true, "been": true, "being": true,
"have": true, "has": true, "had": true, "do": true, "does": true, "did": true,
"will": true, "would": true, "could": true, "should": true, "may": true, "might": true,
"must": true, "shall": true, "can": true, "need": true, "dare": true,
"i": true, "you": true, "he": true, "she": true, "it": true, "we": true, "they": true,
"me": true, "him": true, "her": true, "us": true, "them": true,
"my": true, "your": true, "his": true, "its": true, "our": true, "their": true,
"mine": true, "yours": true, "hers": true, "ours": true, "theirs": true,
"this": true, "that": true, "these": true, "those": true,
"what": true, "which": true, "who": true, "whom": true, "whose": true,
"where": true, "when": true, "why": true, "how": true,
"all": true, "each": true, "every": true, "both": true, "few": true, "more": true,
"most": true, "other": true, "some": true, "such": true, "no": true, "not": true,
"only": true, "same": true, "so": true, "than": true, "too": true, "very": true,
"just": true, "also": true, "now": true, "here": true, "there": true,
"in": true, "on": true, "at": true, "by": true, "for": true, "with": true,
"about": true, "against": true, "between": true, "into": true, "through": true,
"during": true, "before": true, "after": true, "above": true, "below": true,
"to": true, "from": true, "up": true, "down": true, "out": true, "off": true,
"over": true, "under": true, "again": true, "further": true, "then": true, "once": true,
"as": true, "if": true, "because": true, "until": true, "while": true,
// Chinese stop words (most common ~50)
"的": true, "了": true, "和": true, "是": true, "就": true,
"都": true, "而": true, "及": true, "与": true, "着": true,
"或": true, "一个": true, "没有": true, "我们": true, "你们": true,
"他们": true, "它们": true, "这个": true, "那个": true, "这些": true,
"那些": true, "这里": true, "那里": true, "什么": true, "怎么": true,
"为什么": true, "哪里": true, "谁": true, "哪个": true, "多少": true,
"在": true, "有": true, "个": true, "中": true, "为": true,
"以": true, "于": true, "上": true, "下": true, "不": true,
"也": true, "很": true, "到": true, "说": true, "要": true,
"会": true, "可以": true, "这": true, "那": true, "但": true,
"如果": true, "因为": true, "所以": true, "虽然": true, "但是": true,
}

137
agent/search/nlp/keyword/builtin_test.go Normal file
View file

@ -0,0 +1,137 @@
package keyword
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestBuiltinExtractor_Extract(t *testing.T) {
extractor := NewBuiltinExtractor()
tests := []struct {
name string
text string
limit int
minCount int // minimum expected keywords
}{
{
name: "English text",
text: "The quick brown fox jumps over the lazy dog. The fox is very quick.",
limit: 5,
minCount: 3, // fox, quick, etc.
},
{
name: "Chinese text",
text: "人工智能技术正在快速发展，机器学习和深度学习是人工智能的核心技术",
limit: 5,
minCount: 2,
},
{
name: "Mixed text",
text: "AI人工智能 machine learning 机器学习 deep learning 深度学习",
limit: 10,
minCount: 3,
},
{
name: "Empty text",
text: "",
limit: 5,
minCount: 0,
},
{
name: "Only stop words",
text: "the a an is are was were",
limit: 5,
minCount: 0,
},
{
name: "Technical query",
text: "How to implement a search engine with Elasticsearch and Redis caching?",
limit: 5,
minCount: 3, // search, engine, elasticsearch, redis, caching
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
results := extractor.Extract(tt.text, tt.limit)
assert.GreaterOrEqual(t, len(results), tt.minCount, "Expected at least %d keywords", tt.minCount)
assert.LessOrEqual(t, len(results), tt.limit, "Should not exceed limit")
// Check scores are valid
for _, r := range results {
assert.NotEmpty(t, r.Word)
assert.GreaterOrEqual(t, r.Score, 0.0)
assert.LessOrEqual(t, r.Score, 1.0)
}
})
}
}
func TestBuiltinExtractor_ExtractAsStrings(t *testing.T) {
extractor := NewBuiltinExtractor()
text := "Machine learning and deep learning are subfields of artificial intelligence"
keywords := extractor.ExtractAsStrings(text, 5)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 5)
// Check that common ML terms are extracted
keywordSet := make(map[string]bool)
for _, k := range keywords {
keywordSet[k] = true
}
assert.True(t, keywordSet["learning"] || keywordSet["machine"] || keywordSet["artificial"],
"Expected at least one relevant keyword")
}
func TestBuiltinExtractor_StopWords(t *testing.T) {
extractor := NewBuiltinExtractor()
// Test that stop words are filtered
text := "the quick brown fox is very lazy"
results := extractor.Extract(text, 10)
for _, r := range results {
assert.NotEqual(t, "the", r.Word)
assert.NotEqual(t, "is", r.Word)
assert.NotEqual(t, "very", r.Word)
}
}
func TestBuiltinExtractor_Frequency(t *testing.T) {
extractor := NewBuiltinExtractor()
// Word "search" appears 3 times, should rank higher
text := "search engine optimization, search ranking, search results"
results := extractor.Extract(text, 3)
assert.NotEmpty(t, results)
// "search" should be the top keyword
assert.Equal(t, "search", results[0].Word)
assert.Equal(t, 1.0, results[0].Score) // highest frequency = 1.0
}
func TestBuiltinExtractor_ZeroLimit(t *testing.T) {
extractor := NewBuiltinExtractor()
results := extractor.Extract("some text here", 0)
assert.Empty(t, results)
}
func TestBuiltinExtractor_ChineseStopWords(t *testing.T) {
extractor := NewBuiltinExtractor()
// Test that Chinese stop words are filtered
text := "这是一个关于人工智能的文章"
results := extractor.Extract(text, 10)
for _, r := range results {
assert.NotEqual(t, "这", r.Word)
assert.NotEqual(t, "是", r.Word)
assert.NotEqual(t, "一个", r.Word)
assert.NotEqual(t, "的", r.Word)
}
}

106
agent/search/nlp/keyword/extractor.go Normal file
View file

@ -0,0 +1,106 @@
// Package keyword provides keyword extraction for web search optimization
// Supports three modes via uses.keyword configuration:
// - "builtin": Simple frequency-based extraction (no external dependencies)
// - "<assistant-id>": Delegate to an LLM-powered assistant for high-quality extraction
// - "mcp:<server>.<tool>": Call external MCP tool
//
// For production use cases requiring high accuracy, use Agent or MCP mode.
package keyword
import (
"strings"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// Extractor extracts keywords from text
// Mode is determined by uses.keyword configuration
type Extractor struct {
usesKeyword string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
config *types.KeywordConfig // Keyword extraction options
}
// NewExtractor creates a new keyword extractor
// usesKeyword: value from uses.keyword config
// cfg: keyword extraction options from search config
func NewExtractor(usesKeyword string, cfg *types.KeywordConfig) *Extractor {
return &Extractor{
usesKeyword: usesKeyword,
config: cfg,
}
}
// Extract extracts keywords from content based on configured mode
// Returns a list of keywords optimized for search queries
func (e *Extractor) Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) {
// Merge options with config defaults
mergedOpts := e.mergeOptions(opts)
switch {
case e.usesKeyword == "builtin" || e.usesKeyword == "":
return e.builtinExtract(content, mergedOpts)
case strings.HasPrefix(e.usesKeyword, "mcp:"):
return e.mcpExtract(ctx, content, mergedOpts)
default:
// Assume it's an assistant ID for Agent mode
return e.agentExtract(ctx, content, mergedOpts)
}
}
// mergeOptions merges runtime options with config defaults
func (e *Extractor) mergeOptions(opts *types.KeywordOptions) *types.KeywordOptions {
result := &types.KeywordOptions{
MaxKeywords: 10, // default
Language: "auto", // default
}
// Apply config defaults
if e.config != nil {
if e.config.MaxKeywords > 0 {
result.MaxKeywords = e.config.MaxKeywords
}
if e.config.Language != "" {
result.Language = e.config.Language
}
}
// Apply runtime options (highest priority)
if opts != nil {
if opts.MaxKeywords > 0 {
result.MaxKeywords = opts.MaxKeywords
}
if opts.Language != "" {
result.Language = opts.Language
}
}
return result
}
// builtinExtract uses simple frequency-based extraction
// This is a lightweight implementation with no external dependencies.
// For better results, use Agent or MCP mode.
func (e *Extractor) builtinExtract(content string, opts *types.KeywordOptions) ([]string, error) {
extractor := NewBuiltinExtractor()
return extractor.ExtractAsStrings(content, opts.MaxKeywords), nil
}
// agentExtract delegates to an LLM-powered assistant
// The assistant can understand context and extract semantically relevant keywords
func (e *Extractor) agentExtract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) {
provider := NewAgentProvider(e.usesKeyword)
return provider.Extract(ctx, content, opts)
}
// mcpExtract calls an external MCP tool
// Format: "mcp:<server>.<tool>"
func (e *Extractor) mcpExtract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) {
mcpRef := strings.TrimPrefix(e.usesKeyword, "mcp:")
provider, err := NewMCPProvider(mcpRef)
if err != nil {
// Fallback to builtin on invalid MCP format
return e.builtinExtract(content, opts)
}
return provider.Extract(ctx, content, opts)
}

63
agent/search/nlp/keyword/extractor_test.go Normal file
View file

@ -0,0 +1,63 @@
package keyword_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/nlp/keyword"
"github.com/yaoapp/yao/agent/search/types"
)
func TestExtractor_BuiltinMode(t *testing.T) {
// Test builtin mode (no external dependencies)
extractor := keyword.NewExtractor("builtin", &types.KeywordConfig{
MaxKeywords: 5,
Language: "auto",
})
keywords, err := extractor.Extract(nil, "How to build a search engine with Elasticsearch?", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 5)
}
func TestExtractor_EmptyUsesKeyword(t *testing.T) {
// Empty uses.keyword should default to builtin
extractor := keyword.NewExtractor("", nil)
keywords, err := extractor.Extract(nil, "Machine learning algorithms", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords)
}
func TestExtractor_RuntimeOptionsOverride(t *testing.T) {
// Config has max_keywords=10, but runtime opts override to 3
extractor := keyword.NewExtractor("builtin", &types.KeywordConfig{
MaxKeywords: 10,
})
keywords, err := extractor.Extract(nil, "one two three four five six seven eight nine ten", &types.KeywordOptions{
MaxKeywords: 3,
})
assert.NoError(t, err)
assert.LessOrEqual(t, len(keywords), 3)
}
func TestExtractor_ConfigDefaults(t *testing.T) {
// No config, should use defaults
extractor := keyword.NewExtractor("builtin", nil)
keywords, err := extractor.Extract(nil, "Test query for keyword extraction", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 10) // default max_keywords is 10
}
func TestExtractor_InvalidMCPFormat(t *testing.T) {
// Invalid MCP format should fallback to builtin
extractor := keyword.NewExtractor("mcp:invalid", nil)
keywords, err := extractor.Extract(nil, "Test query", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords)
}

123
agent/search/nlp/keyword/mcp.go Normal file
View file

@ -0,0 +1,123 @@
package keyword
import (
"encoding/json"
"fmt"
"strings"
"github.com/yaoapp/gou/mcp"
gouMCPTypes "github.com/yaoapp/gou/mcp/types"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// MCPProvider delegates keyword extraction to an MCP tool
type MCPProvider struct {
serverID string // MCP server ID
toolName string // Tool name to call
}
// NewMCPProvider creates a new MCP-based keyword extractor
// mcpRef format: "server.tool" (e.g., "nlp.extract_keywords")
func NewMCPProvider(mcpRef string) (*MCPProvider, error) {
parts := strings.SplitN(mcpRef, ".", 2)
if len(parts) != 2 {
return nil, fmt.Errorf("invalid MCP format, expected 'server.tool', got '%s'", mcpRef)
}
return &MCPProvider{
serverID: parts[0],
toolName: parts[1],
}, nil
}
// Extract extracts keywords by calling the MCP tool
func (p *MCPProvider) Extract(ctx *agentContext.Context, content string, opts *types.KeywordOptions) ([]string, error) {
// Get MCP client
client, err := mcp.Select(p.serverID)
if err != nil {
return nil, fmt.Errorf("MCP server '%s' not found: %w", p.serverID, err)
}
// Build arguments for the MCP tool
arguments := map[string]interface{}{
"content": content,
"max_keywords": opts.MaxKeywords,
"language": opts.Language,
}
// Call the MCP tool (ctx embeds context.Context)
callResult, err := client.CallTool(ctx, p.toolName, arguments)
if err != nil {
return nil, fmt.Errorf("MCP tool call failed: %w", err)
}
// Parse the result
return p.parseResult(callResult)
}
// parseResult extracts keywords from the MCP tool response
func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse) ([]string, error) {
if result == nil {
return []string{}, nil
}
// Check for errors in result
if result.IsError {
errMsg := "MCP tool returned error"
if len(result.Content) > 0 && result.Content[0].Text != "" {
errMsg = result.Content[0].Text
}
return nil, fmt.Errorf("%s", errMsg)
}
// Parse content - expect JSON data with "keywords" field
if len(result.Content) == 0 {
return []string{}, nil
}
// Try to extract keywords from content
for _, content := range result.Content {
// Check text content type
if content.Type == gouMCPTypes.ToolContentTypeText && content.Text != "" {
// Try to parse as JSON
var data map[string]interface{}
if err := json.Unmarshal([]byte(content.Text), &data); err == nil {
// Look for "keywords" field
if kw, ok := data["keywords"]; ok {
return p.extractKeywordsFromValue(kw)
}
}
// Try to parse as direct array
var keywords []string
if err := json.Unmarshal([]byte(content.Text), &keywords); err == nil {
return keywords, nil
}
}
}
return []string{}, nil
}
// extractKeywordsFromValue extracts string array from various types
func (p *MCPProvider) extractKeywordsFromValue(v interface{}) ([]string, error) {
switch kw := v.(type) {
case []string:
return kw, nil
case []interface{}:
keywords := make([]string, 0, len(kw))
for _, item := range kw {
if s, ok := item.(string); ok {
keywords = append(keywords, s)
}
}
return keywords, nil
case string:
var keywords []string
if err := json.Unmarshal([]byte(kw), &keywords); err == nil {
return keywords, nil
}
return []string{kw}, nil
}
return []string{}, nil
}

155
agent/search/nlp/keyword/mcp_test.go Normal file
View file

@ -0,0 +1,155 @@
package keyword_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/nlp/keyword"
searchTypes "github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
)
func TestMCPProviderWithAssistantConfig(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with MCP mode
extractor := keyword.NewExtractor("mcp:search.extract_keywords", &searchTypes.KeywordConfig{
MaxKeywords: 5,
Language: "auto",
})
// Test extraction
content := "Machine learning and deep learning are subfields of artificial intelligence"
keywords, err := extractor.Extract(ctx, content, nil)
require.NoError(t, err)
assert.NotEmpty(t, keywords, "MCP should return keywords")
assert.LessOrEqual(t, len(keywords), 5, "Should respect max_keywords")
t.Logf("Extracted keywords via MCP: %v", keywords)
}
func TestMCPProviderWithCustomOptions(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with MCP mode
extractor := keyword.NewExtractor("mcp:search.extract_keywords", &searchTypes.KeywordConfig{
MaxKeywords: 10,
})
// Test with runtime options override
content := "Python programming language for data science and web development"
keywords, err := extractor.Extract(ctx, content, &searchTypes.KeywordOptions{
MaxKeywords: 3, // Override to 3
})
require.NoError(t, err)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 3, "Should respect runtime max_keywords override")
t.Logf("Extracted keywords via MCP (max 3): %v", keywords)
}
func TestMCPProviderInvalidFormat(t *testing.T) {
// Test invalid MCP format fallback to builtin
extractor := keyword.NewExtractor("mcp:invalid", nil)
// Should fallback to builtin (no error)
keywords, err := extractor.Extract(nil, "test content for keyword extraction", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords, "Should fallback to builtin and extract keywords")
}
func TestMCPProviderServerNotFound(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with non-existent MCP server
extractor := keyword.NewExtractor("mcp:nonexistent.extract_keywords", &searchTypes.KeywordConfig{})
_, err := extractor.Extract(ctx, "test content", nil)
assert.Error(t, err, "Should error for non-existent MCP server")
assert.Contains(t, err.Error(), "not found")
}
func TestMCPProviderToolNotFound(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with non-existent tool
extractor := keyword.NewExtractor("mcp:search.nonexistent_tool", &searchTypes.KeywordConfig{})
_, err := extractor.Extract(ctx, "test content", nil)
assert.Error(t, err, "Should error for non-existent MCP tool")
}
func TestMCPProviderEmptyContent(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with MCP mode
extractor := keyword.NewExtractor("mcp:search.extract_keywords", nil)
// Test with empty content - MCP tool should return error
_, err := extractor.Extract(ctx, "", nil)
assert.Error(t, err, "Should error for empty content")
}
// newMCPTestContext creates a test context for MCP tests
func newMCPTestContext(t *testing.T) *context.Context {
t.Helper()
authorized := &oauthTypes.AuthorizedInfo{
UserID: "test-user",
}
chatID := "test-chat-mcp-keyword"
ctx := context.New(t.Context(), authorized, chatID)
return ctx
}

102
agent/search/reference.go Normal file
View file

@ -0,0 +1,102 @@
package search
import (
"fmt"
"strings"
"github.com/yaoapp/yao/agent/search/types"
)
// DefaultCitationPrompt is the default prompt for citation instructions
const DefaultCitationPrompt = `You have access to reference data in <references> tags. Each <ref> has:
- id: Citation identifier
- type: Data type (web/kb/db)
- weight: Relevance weight (1.0=highest priority, 0.6=lowest)
- source: Origin (user=user-provided, hook=assistant-searched, auto=auto-searched)
Prioritize higher-weight references when answering.
When citing a reference, use this exact HTML format:
<a class="ref" data-ref-id="{id}" data-ref-type="{type}" href="#ref:{id}">[{id}]</a>
Example: According to the product data<a class="ref" data-ref-id="ref_001" data-ref-type="db" href="#ref:ref_001">[ref_001]</a>, the price is $999.`
// BuildReferences converts search results to unified Reference format
func BuildReferences(results []*types.Result) []*types.Reference {
var refs []*types.Reference
for _, result := range results {
if result == nil {
continue
}
for _, item := range result.Items {
if item == nil {
continue
}
refs = append(refs, &types.Reference{
ID: item.CitationID,
Type: item.Type,
Source: item.Source,
Weight: item.Weight,
Score: item.Score,
Title: item.Title,
Content: item.Content,
URL: item.URL,
})
}
}
return refs
}
// FormatReferencesXML formats references as XML for LLM context
func FormatReferencesXML(refs []*types.Reference) string {
if len(refs) == 0 {
return ""
}
var sb strings.Builder
sb.WriteString("<references>\n")
for _, ref := range refs {
if ref == nil {
continue
}
sb.WriteString(fmt.Sprintf(`<ref id="%s" type="%s" weight="%.1f" source="%s">`,
ref.ID, ref.Type, ref.Weight, ref.Source))
sb.WriteString("\n")
if ref.Title != "" {
sb.WriteString(ref.Title)
sb.WriteString("\n")
}
sb.WriteString(ref.Content)
if ref.URL != "" {
sb.WriteString("\nURL: ")
sb.WriteString(ref.URL)
}
sb.WriteString("\n</ref>\n")
}
sb.WriteString("</references>")
return sb.String()
}
// GetCitationPrompt returns the citation instruction prompt
func GetCitationPrompt(cfg *types.CitationConfig) string {
if cfg == nil {
return DefaultCitationPrompt
}
if cfg.CustomPrompt != "" {
return cfg.CustomPrompt
}
return DefaultCitationPrompt
}
// BuildReferenceContext builds the complete reference context for LLM
func BuildReferenceContext(results []*types.Result, cfg *types.CitationConfig) *types.ReferenceContext {
refs := BuildReferences(results)
return &types.ReferenceContext{
References: refs,
XML: FormatReferencesXML(refs),
Prompt: GetCitationPrompt(cfg),
}
}

484
agent/search/reference_test.go Normal file
View file

@ -0,0 +1,484 @@
package search
import (
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/types"
)
func TestBuildReferences(t *testing.T) {
tests := []struct {
name string
results []*types.Result
expected int
}{
{
name: "nil results",
results: nil,
expected: 0,
},
{
name: "empty results",
results: []*types.Result{},
expected: 0,
},
{
name: "single result with items",
results: []*types.Result{
{
Type: types.SearchTypeWeb,
Query: "test query",
Items: []*types.ResultItem{
{
CitationID: "ref_001",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Weight: 0.6,
Score: 0.9,
Title: "Test Title",
Content: "Test content",
URL: "https://example.com",
},
{
CitationID: "ref_002",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Weight: 0.6,
Score: 0.8,
Title: "Test Title 2",
Content: "Test content 2",
URL: "https://example2.com",
},
},
},
},
expected: 2,
},
{
name: "multiple results",
results: []*types.Result{
{
Type: types.SearchTypeWeb,
Items: []*types.ResultItem{
{CitationID: "ref_001", Type: types.SearchTypeWeb, Content: "Web content"},
},
},
{
Type: types.SearchTypeKB,
Items: []*types.ResultItem{
{CitationID: "ref_002", Type: types.SearchTypeKB, Content: "KB content"},
},
},
{
Type: types.SearchTypeDB,
Items: []*types.ResultItem{
{CitationID: "ref_003", Type: types.SearchTypeDB, Content: "DB content"},
},
},
},
expected: 3,
},
{
name: "result with nil items",
results: []*types.Result{
{
Type: types.SearchTypeWeb,
Items: []*types.ResultItem{
{CitationID: "ref_001", Content: "Content 1"},
nil,
{CitationID: "ref_002", Content: "Content 2"},
},
},
},
expected: 2,
},
{
name: "nil result in slice",
results: []*types.Result{
{
Type: types.SearchTypeWeb,
Items: []*types.ResultItem{
{CitationID: "ref_001", Content: "Content"},
},
},
nil,
{
Type: types.SearchTypeKB,
Items: []*types.ResultItem{
{CitationID: "ref_002", Content: "Content 2"},
},
},
},
expected: 2,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
refs := BuildReferences(tt.results)
assert.Equal(t, tt.expected, len(refs))
})
}
}
func TestBuildReferences_FieldMapping(t *testing.T) {
item := &types.ResultItem{
CitationID: "ref_001",
Type: types.SearchTypeWeb,
Source: types.SourceHook,
Weight: 0.8,
Score: 0.95,
Title: "Test Title",
Content: "Test Content",
URL: "https://example.com",
}
results := []*types.Result{
{Items: []*types.ResultItem{item}},
}
refs := BuildReferences(results)
assert.Equal(t, 1, len(refs))
ref := refs[0]
assert.Equal(t, "ref_001", ref.ID)
assert.Equal(t, types.SearchTypeWeb, ref.Type)
assert.Equal(t, types.SourceHook, ref.Source)
assert.Equal(t, 0.8, ref.Weight)
assert.Equal(t, 0.95, ref.Score)
assert.Equal(t, "Test Title", ref.Title)
assert.Equal(t, "Test Content", ref.Content)
assert.Equal(t, "https://example.com", ref.URL)
}
func TestFormatReferencesXML(t *testing.T) {
tests := []struct {
name string
refs []*types.Reference
contains []string
excludes []string
}{
{
name: "nil refs",
refs: nil,
contains: []string{},
excludes: []string{"<references>"},
},
{
name: "empty refs",
refs: []*types.Reference{},
contains: []string{},
excludes: []string{"<references>"},
},
{
name: "single ref with all fields",
refs: []*types.Reference{
{
ID: "ref_001",
Type: types.SearchTypeWeb,
Source: types.SourceUser,
Weight: 1.0,
Score: 0.9,
Title: "Test Title",
Content: "Test Content",
URL: "https://example.com",
},
},
contains: []string{
"<references>",
"</references>",
`<ref id="ref_001" type="web" weight="1.0" source="user">`,
"</ref>",
"Test Title",
"Test Content",
"URL: https://example.com",
},
},
{
name: "ref without title",
refs: []*types.Reference{
{
ID: "ref_001",
Type: types.SearchTypeKB,
Source: types.SourceHook,
Weight: 0.8,
Content: "Content without title",
},
},
contains: []string{
`<ref id="ref_001" type="kb" weight="0.8" source="hook">`,
"Content without title",
},
excludes: []string{
"URL:",
},
},
{
name: "ref without URL",
refs: []*types.Reference{
{
ID: "ref_001",
Type: types.SearchTypeDB,
Source: types.SourceAuto,
Weight: 0.6,
Title: "DB Record",
Content: "Database content",
},
},
contains: []string{
`<ref id="ref_001" type="db" weight="0.6" source="auto">`,
"DB Record",
"Database content",
},
excludes: []string{
"URL:",
},
},
{
name: "multiple refs",
refs: []*types.Reference{
{ID: "ref_001", Type: types.SearchTypeWeb, Source: types.SourceUser, Weight: 1.0, Content: "Content 1"},
{ID: "ref_002", Type: types.SearchTypeKB, Source: types.SourceHook, Weight: 0.8, Content: "Content 2"},
{ID: "ref_003", Type: types.SearchTypeDB, Source: types.SourceAuto, Weight: 0.6, Content: "Content 3"},
},
contains: []string{
"<references>",
"</references>",
`id="ref_001"`,
`id="ref_002"`,
`id="ref_003"`,
"Content 1",
"Content 2",
"Content 3",
},
},
{
name: "nil ref in slice",
refs: []*types.Reference{
{ID: "ref_001", Type: types.SearchTypeWeb, Weight: 1.0, Content: "Content 1"},
nil,
{ID: "ref_002", Type: types.SearchTypeKB, Weight: 0.8, Content: "Content 2"},
},
contains: []string{
`id="ref_001"`,
`id="ref_002"`,
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
xml := FormatReferencesXML(tt.refs)
for _, s := range tt.contains {
assert.Contains(t, xml, s, "expected XML to contain: %s", s)
}
for _, s := range tt.excludes {
assert.NotContains(t, xml, s, "expected XML to not contain: %s", s)
}
})
}
}
func TestFormatReferencesXML_Structure(t *testing.T) {
refs := []*types.Reference{
{
ID: "ref_001",
Type: types.SearchTypeWeb,
Source: types.SourceUser,
Weight: 1.0,
Title: "Title",
Content: "Content",
URL: "https://example.com",
},
}
xml := FormatReferencesXML(refs)
// Check structure
assert.True(t, strings.HasPrefix(xml, "<references>\n"))
assert.True(t, strings.HasSuffix(xml, "</references>"))
assert.Contains(t, xml, "</ref>\n")
}
func TestGetCitationPrompt(t *testing.T) {
tests := []struct {
name string
cfg *types.CitationConfig
expected string
}{
{
name: "nil config",
cfg: nil,
expected: DefaultCitationPrompt,
},
{
name: "empty config",
cfg: &types.CitationConfig{},
expected: DefaultCitationPrompt,
},
{
name: "config with custom prompt",
cfg: &types.CitationConfig{
CustomPrompt: "Custom citation instructions",
},
expected: "Custom citation instructions",
},
{
name: "config with empty custom prompt",
cfg: &types.CitationConfig{
CustomPrompt: "",
},
expected: DefaultCitationPrompt,
},
{
name: "config with format but no custom prompt",
cfg: &types.CitationConfig{
Format: "[{id}]",
},
expected: DefaultCitationPrompt,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
prompt := GetCitationPrompt(tt.cfg)
assert.Equal(t, tt.expected, prompt)
})
}
}
func TestDefaultCitationPrompt(t *testing.T) {
// Verify default prompt contains key instructions
assert.Contains(t, DefaultCitationPrompt, "<references>")
assert.Contains(t, DefaultCitationPrompt, "id: Citation identifier")
assert.Contains(t, DefaultCitationPrompt, "type: Data type")
assert.Contains(t, DefaultCitationPrompt, "weight: Relevance weight")
assert.Contains(t, DefaultCitationPrompt, "source: Origin")
assert.Contains(t, DefaultCitationPrompt, `<a class="ref"`)
assert.Contains(t, DefaultCitationPrompt, "data-ref-id")
assert.Contains(t, DefaultCitationPrompt, "data-ref-type")
}
func TestBuildReferenceContext(t *testing.T) {
results := []*types.Result{
{
Type: types.SearchTypeWeb,
Items: []*types.ResultItem{
{
CitationID: "ref_001",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Weight: 0.6,
Title: "Test",
Content: "Content",
URL: "https://example.com",
},
},
},
}
t.Run("with nil config", func(t *testing.T) {
ctx := BuildReferenceContext(results, nil)
assert.NotNil(t, ctx)
assert.Equal(t, 1, len(ctx.References))
assert.Contains(t, ctx.XML, "<references>")
assert.Contains(t, ctx.XML, "ref_001")
assert.Equal(t, DefaultCitationPrompt, ctx.Prompt)
})
t.Run("with custom prompt config", func(t *testing.T) {
cfg := &types.CitationConfig{
CustomPrompt: "Custom prompt",
}
ctx := BuildReferenceContext(results, cfg)
assert.NotNil(t, ctx)
assert.Equal(t, "Custom prompt", ctx.Prompt)
})
t.Run("with empty results", func(t *testing.T) {
ctx := BuildReferenceContext([]*types.Result{}, nil)
assert.NotNil(t, ctx)
assert.Equal(t, 0, len(ctx.References))
assert.Equal(t, "", ctx.XML)
assert.Equal(t, DefaultCitationPrompt, ctx.Prompt)
})
}
func TestBuildReferenceContext_Integration(t *testing.T) {
// Simulate a real-world scenario with multiple search types
results := []*types.Result{
{
Type: types.SearchTypeWeb,
Query: "AI developments",
Items: []*types.ResultItem{
{
CitationID: "ref_001",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Weight: 0.6,
Score: 0.95,
Title: "OpenAI Announces GPT-5",
Content: "OpenAI has announced the development of GPT-5...",
URL: "https://news.example.com/gpt5",
},
},
},
{
Type: types.SearchTypeKB,
Query: "AI developments",
Items: []*types.ResultItem{
{
CitationID: "ref_002",
Type: types.SearchTypeKB,
Source: types.SourceHook,
Weight: 0.8,
Score: 0.88,
Title: "Internal AI Research Notes",
Content: "Our internal research on AI capabilities...",
},
},
},
{
Type: types.SearchTypeDB,
Query: "AI developments",
Items: []*types.ResultItem{
{
CitationID: "ref_003",
Type: types.SearchTypeDB,
Source: types.SourceUser,
Weight: 1.0,
Score: 0.92,
Title: "Product: AI Assistant",
Content: "Name: AI Assistant\nPrice: $99\nCategory: Software",
},
},
},
}
ctx := BuildReferenceContext(results, nil)
// Verify all references are included
assert.Equal(t, 3, len(ctx.References))
// Verify XML contains all references
assert.Contains(t, ctx.XML, "ref_001")
assert.Contains(t, ctx.XML, "ref_002")
assert.Contains(t, ctx.XML, "ref_003")
// Verify different source types are represented
assert.Contains(t, ctx.XML, `source="auto"`)
assert.Contains(t, ctx.XML, `source="hook"`)
assert.Contains(t, ctx.XML, `source="user"`)
// Verify different search types are represented
assert.Contains(t, ctx.XML, `type="web"`)
assert.Contains(t, ctx.XML, `type="kb"`)
assert.Contains(t, ctx.XML, `type="db"`)
}

29
agent/search/registry.go Normal file
View file

@ -0,0 +1,29 @@
package search
import (
"github.com/yaoapp/yao/agent/search/interfaces"
"github.com/yaoapp/yao/agent/search/types"
)
// Registry manages search handlers
type Registry struct {
handlers map[types.SearchType]interfaces.Handler
}
// NewRegistry creates a new handler registry
func NewRegistry() *Registry {
return &Registry{
handlers: make(map[types.SearchType]interfaces.Handler),
}
}
// Register registers a handler for a search type
func (r *Registry) Register(handler interfaces.Handler) {
r.handlers[handler.Type()] = handler
}
// Get returns the handler for a search type
func (r *Registry) Get(t types.SearchType) (interfaces.Handler, bool) {
h, ok := r.handlers[t]
return h, ok
}

77
agent/search/registry_test.go Normal file
View file

@ -0,0 +1,77 @@
package search
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/handlers/db"
"github.com/yaoapp/yao/agent/search/handlers/kb"
"github.com/yaoapp/yao/agent/search/handlers/web"
"github.com/yaoapp/yao/agent/search/types"
)
func TestNewRegistry(t *testing.T) {
r := NewRegistry()
assert.NotNil(t, r)
assert.NotNil(t, r.handlers)
assert.Equal(t, 0, len(r.handlers))
}
func TestRegistry_Register(t *testing.T) {
r := NewRegistry()
// Register web handler
webHandler := web.NewHandler("builtin", nil)
r.Register(webHandler)
h, ok := r.Get(types.SearchTypeWeb)
assert.True(t, ok)
assert.Equal(t, types.SearchTypeWeb, h.Type())
}
func TestRegistry_RegisterMultiple(t *testing.T) {
r := NewRegistry()
// Register all handlers
r.Register(web.NewHandler("builtin", nil))
r.Register(kb.NewHandler(nil))
r.Register(db.NewHandler("builtin", nil))
// Verify all are registered
webH, ok := r.Get(types.SearchTypeWeb)
assert.True(t, ok)
assert.Equal(t, types.SearchTypeWeb, webH.Type())
kbH, ok := r.Get(types.SearchTypeKB)
assert.True(t, ok)
assert.Equal(t, types.SearchTypeKB, kbH.Type())
dbH, ok := r.Get(types.SearchTypeDB)
assert.True(t, ok)
assert.Equal(t, types.SearchTypeDB, dbH.Type())
}
func TestRegistry_Get_NotFound(t *testing.T) {
r := NewRegistry()
h, ok := r.Get(types.SearchTypeWeb)
assert.False(t, ok)
assert.Nil(t, h)
}
func TestRegistry_RegisterOverwrite(t *testing.T) {
r := NewRegistry()
// Register first handler
h1 := web.NewHandler("builtin", nil)
r.Register(h1)
// Register second handler (same type)
h2 := web.NewHandler("agent", nil)
r.Register(h2)
// Should get the second handler
h, ok := r.Get(types.SearchTypeWeb)
assert.True(t, ok)
assert.NotNil(t, h)
}

232
agent/search/rerank/agent.go Normal file
View file

@ -0,0 +1,232 @@
package rerank
import (
"encoding/json"
"fmt"
"strings"
"github.com/yaoapp/yao/agent/caller"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// AgentProvider implements reranking by delegating to another agent
// The agent should have a Next Hook that accepts rerank request and returns reordered items
type AgentProvider struct {
agentID string // Assistant ID to delegate to
}
// NewAgentProvider creates a new agent reranker
func NewAgentProvider(agentID string) *AgentProvider {
return &AgentProvider{agentID: agentID}
}
// Rerank delegates reranking to an LLM-powered assistant
// The assistant receives items and query, returns reordered item IDs or items
func (p *AgentProvider) Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
if ctx == nil {
return nil, fmt.Errorf("context is required for agent rerank")
}
// Get agent via caller interface (avoids circular dependency)
agent, err := caller.AgentGetterFunc(p.agentID)
if err != nil {
return nil, fmt.Errorf("failed to get agent %s: %w", p.agentID, err)
}
// Build request message with items to rerank
requestData := map[string]interface{}{
"query": query,
"items": items,
"top_n": opts.TopN,
"action": "rerank",
}
requestJSON, _ := json.Marshal(requestData)
// Create messages for agent
messages := []context.Message{
{
Role: "user",
Content: string(requestJSON),
},
}
// Call agent's Stream method with skip options (no history, no output)
options := &context.Options{
Skip: &context.Skip{
History: true,
Output: true,
},
}
result, err := agent.Stream(ctx, messages, options)
if err != nil {
return nil, fmt.Errorf("agent stream failed: %w", err)
}
// Parse response
return p.parseResponse(result, items, opts)
}
// parseResponse extracts reranked items from agent response
// The response format from agent.Stream is typically:
// { "next": { "data": { "order": [...] } } }
func (p *AgentProvider) parseResponse(result interface{}, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
if result == nil {
return originalItems, nil
}
// Build index map for quick lookup
itemMap := make(map[string]*types.ResultItem)
for _, item := range originalItems {
if item.CitationID != "" {
itemMap[item.CitationID] = item
}
}
// Extract response data
response := extractResponseData(result)
if response == nil {
return originalItems, nil
}
// Try to get reranked order from response
// Expected format: { "order": ["ref_001", "ref_003", "ref_002"] }
// Or: { "items": [{ "citation_id": "ref_001", ... }, ...] }
var reranked []*types.ResultItem
// Try "order" field (list of citation IDs)
if order, ok := response["order"]; ok {
if orderList := toStringSlice(order); len(orderList) > 0 {
for _, id := range orderList {
if item, exists := itemMap[id]; exists {
reranked = append(reranked, item)
delete(itemMap, id) // Avoid duplicates
}
}
// Append remaining items not in order
for _, item := range originalItems {
if _, exists := itemMap[item.CitationID]; exists {
reranked = append(reranked, item)
}
}
}
}
// Try "items" field (full items or items with citation_id)
if len(reranked) == 0 {
if items, ok := response["items"]; ok {
if itemsList := toItemsList(items); len(itemsList) > 0 {
for _, respItem := range itemsList {
// Check if it's just a reference or full item
if citationID, ok := respItem["citation_id"].(string); ok {
if item, exists := itemMap[citationID]; exists {
reranked = append(reranked, item)
delete(itemMap, citationID)
}
}
}
// Append remaining items
for _, item := range originalItems {
if _, exists := itemMap[item.CitationID]; exists {
reranked = append(reranked, item)
}
}
}
}
}
// If no valid response, return original items
if len(reranked) == 0 {
reranked = originalItems
}
// Apply top N
if opts.TopN > 0 && opts.TopN < len(reranked) {
reranked = reranked[:opts.TopN]
}
return reranked, nil
}
// extractResponseData extracts the actual response data from agent.Stream result
// Handles nested structures like { "next": { "data": { ... } } }
func extractResponseData(result interface{}) map[string]interface{} {
switch v := result.(type) {
case map[string]interface{}:
// Check for "next" wrapper (from NextHookResponse)
if next, ok := v["next"].(map[string]interface{}); ok {
// Check for "data" inside next
if data, ok := next["data"].(map[string]interface{}); ok {
return data
}
return next
}
// Check for direct "data" wrapper
if data, ok := v["data"].(map[string]interface{}); ok {
return data
}
return v
case string:
// Try to parse as JSON
var data map[string]interface{}
if err := json.Unmarshal([]byte(v), &data); err == nil {
return extractResponseData(data)
}
}
// Try to handle other types by converting to JSON and back
if result != nil {
if bytes, err := json.Marshal(result); err == nil {
var data map[string]interface{}
if err := json.Unmarshal(bytes, &data); err == nil {
return extractResponseData(data)
}
}
}
return nil
}
// toStringSlice converts interface to string slice
func toStringSlice(v interface{}) []string {
switch val := v.(type) {
case []string:
return val
case []interface{}:
result := make([]string, 0, len(val))
for _, item := range val {
if s, ok := item.(string); ok {
result = append(result, s)
}
}
return result
}
return nil
}
// toItemsList converts interface to list of maps
func toItemsList(v interface{}) []map[string]interface{} {
switch val := v.(type) {
case []map[string]interface{}:
return val
case []interface{}:
result := make([]map[string]interface{}, 0, len(val))
for _, item := range val {
if m, ok := item.(map[string]interface{}); ok {
result = append(result, m)
}
}
return result
}
return nil
}
// extractAgentID extracts assistant ID from uses.rerank value
// For backward compatibility, strips any prefix if present
func extractAgentID(usesRerank string) string {
// Remove any prefix like "agent:" if present
if strings.HasPrefix(usesRerank, "agent:") {
return strings.TrimPrefix(usesRerank, "agent:")
}
return usesRerank
}

123
agent/search/rerank/agent_test.go Normal file
View file

@ -0,0 +1,123 @@
package rerank_test
import (
"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/search/rerank"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
)
func TestAgentProviderWithAssistantConfig(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the rerank-agent assistant
ast, err := assistant.Get("tests.rerank-agent")
require.NoError(t, err)
require.NotNil(t, ast)
// Create test context
ctx := newTestContext(t)
// Create provider with test assistant
provider := rerank.NewAgentProvider("tests.rerank-agent")
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0, Title: "First"},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0, Title: "Second"},
{CitationID: "ref_003", Score: 0.7, Weight: 1.0, Title: "Third"},
}
result, err := provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 10})
require.NoError(t, err)
assert.NotEmpty(t, result)
// The mock agent reverses the order
// So we expect: ref_003, ref_002, ref_001
assert.Len(t, result, 3)
assert.Equal(t, "ref_003", result[0].CitationID)
assert.Equal(t, "ref_002", result[1].CitationID)
assert.Equal(t, "ref_001", result[2].CitationID)
}
func TestAgentProviderWithTopN(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ctx := newTestContext(t)
provider := rerank.NewAgentProvider("tests.rerank-agent")
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
{CitationID: "ref_003", Score: 0.7, Weight: 1.0},
{CitationID: "ref_004", Score: 0.6, Weight: 1.0},
{CitationID: "ref_005", Score: 0.5, Weight: 1.0},
}
// Request top 2 only
result, err := provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 2})
require.NoError(t, err)
assert.Len(t, result, 2)
}
func TestAgentProviderWithoutContext(t *testing.T) {
provider := rerank.NewAgentProvider("tests.rerank-agent")
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
}
_, err := provider.Rerank(nil, "test query", items, &types.RerankOptions{TopN: 10})
assert.Error(t, err)
assert.Contains(t, err.Error(), "context is required")
}
func TestAgentProviderAgentNotFound(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ctx := newTestContext(t)
provider := rerank.NewAgentProvider("non-existent-agent")
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
}
_, err := provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 10})
assert.Error(t, err)
assert.Contains(t, err.Error(), "failed to get agent")
}
func TestAgentProviderEmptyItems(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ctx := newTestContext(t)
provider := rerank.NewAgentProvider("tests.rerank-agent")
result, err := provider.Rerank(ctx, "test query", []*types.ResultItem{}, &types.RerankOptions{TopN: 10})
require.NoError(t, err)
assert.Empty(t, result)
}
// newTestContext creates a test context with required fields
func newTestContext(t *testing.T) *context.Context {
t.Helper()
authorized := &oauthTypes.AuthorizedInfo{
UserID: "test-user",
}
chatID := "test-chat-rerank"
return context.New(t.Context(), authorized, chatID)
}

62
agent/search/rerank/builtin.go Normal file
View file

@ -0,0 +1,62 @@
package rerank
import (
"sort"
"github.com/yaoapp/yao/agent/search/types"
)
// BuiltinReranker implements simple score-based reranking
// For production use cases requiring semantic understanding, use Agent or MCP mode.
type BuiltinReranker struct{}
// NewBuiltinReranker creates a new builtin reranker
func NewBuiltinReranker() *BuiltinReranker {
return &BuiltinReranker{}
}
// Rerank sorts items by weighted score (score * weight) and returns top N
// This is a simple implementation without semantic understanding.
func (r *BuiltinReranker) Rerank(query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
if len(items) == 0 {
return items, nil
}
// Calculate weighted scores
type scoredItem struct {
item *types.ResultItem
weightedScore float64
}
scored := make([]scoredItem, len(items))
for i, item := range items {
// Weighted score = base score * source weight
// Higher weight sources (user=1.0) get priority over lower (auto=0.6)
weight := item.Weight
if weight == 0 {
weight = 0.6 // Default weight for items without weight
}
scored[i] = scoredItem{
item: item,
weightedScore: item.Score * weight,
}
}
// Sort by weighted score descending
sort.Slice(scored, func(i, j int) bool {
return scored[i].weightedScore > scored[j].weightedScore
})
// Get top N
topN := opts.TopN
if topN <= 0 || topN > len(scored) {
topN = len(scored)
}
result := make([]*types.ResultItem, topN)
for i := 0; i < topN; i++ {
result[i] = scored[i].item
}
return result, nil
}

124
agent/search/rerank/builtin_test.go Normal file
View file

@ -0,0 +1,124 @@
package rerank
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/types"
)
func TestBuiltinReranker_EmptyItems(t *testing.T) {
reranker := NewBuiltinReranker()
result, err := reranker.Rerank("test query", []*types.ResultItem{}, &types.RerankOptions{TopN: 5})
assert.NoError(t, err)
assert.Empty(t, result)
}
func TestBuiltinReranker_SortByWeightedScore(t *testing.T) {
reranker := NewBuiltinReranker()
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.8, Weight: 0.6}, // weighted: 0.48
{CitationID: "ref_002", Score: 0.6, Weight: 1.0}, // weighted: 0.60
{CitationID: "ref_003", Score: 0.9, Weight: 0.8}, // weighted: 0.72
{CitationID: "ref_004", Score: 0.5, Weight: 1.0}, // weighted: 0.50
}
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 10})
assert.NoError(t, err)
assert.Len(t, result, 4)
// Should be sorted by weighted score: ref_003 (0.72) > ref_002 (0.60) > ref_004 (0.50) > ref_001 (0.48)
assert.Equal(t, "ref_003", result[0].CitationID)
assert.Equal(t, "ref_002", result[1].CitationID)
assert.Equal(t, "ref_004", result[2].CitationID)
assert.Equal(t, "ref_001", result[3].CitationID)
}
func TestBuiltinReranker_TopN(t *testing.T) {
reranker := NewBuiltinReranker()
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
{CitationID: "ref_003", Score: 0.7, Weight: 1.0},
{CitationID: "ref_004", Score: 0.6, Weight: 1.0},
{CitationID: "ref_005", Score: 0.5, Weight: 1.0},
}
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 3})
assert.NoError(t, err)
assert.Len(t, result, 3)
assert.Equal(t, "ref_001", result[0].CitationID)
assert.Equal(t, "ref_002", result[1].CitationID)
assert.Equal(t, "ref_003", result[2].CitationID)
}
func TestBuiltinReranker_DefaultWeight(t *testing.T) {
reranker := NewBuiltinReranker()
// Items without weight should use default 0.6
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 0}, // weighted: 0.9 * 0.6 = 0.54
{CitationID: "ref_002", Score: 0.5, Weight: 1.0}, // weighted: 0.5 * 1.0 = 0.50
}
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 10})
assert.NoError(t, err)
assert.Len(t, result, 2)
// ref_001 (0.54) > ref_002 (0.50)
assert.Equal(t, "ref_001", result[0].CitationID)
assert.Equal(t, "ref_002", result[1].CitationID)
}
func TestBuiltinReranker_TopNLargerThanItems(t *testing.T) {
reranker := NewBuiltinReranker()
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
}
// TopN > len(items) should return all items
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 10})
assert.NoError(t, err)
assert.Len(t, result, 2)
}
func TestBuiltinReranker_ZeroTopN(t *testing.T) {
reranker := NewBuiltinReranker()
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
}
// TopN = 0 should return all items
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 0})
assert.NoError(t, err)
assert.Len(t, result, 2)
}
func TestBuiltinReranker_SameWeightedScore(t *testing.T) {
reranker := NewBuiltinReranker()
// Items with same weighted score - order should be stable
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.8, Weight: 1.0}, // weighted: 0.80
{CitationID: "ref_002", Score: 0.8, Weight: 1.0}, // weighted: 0.80
{CitationID: "ref_003", Score: 0.4, Weight: 1.0}, // weighted: 0.40
}
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 10})
assert.NoError(t, err)
assert.Len(t, result, 3)
// ref_003 should be last
assert.Equal(t, "ref_003", result[2].CitationID)
}

171
agent/search/rerank/mcp.go Normal file
View file

@ -0,0 +1,171 @@
package rerank
import (
"encoding/json"
"fmt"
"strings"
"github.com/yaoapp/gou/mcp"
gouMCPTypes "github.com/yaoapp/gou/mcp/types"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// MCPProvider implements reranking by calling an MCP tool
type MCPProvider struct {
serverID string // MCP server ID
toolName string // Tool name
}
// NewMCPProvider creates a new MCP reranker
// mcpRef format: "server_id.tool_name"
func NewMCPProvider(mcpRef string) (*MCPProvider, error) {
parts := strings.SplitN(mcpRef, ".", 2)
if len(parts) != 2 {
return nil, fmt.Errorf("invalid MCP format, expected 'server.tool', got '%s'", mcpRef)
}
return &MCPProvider{
serverID: parts[0],
toolName: parts[1],
}, nil
}
// Rerank calls MCP tool to rerank items
func (p *MCPProvider) Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
if ctx == nil {
return nil, fmt.Errorf("context is required for MCP rerank")
}
// Get MCP client
client, err := mcp.Select(p.serverID)
if err != nil {
return nil, fmt.Errorf("MCP server %s not found: %w", p.serverID, err)
}
// Build arguments for MCP tool
args := map[string]interface{}{
"query": query,
"items": items,
"top_n": opts.TopN,
}
// Call MCP tool
result, err := client.CallTool(ctx.Context, p.toolName, args)
if err != nil {
return nil, fmt.Errorf("MCP tool call failed: %w", err)
}
// Parse result
return p.parseResult(result, items, opts)
}
// parseResult extracts reranked items from MCP response
func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
if result == nil || len(result.Content) == 0 {
return originalItems, nil
}
// Build index map for quick lookup
itemMap := make(map[string]*types.ResultItem)
for _, item := range originalItems {
if item.CitationID != "" {
itemMap[item.CitationID] = item
}
}
// Extract text content from MCP response
var textContent string
for _, content := range result.Content {
if content.Type == gouMCPTypes.ToolContentTypeText && content.Text != "" {
textContent = content.Text
break
}
}
if textContent == "" {
return originalItems, nil
}
// Parse JSON response
var response map[string]interface{}
if err := json.Unmarshal([]byte(textContent), &response); err != nil {
// Try parsing as array of IDs
var orderList []string
if err := json.Unmarshal([]byte(textContent), &orderList); err == nil {
return p.reorderByIDs(orderList, itemMap, originalItems, opts)
}
return originalItems, nil
}
// Try "order" field (list of citation IDs)
if order, ok := response["order"]; ok {
if orderList := toStringSlice(order); len(orderList) > 0 {
return p.reorderByIDs(orderList, itemMap, originalItems, opts)
}
}
// Try "items" field
if items, ok := response["items"]; ok {
if itemsList := toItemsList(items); len(itemsList) > 0 {
return p.reorderByItems(itemsList, itemMap, originalItems, opts)
}
}
return originalItems, nil
}
// reorderByIDs reorders items based on list of citation IDs
func (p *MCPProvider) reorderByIDs(order []string, itemMap map[string]*types.ResultItem, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
var result []*types.ResultItem
// Add items in specified order
for _, id := range order {
if item, exists := itemMap[id]; exists {
result = append(result, item)
delete(itemMap, id)
}
}
// Append remaining items
for _, item := range originalItems {
if _, exists := itemMap[item.CitationID]; exists {
result = append(result, item)
}
}
// Apply top N
if opts.TopN > 0 && opts.TopN < len(result) {
result = result[:opts.TopN]
}
return result, nil
}
// reorderByItems reorders items based on list of item references
func (p *MCPProvider) reorderByItems(itemsList []map[string]interface{}, itemMap map[string]*types.ResultItem, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
var result []*types.ResultItem
// Add items in specified order
for _, respItem := range itemsList {
if citationID, ok := respItem["citation_id"].(string); ok {
if item, exists := itemMap[citationID]; exists {
result = append(result, item)
delete(itemMap, citationID)
}
}
}
// Append remaining items
for _, item := range originalItems {
if _, exists := itemMap[item.CitationID]; exists {
result = append(result, item)
}
}
// Apply top N
if opts.TopN > 0 && opts.TopN < len(result) {
result = result[:opts.TopN]
}
return result, nil
}

148
agent/search/rerank/mcp_test.go Normal file
View file

@ -0,0 +1,148 @@
package rerank_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/rerank"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
)
func TestMCPProviderWithSearchRerank(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ctx := newMCPTestContext(t)
provider, err := rerank.NewMCPProvider("search.rerank")
require.NoError(t, err)
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0, Title: "First"},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0, Title: "Second"},
{CitationID: "ref_003", Score: 0.7, Weight: 1.0, Title: "Third"},
}
result, err := provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 10})
require.NoError(t, err)
assert.NotEmpty(t, result)
// The mock MCP reverses the order
// So we expect: ref_003, ref_002, ref_001
assert.Len(t, result, 3)
assert.Equal(t, "ref_003", result[0].CitationID)
assert.Equal(t, "ref_002", result[1].CitationID)
assert.Equal(t, "ref_001", result[2].CitationID)
}
func TestMCPProviderWithTopN(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ctx := newMCPTestContext(t)
provider, err := rerank.NewMCPProvider("search.rerank")
require.NoError(t, err)
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
{CitationID: "ref_003", Score: 0.7, Weight: 1.0},
{CitationID: "ref_004", Score: 0.6, Weight: 1.0},
{CitationID: "ref_005", Score: 0.5, Weight: 1.0},
}
// Request top 2 only
result, err := provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 2})
require.NoError(t, err)
assert.Len(t, result, 2)
}
func TestMCPProviderInvalidFormat(t *testing.T) {
_, err := rerank.NewMCPProvider("invalid-format")
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid MCP format")
}
func TestMCPProviderServerNotFound(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ctx := newMCPTestContext(t)
provider, err := rerank.NewMCPProvider("nonexistent.rerank")
require.NoError(t, err)
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
}
_, err = provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 10})
assert.Error(t, err)
assert.Contains(t, err.Error(), "not found")
}
func TestMCPProviderToolNotFound(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ctx := newMCPTestContext(t)
provider, err := rerank.NewMCPProvider("search.nonexistent_tool")
require.NoError(t, err)
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
}
_, err = provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 10})
assert.Error(t, err)
}
func TestMCPProviderWithoutContext(t *testing.T) {
provider, err := rerank.NewMCPProvider("search.rerank")
require.NoError(t, err)
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
}
_, err = provider.Rerank(nil, "test query", items, &types.RerankOptions{TopN: 10})
assert.Error(t, err)
assert.Contains(t, err.Error(), "context is required")
}
func TestMCPProviderEmptyItems(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ctx := newMCPTestContext(t)
provider, err := rerank.NewMCPProvider("search.rerank")
require.NoError(t, err)
result, err := provider.Rerank(ctx, "test query", []*types.ResultItem{}, &types.RerankOptions{TopN: 10})
require.NoError(t, err)
assert.Empty(t, result)
}
// newMCPTestContext creates a test context with required fields
func newMCPTestContext(t *testing.T) *context.Context {
t.Helper()
authorized := &oauthTypes.AuthorizedInfo{
UserID: "test-user",
}
chatID := "test-chat-rerank-mcp"
return context.New(t.Context(), authorized, chatID)
}

99
agent/search/rerank/reranker.go Normal file
View file

@ -0,0 +1,99 @@
// Package rerank provides result reranking for search module
// Supports three modes via uses.rerank configuration:
// - "builtin": Simple score-based sorting (no external dependencies)
// - "<assistant-id>": Delegate to an LLM-powered assistant for semantic reranking
// - "mcp:<server>.<tool>": Call external MCP tool
//
// For production use cases requiring high accuracy, use Agent or MCP mode.
package rerank
import (
"strings"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// Reranker reorders search results by relevance
// Mode is determined by uses.rerank configuration
type Reranker struct {
usesRerank string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
config *types.RerankConfig // Rerank options
}
// NewReranker creates a new reranker
// usesRerank: value from uses.rerank config
// cfg: rerank options from search config
func NewReranker(usesRerank string, cfg *types.RerankConfig) *Reranker {
return &Reranker{
usesRerank: usesRerank,
config: cfg,
}
}
// Rerank reorders results based on configured mode
// Returns reordered items, potentially truncated to top N
func (r *Reranker) Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
if len(items) == 0 {
return items, nil
}
// Merge options with config defaults
mergedOpts := r.mergeOptions(opts)
switch {
case r.usesRerank == "builtin" || r.usesRerank == "":
return r.builtinRerank(query, items, mergedOpts)
case strings.HasPrefix(r.usesRerank, "mcp:"):
return r.mcpRerank(ctx, query, items, mergedOpts)
default:
// Assume it's an assistant ID for Agent mode
return r.agentRerank(ctx, query, items, mergedOpts)
}
}
// mergeOptions merges runtime options with config defaults
func (r *Reranker) mergeOptions(opts *types.RerankOptions) *types.RerankOptions {
result := &types.RerankOptions{
TopN: 10, // default
}
// Apply config defaults
if r.config != nil {
if r.config.TopN > 0 {
result.TopN = r.config.TopN
}
}
// Apply runtime options (highest priority)
if opts != nil {
if opts.TopN > 0 {
result.TopN = opts.TopN
}
}
return result
}
// builtinRerank uses simple score-based sorting
func (r *Reranker) builtinRerank(query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
reranker := NewBuiltinReranker()
return reranker.Rerank(query, items, opts)
}
// agentRerank delegates to an LLM-powered assistant
func (r *Reranker) agentRerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
provider := NewAgentProvider(r.usesRerank)
return provider.Rerank(ctx, query, items, opts)
}
// mcpRerank calls an external MCP tool
func (r *Reranker) mcpRerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
mcpRef := strings.TrimPrefix(r.usesRerank, "mcp:")
provider, err := NewMCPProvider(mcpRef)
if err != nil {
// Fallback to builtin on invalid MCP format
return r.builtinRerank(query, items, opts)
}
return provider.Rerank(ctx, query, items, opts)
}

99
agent/search/rerank/reranker_test.go Normal file
View file

@ -0,0 +1,99 @@
package rerank
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/types"
)
func TestReranker_BuiltinMode(t *testing.T) {
reranker := NewReranker("builtin", &types.RerankConfig{TopN: 5})
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
{CitationID: "ref_003", Score: 0.7, Weight: 1.0},
}
result, err := reranker.Rerank(nil, "test query", items, nil)
assert.NoError(t, err)
assert.Len(t, result, 3)
assert.Equal(t, "ref_001", result[0].CitationID)
}
func TestReranker_EmptyUsesRerank(t *testing.T) {
// Empty usesRerank should use builtin
reranker := NewReranker("", &types.RerankConfig{TopN: 5})
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
}
result, err := reranker.Rerank(nil, "test query", items, nil)
assert.NoError(t, err)
assert.Len(t, result, 1)
}
func TestReranker_MergeOptions(t *testing.T) {
// Config sets TopN = 5
reranker := NewReranker("builtin", &types.RerankConfig{TopN: 5})
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
{CitationID: "ref_003", Score: 0.7, Weight: 1.0},
{CitationID: "ref_004", Score: 0.6, Weight: 1.0},
{CitationID: "ref_005", Score: 0.5, Weight: 1.0},
{CitationID: "ref_006", Score: 0.4, Weight: 1.0},
}
// Runtime opts override config
result, err := reranker.Rerank(nil, "test query", items, &types.RerankOptions{TopN: 3})
assert.NoError(t, err)
assert.Len(t, result, 3)
}
func TestReranker_ConfigTopN(t *testing.T) {
// Config sets TopN = 3
reranker := NewReranker("builtin", &types.RerankConfig{TopN: 3})
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
{CitationID: "ref_003", Score: 0.7, Weight: 1.0},
{CitationID: "ref_004", Score: 0.6, Weight: 1.0},
{CitationID: "ref_005", Score: 0.5, Weight: 1.0},
}
// No runtime opts, should use config TopN
result, err := reranker.Rerank(nil, "test query", items, nil)
assert.NoError(t, err)
assert.Len(t, result, 3)
}
func TestReranker_NilConfig(t *testing.T) {
reranker := NewReranker("builtin", nil)
items := []*types.ResultItem{
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
}
result, err := reranker.Rerank(nil, "test query", items, nil)
assert.NoError(t, err)
assert.Len(t, result, 1)
}
func TestReranker_EmptyItems(t *testing.T) {
reranker := NewReranker("builtin", &types.RerankConfig{TopN: 5})
result, err := reranker.Rerank(nil, "test query", []*types.ResultItem{}, nil)
assert.NoError(t, err)
assert.Empty(t, result)
}

232
agent/search/search.go Normal file
View file

@ -0,0 +1,232 @@
package search
import (
"sync"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/handlers/db"
"github.com/yaoapp/yao/agent/search/handlers/kb"
"github.com/yaoapp/yao/agent/search/handlers/web"
"github.com/yaoapp/yao/agent/search/interfaces"
"github.com/yaoapp/yao/agent/search/rerank"
"github.com/yaoapp/yao/agent/search/types"
)
// Searcher is the main search implementation
type Searcher struct {
config *types.Config // Merged config (global + assistant)
handlers map[types.SearchType]interfaces.Handler
reranker *rerank.Reranker
citation *CitationGenerator
}
// Uses contains the search-specific uses configuration
// These are extracted from context.Uses and search config
type Uses struct {
Search string // "builtin", "disabled", "<assistant-id>", "mcp:<server>.<tool>"
Web string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
Keyword string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
QueryDSL string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
Rerank string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
}
// New creates a new Searcher instance
// cfg: merged config from agent/load.go + assistant config
// uses: merged uses configuration (global → assistant → hook)
func New(cfg *types.Config, uses *Uses) *Searcher {
if uses == nil {
uses = &Uses{}
}
if cfg == nil {
cfg = &types.Config{}
}
return &Searcher{
config: cfg,
handlers: map[types.SearchType]interfaces.Handler{
types.SearchTypeWeb: web.NewHandler(uses.Web, cfg.Web),
types.SearchTypeKB: kb.NewHandler(cfg.KB),
types.SearchTypeDB: db.NewHandler(uses.QueryDSL, cfg.DB),
},
reranker: rerank.NewReranker(uses.Rerank, cfg.Rerank),
citation: NewCitationGenerator(),
}
}
// Search executes a single search request
func (s *Searcher) Search(ctx *context.Context, req *types.Request) (*types.Result, error) {
handler, ok := s.handlers[req.Type]
if !ok {
return &types.Result{Error: "unsupported search type"}, nil
}
// Execute search
result, err := handler.Search(req)
if err != nil {
return &types.Result{Error: err.Error()}, nil
}
// Assign weights based on source
for _, item := range result.Items {
item.Weight = s.config.GetWeight(req.Source)
}
// Rerank if requested
if req.Rerank != nil && s.reranker != nil {
result.Items, _ = s.reranker.Rerank(ctx, req.Query, result.Items, req.Rerank)
}
// Generate citation IDs
for _, item := range result.Items {
item.CitationID = s.citation.Next()
}
return result, nil
}
// All executes all searches and waits for all to complete (like Promise.all)
func (s *Searcher) All(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
if len(reqs) == 0 {
return []*types.Result{}, nil
}
return s.parallelAll(ctx, reqs)
}
// Any returns as soon as any search succeeds with results (like Promise.any)
func (s *Searcher) Any(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
if len(reqs) == 0 {
return []*types.Result{}, nil
}
return s.parallelAny(ctx, reqs)
}
// Race returns as soon as any search completes (like Promise.race)
func (s *Searcher) Race(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
if len(reqs) == 0 {
return []*types.Result{}, nil
}
return s.parallelRace(ctx, reqs)
}
// parallelAll executes all searches and waits for all to complete (like Promise.all)
func (s *Searcher) parallelAll(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
results := make([]*types.Result, len(reqs))
var wg sync.WaitGroup
var mu sync.Mutex
for i, req := range reqs {
wg.Add(1)
go func(idx int, r *types.Request) {
defer wg.Done()
result, _ := s.Search(ctx, r)
mu.Lock()
results[idx] = result
mu.Unlock()
}(i, req)
}
wg.Wait()
return results, nil
}
// parallelAny returns as soon as any search succeeds (has results) (like Promise.any)
// Other searches continue in background but results are discarded
func (s *Searcher) parallelAny(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
results := make([]*types.Result, len(reqs))
resultChan := make(chan struct {
idx int
result *types.Result
}, len(reqs))
var wg sync.WaitGroup
done := make(chan struct{})
for i, req := range reqs {
wg.Add(1)
go func(idx int, r *types.Request) {
defer wg.Done()
result, _ := s.Search(ctx, r)
select {
case resultChan <- struct {
idx int
result *types.Result
}{idx, result}:
case <-done:
// Already found a successful result, discard this one
}
}(i, req)
}
// Close channel when all goroutines complete
go func() {
wg.Wait()
close(resultChan)
}()
// Collect results until we find one with items (success)
var mu sync.Mutex
for res := range resultChan {
mu.Lock()
results[res.idx] = res.result
// Check if this result has items (success = has results and no error)
if res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" {
mu.Unlock()
close(done) // Signal other goroutines to stop sending
return results, nil
}
mu.Unlock()
}
// No successful result found, return all results
return results, nil
}
// parallelRace returns as soon as any search completes (like Promise.race)
// Returns immediately when first result arrives, regardless of success/failure
func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
results := make([]*types.Result, len(reqs))
resultChan := make(chan struct {
idx int
result *types.Result
}, len(reqs))
var wg sync.WaitGroup
done := make(chan struct{})
for i, req := range reqs {
wg.Add(1)
go func(idx int, r *types.Request) {
defer wg.Done()
result, _ := s.Search(ctx, r)
select {
case resultChan <- struct {
idx int
result *types.Result
}{idx, result}:
case <-done:
// Already got first result, discard this one
}
}(i, req)
}
// Close channel when all goroutines complete
go func() {
wg.Wait()
close(resultChan)
}()
// Return immediately when first result arrives
if res, ok := <-resultChan; ok {
results[res.idx] = res.result
close(done) // Signal other goroutines to stop sending
return results, nil
}
// No results (shouldn't happen with valid requests)
return results, nil
}
// BuildReferences converts search results to unified Reference format
func (s *Searcher) BuildReferences(results []*types.Result) []*types.Reference {
return BuildReferences(results)
}

402
agent/search/search_test.go Normal file
View file

@ -0,0 +1,402 @@
package search
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/types"
)
func TestNew(t *testing.T) {
t.Run("with nil config and uses", func(t *testing.T) {
s := New(nil, nil)
assert.NotNil(t, s)
assert.NotNil(t, s.config)
assert.NotNil(t, s.handlers)
assert.NotNil(t, s.citation)
assert.Equal(t, 3, len(s.handlers)) // web, kb, db
})
t.Run("with config", func(t *testing.T) {
cfg := &types.Config{
Web: &types.WebConfig{
Provider: "tavily",
MaxResults: 10,
},
KB: &types.KBConfig{
Collections: []string{"docs"},
Threshold: 0.8,
},
DB: &types.DBConfig{
Models: []string{"product"},
MaxResults: 20,
},
}
s := New(cfg, nil)
assert.NotNil(t, s)
assert.Equal(t, cfg, s.config)
})
t.Run("with uses", func(t *testing.T) {
uses := &Uses{
Search: "builtin",
Web: "builtin",
Keyword: "builtin",
QueryDSL: "builtin",
Rerank: "builtin",
}
s := New(nil, uses)
assert.NotNil(t, s)
})
}
func TestSearcher_Search_UnsupportedType(t *testing.T) {
s := New(nil, nil)
req := &types.Request{
Type: "unsupported",
Query: "test",
}
result, err := s.Search(nil, req)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.Equal(t, "unsupported search type", result.Error)
}
func TestSearcher_Search_Web(t *testing.T) {
// Note: This test uses skeleton handlers that return empty results
// Real tests with actual API calls are in handlers/web/*_test.go
cfg := &types.Config{
Web: &types.WebConfig{
Provider: "tavily",
},
}
s := New(cfg, &Uses{Web: "builtin"})
req := &types.Request{
Type: types.SearchTypeWeb,
Query: "test query",
Source: types.SourceAuto,
}
result, err := s.Search(nil, req)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.Equal(t, types.SearchTypeWeb, result.Type)
assert.Equal(t, "test query", result.Query)
// Note: actual result depends on API key availability
}
func TestSearcher_Search_KB(t *testing.T) {
cfg := &types.Config{
KB: &types.KBConfig{
Collections: []string{"docs"},
Threshold: 0.7,
},
}
s := New(cfg, nil)
req := &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Source: types.SourceHook,
Collections: []string{"docs"},
}
result, err := s.Search(nil, req)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.Equal(t, types.SearchTypeKB, result.Type)
assert.Equal(t, "test query", result.Query)
assert.Equal(t, types.SourceHook, result.Source)
// Skeleton returns empty items
assert.Equal(t, 0, len(result.Items))
}
func TestSearcher_Search_DB(t *testing.T) {
cfg := &types.Config{
DB: &types.DBConfig{
Models: []string{"product"},
MaxResults: 20,
},
}
s := New(cfg, &Uses{QueryDSL: "builtin"})
req := &types.Request{
Type: types.SearchTypeDB,
Query: "find products under $100",
Source: types.SourceUser,
Models: []string{"product"},
}
result, err := s.Search(nil, req)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.Equal(t, types.SearchTypeDB, result.Type)
assert.Equal(t, "find products under $100", result.Query)
assert.Equal(t, types.SourceUser, result.Source)
// Skeleton returns empty items
assert.Equal(t, 0, len(result.Items))
}
func TestSearcher_Search_WeightAssignment(t *testing.T) {
cfg := &types.Config{
KB: &types.KBConfig{
Collections: []string{"docs"},
},
Weights: &types.WeightsConfig{
User: 1.0,
Hook: 0.8,
Auto: 0.6,
},
}
s := New(cfg, nil)
// Test with different sources
sources := []struct {
source types.SourceType
weight float64
}{
{types.SourceUser, 1.0},
{types.SourceHook, 0.8},
{types.SourceAuto, 0.6},
}
for _, tc := range sources {
req := &types.Request{
Type: types.SearchTypeKB,
Query: "test",
Source: tc.source,
Collections: []string{"docs"},
}
result, err := s.Search(nil, req)
assert.NoError(t, err)
assert.NotNil(t, result)
// Items are empty in skeleton, so weight assignment can't be verified here
// This test ensures the code path works without error
}
}
func TestSearcher_All(t *testing.T) {
cfg := &types.Config{
KB: &types.KBConfig{
Collections: []string{"docs"},
},
DB: &types.DBConfig{
Models: []string{"product"},
},
}
s := New(cfg, nil)
reqs := []*types.Request{
{
Type: types.SearchTypeKB,
Query: "KB query",
Source: types.SourceAuto,
Collections: []string{"docs"},
},
{
Type: types.SearchTypeDB,
Query: "DB query",
Source: types.SourceAuto,
Models: []string{"product"},
},
}
// Test All() - waits for all searches to complete (like Promise.all)
results, err := s.All(nil, reqs)
assert.NoError(t, err)
assert.Equal(t, 2, len(results))
// Verify each result corresponds to its request
assert.Equal(t, types.SearchTypeKB, results[0].Type)
assert.Equal(t, "KB query", results[0].Query)
assert.Equal(t, types.SearchTypeDB, results[1].Type)
assert.Equal(t, "DB query", results[1].Query)
}
func TestSearcher_Any(t *testing.T) {
cfg := &types.Config{
KB: &types.KBConfig{
Collections: []string{"docs"},
},
DB: &types.DBConfig{
Models: []string{"product"},
},
}
s := New(cfg, nil)
reqs := []*types.Request{
{
Type: types.SearchTypeKB,
Query: "KB query",
Source: types.SourceAuto,
Collections: []string{"docs"},
},
{
Type: types.SearchTypeDB,
Query: "DB query",
Source: types.SourceAuto,
Models: []string{"product"},
},
}
// Test Any() - returns when first search has results (like Promise.any)
// Note: With skeleton handlers returning empty results, this will wait for all
results, err := s.Any(nil, reqs)
assert.NoError(t, err)
assert.Equal(t, 2, len(results))
}
func TestSearcher_Race(t *testing.T) {
cfg := &types.Config{
KB: &types.KBConfig{
Collections: []string{"docs"},
},
DB: &types.DBConfig{
Models: []string{"product"},
},
}
s := New(cfg, nil)
reqs := []*types.Request{
{
Type: types.SearchTypeKB,
Query: "KB query",
Source: types.SourceAuto,
Collections: []string{"docs"},
},
{
Type: types.SearchTypeDB,
Query: "DB query",
Source: types.SourceAuto,
Models: []string{"product"},
},
}
// Test Race() - returns when first search completes (like Promise.race)
results, err := s.Race(nil, reqs)
assert.NoError(t, err)
// At least one result should be set
hasResult := false
for _, r := range results {
if r != nil {
hasResult = true
break
}
}
assert.True(t, hasResult)
}
func TestSearcher_All_Empty(t *testing.T) {
s := New(nil, nil)
results, err := s.All(nil, []*types.Request{})
assert.NoError(t, err)
assert.Equal(t, 0, len(results))
}
func TestSearcher_Any_Empty(t *testing.T) {
s := New(nil, nil)
results, err := s.Any(nil, []*types.Request{})
assert.NoError(t, err)
assert.Equal(t, 0, len(results))
}
func TestSearcher_Race_Empty(t *testing.T) {
s := New(nil, nil)
results, err := s.Race(nil, []*types.Request{})
assert.NoError(t, err)
assert.Equal(t, 0, len(results))
}
func TestSearcher_All_ManyRequests(t *testing.T) {
cfg := &types.Config{
KB: &types.KBConfig{
Collections: []string{"docs"},
},
}
s := New(cfg, nil)
// Create multiple requests to test parallel execution
reqs := make([]*types.Request, 10)
for i := 0; i < 10; i++ {
reqs[i] = &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Source: types.SourceAuto,
Collections: []string{"docs"},
}
}
results, err := s.All(nil, reqs)
assert.NoError(t, err)
assert.Equal(t, 10, len(results))
// All results should be valid
for _, result := range results {
assert.NotNil(t, result)
assert.Equal(t, types.SearchTypeKB, result.Type)
}
}
func TestSearcher_BuildReferences(t *testing.T) {
s := New(nil, nil)
results := []*types.Result{
{
Type: types.SearchTypeWeb,
Items: []*types.ResultItem{
{
CitationID: "ref_001",
Type: types.SearchTypeWeb,
Source: types.SourceAuto,
Weight: 0.6,
Title: "Web Result",
Content: "Web content",
URL: "https://example.com",
},
},
},
{
Type: types.SearchTypeKB,
Items: []*types.ResultItem{
{
CitationID: "ref_002",
Type: types.SearchTypeKB,
Source: types.SourceHook,
Weight: 0.8,
Title: "KB Result",
Content: "KB content",
},
},
},
}
refs := s.BuildReferences(results)
assert.Equal(t, 2, len(refs))
assert.Equal(t, "ref_001", refs[0].ID)
assert.Equal(t, "ref_002", refs[1].ID)
}
func TestSearcher_CitationGeneration(t *testing.T) {
s := New(nil, nil)
// Reset citation generator for predictable IDs
s.citation.Reset()
// Note: This test would need actual results with items to verify citation generation
// The skeleton handlers return empty items, so we test the citation generator directly
id1 := s.citation.Next()
id2 := s.citation.Next()
id3 := s.citation.Next()
assert.Equal(t, "ref_001", id1)
assert.Equal(t, "ref_002", id2)
assert.Equal(t, "ref_003", id3)
}

335
agent/search/search_web_test.go Normal file
View file

@ -0,0 +1,335 @@
package search_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/search"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
)
// =============================================================================
// Web Search Integration Tests - Single Search
// =============================================================================
// TestWebSearch_Tavily tests web search using Tavily provider via assistant config
func TestWebSearch_Tavily(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-tavily test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Verify assistant config
assert.Equal(t, "tavily", ast.Search.Web.Provider)
// Create Searcher with assistant's config
uses := &search.Uses{Web: "builtin"}
s := search.New(ast.Search, uses)
// Execute search
req := &types.Request{
Type: types.SearchTypeWeb,
Query: "Yao App Engine low-code platform",
Source: types.SourceAuto,
Limit: 5,
}
result, err := s.Search(nil, req)
require.NoError(t, err)
require.NotNil(t, result)
require.Empty(t, result.Error, "Search should succeed, got error: %s", result.Error)
// Verify results
assert.NotEmpty(t, result.Items, "Should have search results")
for _, item := range result.Items {
assert.NotEmpty(t, item.CitationID, "Each item should have citation ID")
assert.NotEmpty(t, item.Content, "Each item should have content")
t.Logf(" [%s] %s - %s", item.CitationID, item.Title, item.URL)
}
t.Logf("Tavily search returned %d results", len(result.Items))
}
// TestWebSearch_Serper tests web search using Serper provider via assistant config
func TestWebSearch_Serper(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serper test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Verify assistant config
assert.Equal(t, "serper", ast.Search.Web.Provider)
// Create Searcher with assistant's config
uses := &search.Uses{Web: "builtin"}
s := search.New(ast.Search, uses)
// Execute search
req := &types.Request{
Type: types.SearchTypeWeb,
Query: "Go programming language concurrency",
Source: types.SourceAuto,
Limit: 5,
}
result, err := s.Search(nil, req)
require.NoError(t, err)
require.NotNil(t, result)
require.Empty(t, result.Error, "Search should succeed, got error: %s", result.Error)
// Verify results
assert.NotEmpty(t, result.Items, "Should have search results")
for _, item := range result.Items {
assert.NotEmpty(t, item.CitationID, "Each item should have citation ID")
t.Logf(" [%s] %s - %s", item.CitationID, item.Title, item.URL)
}
t.Logf("Serper search returned %d results", len(result.Items))
}
// TestWebSearch_SerpAPI tests web search using SerpAPI provider via assistant config
func TestWebSearch_SerpAPI(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serpapi test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
require.NoError(t, err)
require.NotNil(t, ast.Search)
require.NotNil(t, ast.Search.Web)
// Verify assistant config
assert.Equal(t, "serpapi", ast.Search.Web.Provider)
// Create Searcher with assistant's config
uses := &search.Uses{Web: "builtin"}
s := search.New(ast.Search, uses)
// Execute search
req := &types.Request{
Type: types.SearchTypeWeb,
Query: "Kubernetes container orchestration",
Source: types.SourceAuto,
Limit: 5,
}
result, err := s.Search(nil, req)
require.NoError(t, err)
require.NotNil(t, result)
require.Empty(t, result.Error, "Search should succeed, got error: %s", result.Error)
// Verify results
assert.NotEmpty(t, result.Items, "Should have search results")
for _, item := range result.Items {
assert.NotEmpty(t, item.CitationID, "Each item should have citation ID")
t.Logf(" [%s] %s - %s", item.CitationID, item.Title, item.URL)
}
t.Logf("SerpAPI search returned %d results", len(result.Items))
}
// =============================================================================
// Web Search Integration Tests - Parallel Search
// =============================================================================
// TestWebSearch_All tests parallel web search with All() - like Promise.all
func TestWebSearch_All(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-tavily test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
require.NoError(t, err)
require.NotNil(t, ast.Search)
// Create Searcher
uses := &search.Uses{Web: "builtin"}
s := search.New(ast.Search, uses)
// Multiple queries
reqs := []*types.Request{
{Type: types.SearchTypeWeb, Query: "artificial intelligence", Source: types.SourceAuto, Limit: 3},
{Type: types.SearchTypeWeb, Query: "machine learning", Source: types.SourceAuto, Limit: 3},
{Type: types.SearchTypeWeb, Query: "deep learning", Source: types.SourceAuto, Limit: 3},
}
// Execute parallel search with All() - waits for all searches to complete
results, err := s.All(nil, reqs)
require.NoError(t, err)
require.Len(t, results, 3, "Should have 3 results")
// Verify all results
for i, result := range results {
require.NotNil(t, result, "Result %d should not be nil", i)
if result.Error == "" {
assert.NotEmpty(t, result.Items, "Result %d should have items", i)
t.Logf("Query '%s': %d results", reqs[i].Query, len(result.Items))
} else {
t.Logf("Query '%s': error - %s", reqs[i].Query, result.Error)
}
}
}
// TestWebSearch_Any tests parallel web search with Any() - like Promise.any
func TestWebSearch_Any(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serper test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
require.NoError(t, err)
require.NotNil(t, ast.Search)
// Create Searcher
uses := &search.Uses{Web: "builtin"}
s := search.New(ast.Search, uses)
// Multiple queries
reqs := []*types.Request{
{Type: types.SearchTypeWeb, Query: "golang channels", Source: types.SourceAuto, Limit: 3},
{Type: types.SearchTypeWeb, Query: "rust ownership", Source: types.SourceAuto, Limit: 3},
{Type: types.SearchTypeWeb, Query: "python asyncio", Source: types.SourceAuto, Limit: 3},
}
// Execute parallel search with Any() - returns when first search succeeds
results, err := s.Any(nil, reqs)
require.NoError(t, err)
// Any() returns as soon as any search succeeds
hasSuccess := false
for _, result := range results {
if result != nil && len(result.Items) > 0 && result.Error == "" {
hasSuccess = true
t.Logf("First success: '%s' with %d results", result.Query, len(result.Items))
break
}
}
assert.True(t, hasSuccess, "At least one search should succeed")
}
// TestWebSearch_Race tests parallel web search with Race() - like Promise.race
func TestWebSearch_Race(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-tavily test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
require.NoError(t, err)
require.NotNil(t, ast.Search)
// Create Searcher
uses := &search.Uses{Web: "builtin"}
s := search.New(ast.Search, uses)
// Multiple queries
reqs := []*types.Request{
{Type: types.SearchTypeWeb, Query: "docker containers", Source: types.SourceAuto, Limit: 3},
{Type: types.SearchTypeWeb, Query: "kubernetes pods", Source: types.SourceAuto, Limit: 3},
}
// Execute parallel search with Race() - returns when first search completes
results, err := s.Race(nil, reqs)
require.NoError(t, err)
// Race() returns immediately when first result arrives
hasResult := false
for _, result := range results {
if result != nil {
hasResult = true
t.Logf("First to complete: '%s'", result.Query)
break
}
}
assert.True(t, hasResult, "Should have at least one result")
}
// =============================================================================
// Web Search - Citation and Reference Tests
// =============================================================================
// TestWebSearch_BuildReferences tests building references from web search results
func TestWebSearch_BuildReferences(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-tavily test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
require.NoError(t, err)
require.NotNil(t, ast.Search)
// Create Searcher with weights config
uses := &search.Uses{Web: "builtin"}
s := search.New(ast.Search, uses)
// Execute search
req := &types.Request{
Type: types.SearchTypeWeb,
Query: "OpenAI GPT-4",
Source: types.SourceAuto,
Limit: 5,
}
result, err := s.Search(nil, req)
require.NoError(t, err)
require.NotNil(t, result)
require.Empty(t, result.Error, "Search should succeed")
require.NotEmpty(t, result.Items, "Should have results")
// Build references
refs := s.BuildReferences([]*types.Result{result})
assert.NotEmpty(t, refs, "Should have references")
for _, ref := range refs {
assert.NotEmpty(t, ref.ID, "Reference should have ID")
assert.Equal(t, types.SearchTypeWeb, ref.Type, "Reference type should be web")
assert.Equal(t, types.SourceAuto, ref.Source, "Reference source should be auto")
t.Logf(" Ref: %s - %s (weight: %.2f)", ref.ID, ref.Title, ref.Weight)
}
}
// =============================================================================
// Web Search - Error Handling Tests
// =============================================================================
// TestWebSearch_SiteRestriction tests web search with site restriction
func TestWebSearch_SiteRestriction(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the web-serper test assistant
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
require.NoError(t, err)
require.NotNil(t, ast.Search)
// Create Searcher
uses := &search.Uses{Web: "builtin"}
s := search.New(ast.Search, uses)
// Execute search with site restriction
req := &types.Request{
Type: types.SearchTypeWeb,
Query: "yao-app-engine",
Source: types.SourceAuto,
Sites: []string{"github.com"},
Limit: 5,
}
result, err := s.Search(nil, req)
require.NoError(t, err)
require.NotNil(t, result)
if result.Error == "" && len(result.Items) > 0 {
// Log results
for _, item := range result.Items {
t.Logf(" %s - %s", item.Title, item.URL)
}
}
}

118
agent/search/types/config.go Normal file
View file

@ -0,0 +1,118 @@
package types
// Config represents the complete search configuration
type Config struct {
Web *WebConfig `json:"web,omitempty" yaml:"web,omitempty"`
KB *KBConfig `json:"kb,omitempty" yaml:"kb,omitempty"`
DB *DBConfig `json:"db,omitempty" yaml:"db,omitempty"`
Keyword *KeywordConfig `json:"keyword,omitempty" yaml:"keyword,omitempty"`
QueryDSL *QueryDSLConfig `json:"querydsl,omitempty" yaml:"querydsl,omitempty"`
Rerank *RerankConfig `json:"rerank,omitempty" yaml:"rerank,omitempty"`
Citation *CitationConfig `json:"citation,omitempty" yaml:"citation,omitempty"`
Weights *WeightsConfig `json:"weights,omitempty" yaml:"weights,omitempty"`
Options *OptionsConfig `json:"options,omitempty" yaml:"options,omitempty"`
}
// WebConfig for web search settings
// Note: uses.web determines the mode (builtin/agent/mcp)
// Provider is only used when uses.web = "builtin"
type WebConfig struct {
Provider string `json:"provider,omitempty" yaml:"provider,omitempty"` // "tavily", "serper", or "serpapi" (for builtin mode)
APIKeyEnv string `json:"api_key_env,omitempty" yaml:"api_key_env,omitempty"` // Environment variable for API key
MaxResults int `json:"max_results,omitempty" yaml:"max_results,omitempty"` // Max results (default: 10)
Engine string `json:"engine,omitempty" yaml:"engine,omitempty"` // Search engine for SerpAPI: "google", "bing", "baidu", "yandex", etc. (default: "google")
}
// KBConfig for knowledge base search settings
type KBConfig struct {
Collections []string `json:"collections,omitempty" yaml:"collections,omitempty"` // Default collections
Threshold float64 `json:"threshold,omitempty" yaml:"threshold,omitempty"` // Similarity threshold (default: 0.7)
Graph bool `json:"graph,omitempty" yaml:"graph,omitempty"` // Enable GraphRAG (default: false)
}
// DBConfig for database search settings
type DBConfig struct {
Models []string `json:"models,omitempty" yaml:"models,omitempty"` // Default models
MaxResults int `json:"max_results,omitempty" yaml:"max_results,omitempty"` // Max results (default: 20)
}
// KeywordConfig for keyword extraction
type KeywordConfig struct {
MaxKeywords int `json:"max_keywords,omitempty" yaml:"max_keywords,omitempty"` // Max keywords (default: 10)
Language string `json:"language,omitempty" yaml:"language,omitempty"` // "auto", "en", "zh", etc.
}
// KeywordOptions for keyword extraction (runtime options)
type KeywordOptions struct {
MaxKeywords int `json:"max_keywords,omitempty"`
Language string `json:"language,omitempty"`
}
// QueryDSLConfig for QueryDSL generation from natural language
type QueryDSLConfig struct {
Strict bool `json:"strict,omitempty" yaml:"strict,omitempty"` // Fail if generation fails (default: false)
}
// RerankConfig for reranking
type RerankConfig struct {
TopN int `json:"top_n,omitempty" yaml:"top_n,omitempty"` // Return top N (default: 10)
}
// CitationConfig for citation format
type CitationConfig struct {
Format string `json:"format,omitempty" yaml:"format,omitempty"` // Default: "#ref:{id}"
AutoInjectPrompt bool `json:"auto_inject_prompt,omitempty" yaml:"auto_inject_prompt,omitempty"` // Auto-inject prompt (default: true)
CustomPrompt string `json:"custom_prompt,omitempty" yaml:"custom_prompt,omitempty"` // Custom prompt template
}
// WeightsConfig for source weighting
type WeightsConfig struct {
User float64 `json:"user,omitempty" yaml:"user,omitempty"` // User-provided (default: 1.0)
Hook float64 `json:"hook,omitempty" yaml:"hook,omitempty"` // Hook results (default: 0.8)
Auto float64 `json:"auto,omitempty" yaml:"auto,omitempty"` // Auto search (default: 0.6)
}
// OptionsConfig for search behavior
type OptionsConfig struct {
SkipThreshold int `json:"skip_threshold,omitempty" yaml:"skip_threshold,omitempty"` // Skip auto search if user provides >= N results
}
// GetWeight returns the weight for a source type
func (c *Config) GetWeight(source SourceType) float64 {
if c == nil || c.Weights == nil {
return getDefaultWeight(source)
}
switch source {
case SourceUser:
if c.Weights.User > 0 {
return c.Weights.User
}
return 1.0
case SourceHook:
if c.Weights.Hook > 0 {
return c.Weights.Hook
}
return 0.8
case SourceAuto:
if c.Weights.Auto > 0 {
return c.Weights.Auto
}
return 0.6
default:
return 0.6
}
}
// getDefaultWeight returns default weight for a source type
func getDefaultWeight(source SourceType) float64 {
switch source {
case SourceUser:
return 1.0
case SourceHook:
return 0.8
case SourceAuto:
return 0.6
default:
return 0.6
}
}

12
agent/search/types/graph.go Normal file
View file

@ -0,0 +1,12 @@
package types
// GraphNode represents a related entity from knowledge graph
type GraphNode struct {
ID string `json:"id"`
Type string `json:"type"` // Entity type
Name string `json:"name"` // Entity name
Description string `json:"description,omitempty"` // Entity description
Relation string `json:"relation,omitempty"` // Relationship to query
Score float64 `json:"score,omitempty"` // Relevance score
Metadata map[string]interface{} `json:"metadata,omitempty"`
}

22
agent/search/types/reference.go Normal file
View file

@ -0,0 +1,22 @@
package types
// Reference is the unified structure for all data sources
// Used to build LLM context from search results
type Reference struct {
ID string `json:"id"` // Unique citation ID: "ref_001", "ref_002"
Type SearchType `json:"type"` // Data type: "web", "kb", "db"
Source SourceType `json:"source"` // Origin: "user", "hook", "auto"
Weight float64 `json:"weight"` // Relevance weight (1.0=highest, 0.6=lowest)
Score float64 `json:"score"` // Relevance score (0-1)
Title string `json:"title"` // Optional title
Content string `json:"content"` // Main content
URL string `json:"url"` // Optional URL
Meta map[string]interface{} `json:"meta"` // Additional metadata
}
// ReferenceContext holds the formatted references for LLM input
type ReferenceContext struct {
References []*Reference `json:"references"` // All references
XML string `json:"xml"` // Formatted <references> XML
Prompt string `json:"prompt"` // Citation instruction prompt
}

114
agent/search/types/types.go Normal file
View file

@ -0,0 +1,114 @@
package types
import (
"github.com/yaoapp/gou/query/gou"
)
// SearchType represents the type of search
type SearchType string
// SearchType constants
const (
SearchTypeWeb SearchType = "web" // Web/Internet search
SearchTypeKB SearchType = "kb" // Knowledge base vector search
SearchTypeDB SearchType = "db" // Database search (Yao Model/QueryDSL)
)
// SourceType represents where the search result came from
type SourceType string
// SourceType constants
const (
SourceUser SourceType = "user" // User-provided DataContent (highest priority)
SourceHook SourceType = "hook" // Hook ctx.search.*() results
SourceAuto SourceType = "auto" // Auto search results (lowest priority)
)
// Request represents a search request
type Request struct {
// Common fields
Query string `json:"query"` // Search query (natural language)
Type SearchType `json:"type"` // Search type: "web", "kb", or "db"
Limit int `json:"limit,omitempty"` // Max results (default: 10)
Source SourceType `json:"source"` // Source of this request (user/hook/auto)
// Web search specific
Sites []string `json:"sites,omitempty"` // Restrict to specific sites
TimeRange string `json:"time_range,omitempty"` // "day", "week", "month", "year"
// Knowledge base specific
Collections []string `json:"collections,omitempty"` // KB collection IDs
Threshold float64 `json:"threshold,omitempty"` // Similarity threshold (0-1)
Graph bool `json:"graph,omitempty"` // Enable graph association
// Database search specific
Models []string `json:"models,omitempty"` // Model IDs (e.g., "user", "agents.mybot.product")
Wheres []gou.Where `json:"wheres,omitempty"` // Pre-defined filters (optional), uses GOU QueryDSL Where
Orders gou.Orders `json:"orders,omitempty"` // Sort orders (optional), uses GOU QueryDSL Orders
Select []string `json:"select,omitempty"` // Fields to return (optional)
// Reranking
Rerank *RerankOptions `json:"rerank,omitempty"`
}
// RerankOptions controls result reranking
// Reranker type is determined by uses.rerank in agent/agent.yml
type RerankOptions struct {
TopN int `json:"top_n,omitempty"` // Return top N after reranking
}
// Result represents the search result
type Result struct {
Type SearchType `json:"type"` // Search type
Query string `json:"query"` // Original query
Source SourceType `json:"source"` // Source of this result
Items []*ResultItem `json:"items"` // Result items
Total int `json:"total"` // Total matches
Duration int64 `json:"duration_ms"` // Search duration in ms
Error string `json:"error,omitempty"` // Error message if failed
// Graph associations (KB only, if enabled)
GraphNodes []*GraphNode `json:"graph_nodes,omitempty"`
}
// ResultItem represents a single search result item
type ResultItem struct {
// Citation
CitationID string `json:"citation_id"` // Unique ID for LLM reference: "ref_001"
// Weighting
Source SourceType `json:"source"` // Source type: "user", "hook", "auto"
Weight float64 `json:"weight"` // Source weight (from config)
Score float64 `json:"score,omitempty"` // Relevance score (0-1)
// Common fields
Type SearchType `json:"type"` // Search type for this item
Title string `json:"title,omitempty"` // Title/headline
Content string `json:"content"` // Main content/snippet
URL string `json:"url,omitempty"` // Source URL
// KB specific
DocumentID string `json:"document_id,omitempty"` // Source document ID
Collection string `json:"collection,omitempty"` // Collection name
// DB specific
Model string `json:"model,omitempty"` // Model ID
RecordID interface{} `json:"record_id,omitempty"` // Record primary key
Data map[string]interface{} `json:"data,omitempty"` // Full record data
// Metadata
Metadata map[string]interface{} `json:"metadata,omitempty"` // Additional metadata
}
// ProcessedQuery represents a processed query ready for execution
type ProcessedQuery struct {
Type SearchType `json:"type"`
Keywords []string `json:"keywords,omitempty"` // For web search
Vector []float32 `json:"vector,omitempty"` // For KB search
DSL *gou.QueryDSL `json:"dsl,omitempty"` // For DB search, uses GOU QueryDSL
}
// Note: For QueryDSL and Model types, use GOU types directly:
// - github.com/yaoapp/gou/query/gou.QueryDSL
// - github.com/yaoapp/gou/model.Model
// - github.com/yaoapp/gou/model.Column

58
agent/store/types/convert.go
View file

@ -9,7 +9,9 @@ import (
"github.com/spf13/cast"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
searchTypes "github.com/yaoapp/yao/agent/search/types"
)
// ToKnowledgeBase converts various types to KnowledgeBase
@ -616,3 +618,59 @@ func ToPromptPresets(v interface{}) (map[string][]Prompt, error) {
return result, nil
}
}
// ToUses converts various types to context.Uses
func ToUses(v interface{}) (*context.Uses, error) {
if v == nil {
return nil, nil
}
switch uses := v.(type) {
case *context.Uses:
return uses, nil
case context.Uses:
return &uses, nil
default:
raw, err := jsoniter.Marshal(uses)
if err != nil {
return nil, fmt.Errorf("uses format error: %s", err.Error())
}
var result context.Uses
err = jsoniter.Unmarshal(raw, &result)
if err != nil {
return nil, fmt.Errorf("uses format error: %s", err.Error())
}
return &result, nil
}
}
// ToSearchConfig converts various types to searchTypes.Config
func ToSearchConfig(v interface{}) (*searchTypes.Config, error) {
if v == nil {
return nil, nil
}
switch cfg := v.(type) {
case *searchTypes.Config:
return cfg, nil
case searchTypes.Config:
return &cfg, nil
default:
raw, err := jsoniter.Marshal(cfg)
if err != nil {
return nil, fmt.Errorf("search config format error: %s", err.Error())
}
var result searchTypes.Config
err = jsoniter.Unmarshal(raw, &result)
if err != nil {
return nil, fmt.Errorf("search config format error: %s", err.Error())
}
return &result, nil
}
}

2
agent/store/types/fields.go
View file

@ -33,6 +33,7 @@ var AssistantAllowedFields = map[string]bool{
"share": true,
"locales": true,
"uses": true,
"search": true,
"automated": true,
"mentionable": true,
"created_at": true,
@ -104,6 +105,7 @@ var AssistantFullFields = []string{
"share",
"locales",
"uses",
"search",
"automated",
"mentionable",
"created_at",

2
agent/store/types/types.go
View file

@ -9,6 +9,7 @@ import (
"github.com/yaoapp/xun/dbal/query"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
searchTypes "github.com/yaoapp/yao/agent/search/types"
)
// Setting represents the conversation configuration structure
@ -436,6 +437,7 @@ type AssistantModel struct {
Source string `json:"source,omitempty"` // Hook script source code
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings
Search *searchTypes.Config `json:"search,omitempty"` // Search configuration (web, kb, db, citation, weights, etc.)
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp

18
agent/store/xun/assistant.go
View file

@ -10,6 +10,7 @@ import (
"github.com/yaoapp/xun/dbal/query"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
searchTypes "github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/store/types"
)
@ -183,6 +184,7 @@ func (store *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
"placeholder": assistant.Placeholder,
"locales": assistant.Locales,
"uses": assistant.Uses,
"search": assistant.Search,
}
for field, value := range jsonFields {
@ -241,7 +243,7 @@ func (store *Xun) UpdateAssistant(assistantID string, updates map[string]interfa
data := make(map[string]interface{})
// List of fields that need JSON marshaling
jsonFields := []string{"options", "tags", "modes", "prompts", "prompt_presets", "connector_options", "kb", "db", "mcp", "workflow", "placeholder", "locales", "uses"}
jsonFields := []string{"options", "tags", "modes", "prompts", "prompt_presets", "connector_options", "kb", "db", "mcp", "workflow", "placeholder", "locales", "uses", "search"}
jsonFieldSet := make(map[string]bool)
for _, field := range jsonFields {
jsonFieldSet[field] = true
@ -441,7 +443,7 @@ func (store *Xun) GetAssistants(filter types.AssistantFilter, locale ...string)
// Convert rows to types.AssistantModel slice
assistants := make([]*types.AssistantModel, 0, len(rows))
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "mcp", "placeholder", "locales", "uses"}
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "mcp", "placeholder", "locales", "uses", "search"}
for _, row := range rows {
data := row.ToMap()
@ -514,7 +516,7 @@ func (store *Xun) GetAssistant(assistantID string, fields []string, locale ...st
}
// Parse JSON fields
jsonFields := []string{"tags", "modes", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "db", "mcp", "placeholder", "locales", "uses"}
jsonFields := []string{"tags", "modes", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "db", "mcp", "placeholder", "locales", "uses", "search"}
store.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel
@ -659,6 +661,16 @@ func (store *Xun) GetAssistant(assistantID string, fields []string, locale ...st
}
}
if search, has := data["search"]; has && search != nil {
raw, err := jsoniter.Marshal(search)
if err == nil {
var s searchTypes.Config
if err := jsoniter.Unmarshal(raw, &s); err == nil {
model.Search = &s
}
}
}
// Apply i18n translation if locale is provided
if len(locale) > 0 && locale[0] != "" {
store.translate(model, assistantID, locale[0])

315
agent/store/xun/assistant_test.go
View file

@ -10,6 +10,7 @@ import (
"github.com/yaoapp/xun/dbal/query"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
searchTypes "github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/agent/store/xun"
"github.com/yaoapp/yao/config"
@ -452,6 +453,198 @@ func TestSaveAssistant(t *testing.T) {
}
})
t.Run("SearchConfiguration", func(t *testing.T) {
// Test assistant with Search configuration
assistant := &types.AssistantModel{
Name: "Search Config Test Assistant",
Type: "assistant",
Connector: "openai",
Share: "private",
Search: &searchTypes.Config{
Web: &searchTypes.WebConfig{
Provider: "tavily",
MaxResults: 15,
},
KB: &searchTypes.KBConfig{
Collections: []string{"docs", "faq"},
Threshold: 0.8,
Graph: true,
},
DB: &searchTypes.DBConfig{
Models: []string{"user", "product"},
MaxResults: 50,
},
Citation: &searchTypes.CitationConfig{
Format: "#ref:{id}",
AutoInjectPrompt: true,
},
Weights: &searchTypes.WeightsConfig{
User: 1.0,
Hook: 0.9,
Auto: 0.7,
},
},
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with search config: %v", err)
}
// Retrieve and verify search configuration - search is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Search == nil {
t.Fatal("Expected search to be set")
}
// Verify Web config
if retrieved.Search.Web == nil {
t.Fatal("Expected search.web to be set")
}
if retrieved.Search.Web.Provider != "tavily" {
t.Errorf("Expected web provider 'tavily', got '%s'", retrieved.Search.Web.Provider)
}
if retrieved.Search.Web.MaxResults != 15 {
t.Errorf("Expected web max_results 15, got %d", retrieved.Search.Web.MaxResults)
}
// Verify KB config
if retrieved.Search.KB == nil {
t.Fatal("Expected search.kb to be set")
}
if len(retrieved.Search.KB.Collections) != 2 {
t.Errorf("Expected 2 KB collections, got %d", len(retrieved.Search.KB.Collections))
}
if retrieved.Search.KB.Collections[0] != "docs" {
t.Errorf("Expected first collection 'docs', got '%s'", retrieved.Search.KB.Collections[0])
}
if retrieved.Search.KB.Threshold != 0.8 {
t.Errorf("Expected KB threshold 0.8, got %f", retrieved.Search.KB.Threshold)
}
if !retrieved.Search.KB.Graph {
t.Error("Expected KB graph to be true")
}
// Verify DB config
if retrieved.Search.DB == nil {
t.Fatal("Expected search.db to be set")
}
if len(retrieved.Search.DB.Models) != 2 {
t.Errorf("Expected 2 DB models, got %d", len(retrieved.Search.DB.Models))
}
if retrieved.Search.DB.MaxResults != 50 {
t.Errorf("Expected DB max_results 50, got %d", retrieved.Search.DB.MaxResults)
}
// Verify Citation config
if retrieved.Search.Citation == nil {
t.Fatal("Expected search.citation to be set")
}
if retrieved.Search.Citation.Format != "#ref:{id}" {
t.Errorf("Expected citation format '#ref:{id}', got '%s'", retrieved.Search.Citation.Format)
}
if !retrieved.Search.Citation.AutoInjectPrompt {
t.Error("Expected citation auto_inject_prompt to be true")
}
// Verify Weights config
if retrieved.Search.Weights == nil {
t.Fatal("Expected search.weights to be set")
}
if retrieved.Search.Weights.User != 1.0 {
t.Errorf("Expected weights.user 1.0, got %f", retrieved.Search.Weights.User)
}
if retrieved.Search.Weights.Hook != 0.9 {
t.Errorf("Expected weights.hook 0.9, got %f", retrieved.Search.Weights.Hook)
}
if retrieved.Search.Weights.Auto != 0.7 {
t.Errorf("Expected weights.auto 0.7, got %f", retrieved.Search.Weights.Auto)
}
t.Logf("Successfully saved and retrieved assistant with search configuration")
})
t.Run("NilSearchConfiguration", func(t *testing.T) {
// Test assistant without Search configuration
assistant := &types.AssistantModel{
Name: "No Search Config Assistant",
Type: "assistant",
Connector: "openai",
Share: "private",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant without search: %v", err)
}
// Retrieve and verify search is nil - request all fields to check search
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Search != nil {
t.Errorf("Expected search to be nil, got %+v", retrieved.Search)
}
})
t.Run("PartialSearchConfiguration", func(t *testing.T) {
// Test assistant with partial Search configuration
assistant := &types.AssistantModel{
Name: "Partial Search Config Assistant",
Type: "assistant",
Connector: "openai",
Share: "private",
Search: &searchTypes.Config{
Web: &searchTypes.WebConfig{
Provider: "serper",
},
// KB, DB, Citation, Weights not set
},
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with partial search: %v", err)
}
// Retrieve and verify - request all fields for search
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Search == nil {
t.Fatal("Expected search to be set")
}
if retrieved.Search.Web == nil {
t.Fatal("Expected search.web to be set")
}
if retrieved.Search.Web.Provider != "serper" {
t.Errorf("Expected web provider 'serper', got '%s'", retrieved.Search.Web.Provider)
}
// Other fields should be nil
if retrieved.Search.KB != nil {
t.Errorf("Expected search.kb to be nil, got %+v", retrieved.Search.KB)
}
if retrieved.Search.DB != nil {
t.Errorf("Expected search.db to be nil, got %+v", retrieved.Search.DB)
}
if retrieved.Search.Citation != nil {
t.Errorf("Expected search.citation to be nil, got %+v", retrieved.Search.Citation)
}
if retrieved.Search.Weights != nil {
t.Errorf("Expected search.weights to be nil, got %+v", retrieved.Search.Weights)
}
})
t.Run("ConnectorOptions", func(t *testing.T) {
// Test assistant with connector options
optionalTrue := true
@ -2725,6 +2918,128 @@ func TestUpdateAssistant(t *testing.T) {
}
})
t.Run("UpdateSearch", func(t *testing.T) {
// Create assistant without search
assistant := &types.AssistantModel{
Name: "Search Update Test",
Type: "assistant",
Connector: "openai",
Share: "private",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to create assistant: %v", err)
}
// Update with search configuration
updates := map[string]interface{}{
"search": &searchTypes.Config{
Web: &searchTypes.WebConfig{
Provider: "tavily",
MaxResults: 20,
},
KB: &searchTypes.KBConfig{
Collections: []string{"knowledge"},
Threshold: 0.75,
},
},
}
err = store.UpdateAssistant(id, updates)
if err != nil {
t.Fatalf("Failed to update search: %v", err)
}
// Verify updates - search is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Search == nil {
t.Fatal("Expected search to be set")
}
if retrieved.Search.Web == nil {
t.Fatal("Expected search.web to be set")
}
if retrieved.Search.Web.Provider != "tavily" {
t.Errorf("Expected web provider 'tavily', got '%s'", retrieved.Search.Web.Provider)
}
if retrieved.Search.Web.MaxResults != 20 {
t.Errorf("Expected web max_results 20, got %d", retrieved.Search.Web.MaxResults)
}
if retrieved.Search.KB == nil {
t.Fatal("Expected search.kb to be set")
}
if len(retrieved.Search.KB.Collections) != 1 {
t.Errorf("Expected 1 KB collection, got %d", len(retrieved.Search.KB.Collections))
}
if retrieved.Search.KB.Threshold != 0.75 {
t.Errorf("Expected KB threshold 0.75, got %f", retrieved.Search.KB.Threshold)
}
// Update to change search configuration
updates2 := map[string]interface{}{
"search": &searchTypes.Config{
Web: &searchTypes.WebConfig{
Provider: "serper",
MaxResults: 30,
},
Citation: &searchTypes.CitationConfig{
Format: "#cite:{id}",
AutoInjectPrompt: false,
},
},
}
err = store.UpdateAssistant(id, updates2)
if err != nil {
t.Fatalf("Failed to update search again: %v", err)
}
// Verify second update - search is NOT in default fields
retrieved2, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved2.Search.Web.Provider != "serper" {
t.Errorf("Expected web provider 'serper', got '%s'", retrieved2.Search.Web.Provider)
}
if retrieved2.Search.Web.MaxResults != 30 {
t.Errorf("Expected web max_results 30, got %d", retrieved2.Search.Web.MaxResults)
}
if retrieved2.Search.Citation == nil {
t.Fatal("Expected search.citation to be set")
}
if retrieved2.Search.Citation.Format != "#cite:{id}" {
t.Errorf("Expected citation format '#cite:{id}', got '%s'", retrieved2.Search.Citation.Format)
}
// Update to remove search (set to nil)
updates3 := map[string]interface{}{
"search": nil,
}
err = store.UpdateAssistant(id, updates3)
if err != nil {
t.Fatalf("Failed to set search to nil: %v", err)
}
// Verify search is nil - search is NOT in default fields
retrieved3, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved3.Search != nil {
t.Errorf("Expected search to be nil, got %+v", retrieved3.Search)
}
})
t.Run("UpdatePermissionFields", func(t *testing.T) {
// Create assistant with permission fields
assistant := &types.AssistantModel{

8
agent/types/types.go
View file

@ -3,6 +3,7 @@ package types
import (
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/assistant"
searchTypes "github.com/yaoapp/yao/agent/search/types"
store "github.com/yaoapp/yao/agent/store/types"
)
@ -19,6 +20,7 @@ type DSL struct {
// ===============================
Models map[string]openai.Capabilities `json:"models,omitempty" yaml:"models,omitempty"` // The model capabilities configuration
KB *store.KBSetting `json:"kb,omitempty" yaml:"kb,omitempty"` // The knowledge base configuration loaded from agent/kb.yml
Search *searchTypes.Config `json:"search,omitempty" yaml:"search,omitempty"` // The search configuration loaded from agent/search.yao
// Internal
// ===============================
@ -38,6 +40,12 @@ type Uses struct {
Audio string `json:"audio,omitempty" yaml:"audio,omitempty"` // The assistant for processing audio (speech-to-text, text-to-speech). If the model doesn't support audio, use this to convert audio to text. Format: "agent" or "mcp:mcp_server_id"
Search string `json:"search,omitempty" yaml:"search,omitempty"` // The assistant for searching the knowledge, global web search. If not set, and the assistant enable the knowledge, it will search the result from the knowledge automatically.
Fetch string `json:"fetch,omitempty" yaml:"fetch,omitempty"` // The assistant for fetching the http/https/ftp/sftp/etc. file, and return the file's content. if not set, use the http process to fetch the file.
// Search-related processing tools (NLP)
Web string `json:"web,omitempty" yaml:"web,omitempty"` // Web search handler: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
Keyword string `json:"keyword,omitempty" yaml:"keyword,omitempty"` // Keyword extraction: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
QueryDSL string `json:"querydsl,omitempty" yaml:"querydsl,omitempty"` // QueryDSL generation: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
Rerank string `json:"rerank,omitempty" yaml:"rerank,omitempty"` // Result reranking: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
}
// Mention Structure

286
data/bindata.go
View file

File diff suppressed because one or more lines are too long

7
yao/models/agent/assistant.mod.yao
View file

@ -235,6 +235,13 @@
"comment": "Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings",
"nullable": true
},
{
"name": "search",
"type": "json",
"label": "Search",
"comment": "Search configuration (web, kb, db, citation, weights, etc.)",
"nullable": true
},
{
"name": "automated",
"type": "boolean",