diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index cb659647..200c807e 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -203,7 +203,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Execute Auto Search (if enabled) // ================================================ if ast.shouldAutoSearch(ctx, createResponse) { - refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse) + refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse, opts) if refCtx != nil && len(refCtx.References) > 0 { completionMessages = ast.injectSearchContext(completionMessages, refCtx) } diff --git a/agent/assistant/search.go b/agent/assistant/search.go index dd9a442a..f6164a63 100644 --- a/agent/assistant/search.go +++ b/agent/assistant/search.go @@ -5,6 +5,7 @@ import ( "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" ) @@ -73,7 +74,8 @@ func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResp // executeAutoSearch executes auto search based on configuration // Returns ReferenceContext with results and formatted context -func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) *searchTypes.ReferenceContext { +// 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") @@ -103,6 +105,30 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context 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 { diff --git a/agent/assistant/search_auto_keyword_test.go b/agent/assistant/search_auto_keyword_test.go new file mode 100644 index 00000000..d5e2a1b0 --- /dev/null +++ b/agent/assistant/search_auto_keyword_test.go @@ -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") + }) +} diff --git a/agent/context/types.go b/agent/context/types.go index e136b365..5c78ffe8 100644 --- a/agent/context/types.go +++ b/agent/context/types.go @@ -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 diff --git a/agent/search/DESIGN.md b/agent/search/DESIGN.md index 8b0cd4a3..077a3a69 100644 --- a/agent/search/DESIGN.md +++ b/agent/search/DESIGN.md @@ -1559,6 +1559,7 @@ Stream(ctx, messages, options) │ ├── IF Create Hook returned uses.search="disabled" → SKIP │ └── ELSE → Execute Auto Search (executeAutoSearch) │ ├── Read assistant's search config (GetMergedSearchConfig) + │ ├── Extract keywords (if uses.keyword && !Skip.Keyword) │ ├── Build search requests (buildSearchRequests) │ ├── Execute web/kb/db in parallel (searcher.All) │ ├── Build reference context (BuildReferenceContext) @@ -1586,7 +1587,8 @@ Stream(ctx, messages, options) func (ast *Assistant) shouldAutoSearch(ctx *context.Context, createResponse *context.HookCreateResponse) bool // executeAutoSearch executes auto search based on configuration -func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) *searchTypes.ReferenceContext +// opts is optional, used to check Skip.Keyword for keyword extraction +func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, opts ...*context.Options) *searchTypes.ReferenceContext // injectSearchContext injects search results into messages func (ast *Assistant) injectSearchContext(messages []context.Message, refCtx *searchTypes.ReferenceContext) []context.Message @@ -1598,18 +1600,49 @@ func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResp func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Config) []*searchTypes.Request ``` +**Keyword Extraction in executeAutoSearch:** + +When `uses.keyword` is configured and `opts.Skip.Keyword` is not true, keyword extraction is performed before web search: + +```go +// Extract keywords for web search if: +// 1. uses.keyword is configured (not empty) +// 2. Skip.Keyword is not true +// 3. Web search is enabled +if webSearchEnabled && !skipKeyword && searchUses.Keyword != "" { + extractor := keyword.NewExtractor(searchUses.Keyword, searchConfig.Keyword) + keywords, err := extractor.Extract(ctx, query, nil) + if err == nil && len(keywords) > 0 { + query = strings.Join(keywords, " ") + } +} +``` + **Integration in agent.go:** ```go // In Stream(), after BuildContent: if ast.shouldAutoSearch(ctx, createResponse) { - refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse) + refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse, opts) if refCtx != nil && len(refCtx.References) > 0 { completionMessages = ast.injectSearchContext(completionMessages, refCtx) } } ``` +**Skip.Keyword Option (`context.Options.Skip`):** + +```go +type Skip struct { + History bool `json:"history"` // Skip saving chat history + Trace bool `json:"trace"` // Skip trace logging + Output bool `json:"output"` // Skip output to client + Keyword bool `json:"keyword"` // Skip keyword extraction for web search +} +``` + +Use `Skip.Keyword = true` when you want to use the raw query directly without keyword extraction. + ### Control via Uses.Search Search is controlled via the `Uses` mechanism, following the merge hierarchy: