diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 12ed9669..cb659647 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -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) + 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, diff --git a/agent/assistant/search.go b/agent/assistant/search.go new file mode 100644 index 00000000..dd9a442a --- /dev/null +++ b/agent/assistant/search.go @@ -0,0 +1,275 @@ +package assistant + +import ( + "strings" + + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search" + 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 +func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) *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 + } + + // 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] + "..." +} diff --git a/agent/assistant/search_auto_disabled_test.go b/agent/assistant/search_auto_disabled_test.go new file mode 100644 index 00000000..5ab7ced6 --- /dev/null +++ b/agent/assistant/search_auto_disabled_test.go @@ -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)") + }) +} diff --git a/agent/assistant/search_auto_full_test.go b/agent/assistant/search_auto_full_test.go new file mode 100644 index 00000000..204d46d9 --- /dev/null +++ b/agent/assistant/search_auto_full_test.go @@ -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)") + }) +} diff --git a/agent/assistant/search_auto_hook_disable_test.go b/agent/assistant/search_auto_hook_disable_test.go new file mode 100644 index 00000000..fed6c8cd --- /dev/null +++ b/agent/assistant/search_auto_hook_disable_test.go @@ -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") + }) +} diff --git a/agent/assistant/search_auto_web_test.go b/agent/assistant/search_auto_web_test.go new file mode 100644 index 00000000..22ec0896 --- /dev/null +++ b/agent/assistant/search_auto_web_test.go @@ -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") + }) +} diff --git a/agent/search/DESIGN.md b/agent/search/DESIGN.md index b538f1d0..8b0cd4a3 100644 --- a/agent/search/DESIGN.md +++ b/agent/search/DESIGN.md @@ -1552,22 +1552,62 @@ Stream(ctx, messages, options) ├── 2. Create Hook (optional) │ └── Can call ctx.search.* and return search results │ - ├── 3. Auto Search Decision + ├── 3. BuildRequest + BuildContent + │ + ├── 4. Auto Search Decision (shouldAutoSearch) │ ├── IF Uses.Search == "disabled" → SKIP │ ├── IF Create Hook returned uses.search="disabled" → SKIP - │ └── ELSE → Execute Auto Search (based on Uses.Search mode) - │ ├── Read assistant's search config - │ ├── Execute web/kb/db in parallel - │ ├── Send search_start/search_result/search_complete to output - │ ├── Rerank results - │ ├── Generate citation IDs - │ └── Inject search context + citation prompt to messages + │ └── ELSE → Execute Auto Search (executeAutoSearch) + │ ├── Read assistant's search config (GetMergedSearchConfig) + │ ├── Build search requests (buildSearchRequests) + │ ├── Execute web/kb/db in parallel (searcher.All) + │ ├── Build reference context (BuildReferenceContext) + │ └── Inject search context to messages (injectSearchContext) │ - ├── 4. LLM Call (with search context if any) + ├── 5. LLM Call (with search context if any) │ - ├── 5. Next Hook (optional) + ├── 6. Next Hook (optional) │ - └── 6. Output (response may contain #ref:xxx citations) + └── 7. Output (response may contain #ref:xxx citations) +``` + +**Implementation Files:** + +| File | Description | +| --------------------- | ----------------------------------------------- | +| `assistant/search.go` | Core integration logic (shouldAutoSearch, etc.) | +| `assistant/agent.go` | Stream() integration point (after BuildContent) | +| `search/reference.go` | BuildReferenceContext, FormatReferencesXML | + +**Key Functions (`assistant/search.go`):** + +```go +// shouldAutoSearch determines if auto search should be executed +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 + +// injectSearchContext injects search results into messages +func (ast *Assistant) injectSearchContext(messages []context.Message, refCtx *searchTypes.ReferenceContext) []context.Message + +// getMergedSearchUses returns the merged uses configuration for search +func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResponse) *context.Uses + +// buildSearchRequests builds search requests based on assistant configuration +func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Config) []*searchTypes.Request +``` + +**Integration in agent.go:** + +```go +// In Stream(), after BuildContent: +if ast.shouldAutoSearch(ctx, createResponse) { + refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse) + if refCtx != nil && len(refCtx.References) > 0 { + completionMessages = ast.injectSearchContext(completionMessages, refCtx) + } +} ``` ### Control via Uses.Search