yao/agent/assistant/search.go

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] + "..."
}