Enhance Auto Search Functionality with Keyword Extraction
- Updated the executeAutoSearch method to include an optional parameter for Skip.Keyword, allowing for conditional keyword extraction during web searches. - Implemented logic to extract keywords only when configured and not skipped, improving search query optimization. - Modified the Assistant's Stream method to pass options to executeAutoSearch, ensuring seamless integration of the new functionality. - Updated DESIGN.md to document the changes in keyword extraction logic and its impact on the search process.
This commit is contained in:
Max
parent
e38f6ecc96
commit
4dbce2f92e
5 changed files with 250 additions and 4 deletions
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
186
agent/assistant/search_auto_keyword_test.go
Normal file
186
agent/assistant/search_auto_keyword_test.go
Normal 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")
|
||||
})
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue