diff --git a/agent/assistant/agent_next_test.go b/agent/assistant/agent_next_test.go index 8099c943..997b19b2 100644 --- a/agent/assistant/agent_next_test.go +++ b/agent/assistant/agent_next_test.go @@ -168,7 +168,9 @@ func TestAgentNextConditional(t *testing.T) { ctx := newAgentNextTestContext("test-conditional", "tests.realworld-next") messages := []context.Message{ - {Role: context.RoleUser, Content: "scenario: conditional - Task completed"}, + // Use conditional_success sub-scenario for deterministic behavior + // This avoids test flakiness caused by LLM response unpredictability + {Role: context.RoleUser, Content: "scenario: conditional_success - Task completed"}, } response, err := agent.Stream(ctx, messages) diff --git a/agent/search/DESIGN.md b/agent/search/DESIGN.md new file mode 100644 index 00000000..e2b770bc --- /dev/null +++ b/agent/search/DESIGN.md @@ -0,0 +1,2419 @@ +# Search Module Design + +## Overview + +The Search module provides a unified RAG (Retrieval-Augmented Generation) interface for Yao Agent, supporting three search types: + +| Type | Source | Use Case | +| ----- | -------------- | ---------------------------------------------------- | +| `web` | Internet | Real-time information, news, external knowledge | +| `kb` | Knowledge Base | Documents, FAQs, internal knowledge (vector + graph) | +| `db` | Database | Structured data from Yao Models (QueryDSL) | + +The module follows the **Handler + Registry** pattern consistent with the `content` module, and exposes JSAPI for flexible usage in Create/Next hooks. + +## Key Features + +- **Unified JSAPI**: `ctx.search.Web()`, `ctx.search.KB()`, `ctx.search.DB()`, `ctx.search.Parallel()` +- **Citation System**: Auto-generate citation IDs (`#ref:xxx`) for LLM reference +- **Real-time Output**: Stream search progress to client +- **Trace Integration**: Report search operations to user for transparency +- **Reranking**: Builtin, Agent, or MCP-based result reranking +- **Graceful Degradation**: Search errors don't block agent flow + +## Quick Start + +```typescript +// In Create hook (assistants/my-assistant/index.ts) +function Create(ctx, messages, options) { + const query = messages[messages.length - 1].content; + + // Simple web search + const result = ctx.search.Web(query, { limit: 5 }); + + // Or parallel search across all sources + const [web, kb, db] = ctx.search.Parallel([ + { type: "web", query, limit: 5 }, + { type: "kb", query, collections: ["docs"] }, + { type: "db", query, models: ["product"] }, + ]); + + return { + messages: [{ role: "system", content: formatContext(web, kb, db) }], + uses: { search: "disabled" }, // Disable auto search since hook handled it + }; +} +``` + +## Goals + +1. **Unified Interface**: Single API for web, knowledge base, and database search +2. **Flexibility**: Support built-in handlers and external tools (MCP/Agent delegation) +3. **JSAPI Support**: Enable search calls from Create/Next hooks via JavaScript +4. **Parallel Execution**: Support concurrent web + KB + DB searches +5. **Graceful Degradation**: Search failures should not block the main agent flow +6. **Real-time Feedback**: Stream search progress and results to users via output +7. **Traceability**: Report search operations to users for transparency +8. **Citation Support**: Enable LLM to reference search results with trackable citations + +## Architecture + +### Search Flow Diagram + +```mermaid +flowchart TD + A[Stream Start] --> B{Uses.Search?} + B -->|disabled| C[Skip Search] + B -->|builtin/agent/mcp| D{Hook Handled?} + D -->|"Yes (uses.search=disabled)"| C + D -->|No| E[Auto Search] + + E --> F{Check Assistant Config} + F --> G[Web Search] + F --> H[KB Search] + F --> I[DB Search] + + G --> J[Parallel Execute] + H --> J + I --> J + + J --> K[Merge Results] + K --> L[Rerank] + L --> M[Generate Citations] + M --> N[Inject to System Prompt] + + C --> O[LLM Call] + N --> O + O --> P[Output with Citations] +``` + +### Integration in Stream() + +```mermaid +sequenceDiagram + participant Client + participant Stream + participant CreateHook + participant Search + participant LLM + participant Output + + Client->>Stream: Stream(ctx, messages, options) + Stream->>Stream: Initialize + + alt Has Create Hook + Stream->>CreateHook: Create(ctx, messages, options) + CreateHook-->>Stream: response (may include search results) + end + + alt Uses.Search != "disabled" AND not handled by Hook + Stream->>Search: AutoSearch(ctx, messages) + Search->>Search: Web/KB/DB in parallel + Search->>Search: Rerank & Citations + Search->>Output: search_start, search_result, search_complete + Search-->>Stream: Inject search context to messages + end + + Stream->>LLM: Execute with search context + LLM->>Output: Stream response with #ref:xxx + Stream-->>Client: Complete +``` + +### Directory Structure + +``` +agent/search/ +├── DESIGN.md # This document +├── search.go # Main Searcher implementation and public API +├── registry.go # Handler registry (manages web/kb/db handlers) +├── jsapi.go # JavaScript API bindings for hooks +├── trace.go # Trace node creation and management +├── output.go # Real-time output/streaming to client +├── citation.go # Citation ID generation and tracking +├── reference.go # Reference building and LLM context formatting +│ +├── types/ # Type definitions (no dependencies on other search packages) +│ ├── types.go # Core types (SearchType, Request, Result, ResultItem, etc.) +│ ├── config.go # Configuration types (Config, CitationConfig, WeightsConfig, etc.) +│ ├── reference.go # Reference type for unified context protocol +│ └── graph.go # Graph-related types (GraphNode) +│ +├── interfaces/ # Interface definitions (depends only on types/) +│ ├── handler.go # Handler interface +│ ├── searcher.go # Searcher interface (public API) +│ ├── reranker.go # Reranker interface +│ └── nlp.go # NLP interfaces (KeywordExtractor, QueryDSLGenerator) +│ +├── rerank/ # Result reranking implementations +│ ├── rerank.go # Reranker factory and common logic +│ ├── builtin.go # Built-in score-based reranking (default) +│ ├── agent.go # Agent-based reranking (delegate to another assistant) +│ └── mcp.go # MCP-based reranking (call MCP server tool) +│ +├── nlp/ # Natural language processing for search +│ ├── nlp.go # NLP factory and common logic +│ ├── keyword.go # Keyword extraction for web search +│ └── querydsl.go # QueryDSL generation for DB search +│ # Note: Embedding follows KB collection config, not in this package +│ +├── handlers/ # Search handler implementations +│ ├── web/ # Web search +│ │ ├── handler.go # Web search handler (mode dispatch) +│ │ ├── tavily.go # Tavily provider (builtin) +│ │ ├── serper.go # Serper provider (builtin) +│ │ ├── agent.go # Agent mode (AI Search) +│ │ └── mcp.go # MCP mode (external service) +│ │ +│ ├── kb/ # Knowledge base search +│ │ ├── handler.go # KB search handler +│ │ ├── vector.go # Vector similarity search +│ │ └── graph.go # Graph-based association (GraphRAG) +│ │ +│ └── db/ # Database search (Yao Model/QueryDSL) +│ ├── handler.go # DB search handler +│ ├── query.go # QueryDSL builder +│ └── schema.go # Model schema introspection +│ +└── defaults/ # Default configuration values + └── defaults.go # System built-in defaults (used by agent/load.go) +``` + +### Dependency Graph + +``` + ┌─────────────┐ + │ types/ │ ← No internal dependencies + └──────┬──────┘ + │ + ┌──────▼──────┐ + │ interfaces/ │ ← Depends only on types/ + └──────┬──────┘ + │ + ┌─────────────────┼─────────────────┐ + │ │ │ + ┌─────▼─────┐ ┌──────▼──────┐ ┌──────▼──────┐ + │ rerank/ │ │ nlp/ │ │ defaults/ │ + └─────┬─────┘ └──────┬──────┘ └──────┬──────┘ + │ │ │ + └────────┬────────┴────────┬────────┘ + │ │ + ┌──────▼──────┐ ┌──────▼──────┐ + │ handlers/ │ │ (root pkg) │ + │ web/kb/db │ │ search.go │ + └──────┬──────┘ │ registry │ + │ │ jsapi, etc │ + └────┬─────┴─────────────┘ + │ + ┌─────▼─────┐ + │ External │ + │ Packages │ + └───────────┘ +``` + +### Package Import Rules + +1. **`types/`** - Zero internal dependencies, only stdlib and external packages +2. **`interfaces/`** - Imports only `types/` +3. **`rerank/`**, **`nlp/`**, **`defaults/`** - Import `types/` and `interfaces/` +4. **`handlers/*`** - Import `types/`, `interfaces/`, and may use `nlp/` for NL processing +5. **Root package** - Imports all sub-packages, provides public API + +### Main Searcher Implementation (`search.go`) + +Configuration is loaded by `agent/load.go` (global) and `agent/assistant/load.go` (assistant-level), following the existing pattern. The Search package directly uses the loaded configuration. + +```go +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 interfaces.Reranker + citation *CitationGenerator +} + +// SearchUses contains the search-specific uses configuration +// These are extracted from context.Uses and search config +type SearchUses struct { + Search string // "builtin", "disabled", "", "mcp:." + Web string // "builtin", "", "mcp:." + Keyword string // "builtin", "", "mcp:." + QueryDSL string // "builtin", "", "mcp:." + Rerank string // "builtin", "", "mcp:." +} + +// 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 *SearchUses) *Searcher { + return &Searcher{ + config: cfg, + handlers: map[types.SearchType]interfaces.Handler{ + types.SearchTypeWeb: web.NewHandler(uses.Web, cfg.Web), + types.SearchTypeKB: kb.NewHandler(cfg.KB), // KB always builtin + types.SearchTypeDB: db.NewHandler(uses.QueryDSL, cfg.DB), + }, + reranker: rerank.NewReranker(uses.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(ctx, 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 { + 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 +} + +// SearchMultiple executes multiple searches in parallel +func (s *Searcher) SearchMultiple(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 +} + +// BuildReferences converts search results to unified Reference format +func (s *Searcher) BuildReferences(results []*types.Result) []*types.Reference { + var refs []*types.Reference + for _, result := range results { + for _, item := range result.Items { + 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 +} +``` + +### Registry (`registry.go`) + +```go +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 +} +``` + +## Core Interfaces + +All interfaces are defined in `search/interfaces/` package to prevent circular dependencies. + +### Handler Interface (`interfaces/handler.go`) + +```go +package interfaces + +import ( + "github.com/yaoapp/yao/agent/context" + "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 + + // CanHandle checks if this handler can process the given request + CanHandle(ctx *context.Context, req *types.Request) bool + + // Search executes the search and returns results + Search(ctx *context.Context, req *types.Request) (*types.Result, error) +} +``` + +### Searcher Interface (`interfaces/searcher.go`) + +```go +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) + + // SearchMultiple executes multiple searches (potentially in parallel) + SearchMultiple(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 +} +``` + +### NLP Interfaces (`interfaces/nlp.go`) + +```go +package interfaces + +import ( + "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 + 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(ctx *context.Context, query string, schemas []*types.ModelSchema) (*types.QueryDSL, error) +} + +// Note: Embedding is handled by KB collection's own config (embedding provider + model), +// not defined here. See KB handler for details. +``` + +### Reranker Interface (`interfaces/reranker.go`) + +```go +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 + Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) +} +``` + +## Types + +All types are defined in `search/types/` package to prevent circular dependencies. + +### Core Types (`types/types.go`) + +```go +package types + +// SearchType represents the type of search +type SearchType string + +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 + +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 []QueryWhere `json:"wheres,omitempty"` // Pre-defined filters (optional) + Orders []QueryOrder `json:"orders,omitempty"` // Sort orders (optional) + Select []string `json:"select,omitempty"` // Fields to return (optional) + + // Reranking + Rerank *RerankOptions `json:"rerank,omitempty"` +} + +// QueryWhere represents a filter condition for DB search +type QueryWhere struct { + Field string `json:"field"` // Field name + Op string `json:"op,omitempty"` // Operator: "=", "like", ">", "<", "in", etc. (default: "=") + Value interface{} `json:"value"` // Filter value +} + +// QueryOrder represents a sort order for DB search +type QueryOrder struct { + Field string `json:"field"` // Field name + Order string `json:"order,omitempty"` // "asc" or "desc" (default: "desc") +} + +// 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 *QueryDSL `json:"dsl,omitempty"` // For DB search +} + +// QueryDSL represents a Yao QueryDSL for database search +type QueryDSL struct { + Model string `json:"model"` // Target model + Select []string `json:"select,omitempty"` // Fields to return + Wheres []QueryWhere `json:"wheres,omitempty"` // Filter conditions + Orders []QueryOrder `json:"orders,omitempty"` // Sort orders + Limit int `json:"limit,omitempty"` // Max results +} + +// ModelSchema represents a Yao Model schema for DSL generation +type ModelSchema struct { + ID string `json:"id"` // Model ID + Name string `json:"name"` // Model name + Description string `json:"description"` // Model description + Fields []FieldSchema `json:"fields"` // Field definitions +} + +// FieldSchema represents a field in the model schema +type FieldSchema struct { + Name string `json:"name"` // Field name + Type string `json:"type"` // Field type + Description string `json:"description"` // Field description + Searchable bool `json:"searchable"` // Whether field is searchable +} +``` + +### Graph Types (`types/graph.go`) + +```go +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"` +} +``` + +### Reference Types (`types/reference.go`) + +```go +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 XML + Prompt string `json:"prompt"` // Citation instruction prompt +} +``` + +### Configuration Types (`types/config.go`) + +```go +package types + +// Config represents the complete search configuration +type Config struct { + Web *WebConfig `json:"web,omitempty"` + KB *KBConfig `json:"kb,omitempty"` + DB *DBConfig `json:"db,omitempty"` + Keyword *KeywordConfig `json:"keyword,omitempty"` + QueryDSL *QueryDSLConfig `json:"querydsl,omitempty"` + Rerank *RerankConfig `json:"rerank,omitempty"` + Citation *CitationConfig `json:"citation,omitempty"` + Weights *WeightsConfig `json:"weights,omitempty"` + Options *OptionsConfig `json:"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"` // "tavily" or "serper" (for builtin mode) + APIKeyEnv string `json:"api_key_env,omitempty"` // Environment variable for API key + MaxResults int `json:"max_results,omitempty"` // Max results (default: 10) +} + +// KBConfig for knowledge base search settings +type KBConfig struct { + Collections []string `json:"collections,omitempty"` // Default collections + Threshold float64 `json:"threshold,omitempty"` // Similarity threshold (default: 0.7) + Graph bool `json:"graph,omitempty"` // Enable GraphRAG (default: false) +} + +// DBConfig for database search settings +type DBConfig struct { + Models []string `json:"models,omitempty"` // Default models + MaxResults int `json:"max_results,omitempty"` // Max results (default: 20) +} + +// KeywordConfig for keyword extraction +type KeywordConfig struct { + MaxKeywords int `json:"max_keywords,omitempty"` // Max keywords (default: 10) + Language string `json:"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"` // Fail if generation fails (default: false) +} + +// RerankConfig for reranking +type RerankConfig struct { + TopN int `json:"top_n,omitempty"` // Return top N (default: 10) +} + +// CitationConfig for citation format +type CitationConfig struct { + Format string `json:"format,omitempty"` // Default: "#ref:{id}" + AutoInjectPrompt bool `json:"auto_inject_prompt,omitempty"` // Auto-inject prompt (default: true) + CustomPrompt string `json:"custom_prompt,omitempty"` // Custom prompt template +} + +// WeightsConfig for source weighting +type WeightsConfig struct { + User float64 `json:"user,omitempty"` // User-provided (default: 1.0) + Hook float64 `json:"hook,omitempty"` // Hook results (default: 0.8) + Auto float64 `json:"auto,omitempty"` // Auto search (default: 0.6) +} + +// OptionsConfig for search behavior +type OptionsConfig struct { + SkipThreshold int `json:"skip_threshold,omitempty"` // Skip auto search if user provides >= N results +} +``` + +### Note on Reranker + +Reranker type is determined by `uses.rerank` in `agent/agent.yml`: + +- `"builtin"` - Simple score-based sorting +- `""` - Delegate to an assistant (Agent) +- `"mcp:."` - Call MCP tool (e.g., `"mcp:my-server.rerank"`) + +## Citation System + +Each search result has a unique `CitationID` for LLM reference. Citation logic is implemented in `search/citation.go`. + +### Citation ID Generation + +Citation IDs are generated sequentially: `ref_001`, `ref_002`, etc. + +```go +// citation.go +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) +} +``` + +### Citation Config (in `types/config.go`) + +```go +type CitationConfig struct { + Format string `json:"format,omitempty"` // Default: "#ref:{id}" + AutoInjectPrompt bool `json:"auto_inject_prompt,omitempty"` // Auto-add instructions to system prompt + CustomPrompt string `json:"custom_prompt,omitempty"` // Override default prompt template +} +``` + +### Default Citation Prompt + +When `AutoInjectPrompt` is enabled (default), the system prompt includes: + +``` +You have access to reference data in tags. Each 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: +[{id}] + +Example: According to the product data[ref_001], the price is $999. +``` + +### Custom Prompt in Config + +```yaml +# assistants/my-assistant.yml +search: + citation: + format: "[{id}]" + auto_inject_prompt: true + custom_prompt: "Cite using [{id}]. Sources: ..." +``` + +## Trace Integration + +Search operations create trace nodes to report execution details to users, providing transparency about what the agent is doing. + +### Trace Node Structure + +``` +search (type: "search") +├── query // Original query +├── search_type // "web", "kb", or "db" +├── duration_ms +├── status // "success", "failed" +├── result_count +└── children // Sub-operations + ├── embedding (kb only) + ├── vector_search (kb only) + ├── graph_search (kb, if enabled) + ├── querydsl_build (db only) + ├── db_query (db only) + └── rerank (if enabled) +``` + +## Real-time Output + +Search progress is streamed to the client via the output system. + +### Output Message Types + +```go +const ( + TypeSearchStart = "search_start" // Search initiated + TypeSearchResult = "search_result" // Result item (streamed) + TypeSearchComplete = "search_complete" // Search completed +) +``` + +### Client Display Example + +``` +🔍 Searching "latest AI developments"... + +📄 Found 5 results: + 1. #ref:a1b2 - OpenAI Announces GPT-5 + 2. #ref:c3d4 - Google's New AI Model + ... + +✅ Search complete (1.2s) +``` + +## JSAPI Integration + +The Search module is exposed via `ctx.search` object in hook scripts. + +### API Methods + +```typescript +// In hook scripts (index.ts) + +// Web search +ctx.search.Web(query: string, options?: WebOptions): Result + +// Knowledge base search +ctx.search.KB(query: string, options?: KBOptions): Result + +// Database search (Yao Model/QueryDSL) +ctx.search.DB(query: string, options?: DBOptions): Result + +// Parallel search (multiple types) +ctx.search.Parallel(requests: Request[]): Result[] +``` + +### Options Types + +```typescript +interface WebOptions { + limit?: number; // Max results (default: 10) + sites?: string[]; // Restrict to sites + timeRange?: string; // "day", "week", "month", "year" + rerank?: RerankOptions; +} + +interface KBOptions { + collections?: string[]; // Collection IDs + threshold?: number; // Similarity threshold (0-1) + limit?: number; // Max results + graph?: boolean; // Enable graph association + rerank?: RerankOptions; +} + +interface DBOptions { + models?: string[]; // Model IDs (default: use assistant's db.models) + wheres?: QueryWhere[]; // Pre-defined filters + orders?: QueryOrder[]; // Sort orders + select?: string[]; // Fields to return + limit?: number; // Max results (default: 10) + rerank?: RerankOptions; +} + +interface QueryWhere { + field: string; + op?: string; // "=", "like", ">", "<", "in", etc. + value: any; +} + +interface QueryOrder { + field: string; + order?: string; // "asc" or "desc" +} + +interface RerankOptions { + topN?: number; // Return top N after reranking + // Note: Reranker type is determined by uses.rerank in agent/agent.yml +} +``` + +### Usage Examples + +#### Example 1: Web Search + +```typescript +function Create(ctx, messages, options) { + const query = messages[messages.length - 1].content; + + const result = ctx.search.Web(query, { + limit: 5, + timeRange: "week", + }); + + if (result.items.length > 0) { + return { + messages: [ + { + role: "system", + content: formatSearchContext(result), + }, + ], + uses: { search: "disabled" }, // Disable auto search + }; + } + + return { messages: [] }; // Let auto search handle it +} +``` + +#### Example 2: Knowledge Base Search with Graph + +```typescript +function Create(ctx, messages, options) { + const query = messages[messages.length - 1].content; + + const result = ctx.search.KB(query, { + collections: ["docs", "faq"], + threshold: 0.7, + limit: 10, + graph: true, // Enable graph association + }); + + if (result.items.length > 0) { + return { + messages: [ + { + role: "system", + content: formatKBContext(result), + }, + ], + uses: { search: "disabled" }, // Disable auto search + }; + } + + return { messages: [] }; // Let auto search handle it +} +``` + +#### Example 3: Database Search + +```typescript +function Create(ctx, messages, options) { + const query = messages[messages.length - 1].content; + + // Search in assistant's models (uses db.models from assistant config) + const result = ctx.search.DB(query, { + models: ["product", "agents.mybot.order"], // Optional: override models + wheres: [{ field: "status", value: "active" }], // Pre-filter + limit: 20, + }); + + if (result.items.length > 0) { + return { + messages: [ + { + role: "system", + content: formatDBContext(result), + }, + ], + uses: { search: "disabled" }, // Disable auto search + }; + } + + return { messages: [] }; // Let auto search handle it +} +``` + +#### Example 4: Parallel Web + KB + DB Search + +```typescript +function Create(ctx, messages, options) { + const query = messages[messages.length - 1].content; + + // Execute web, KB, and DB search in parallel + const [webResult, kbResult, dbResult] = ctx.search.Parallel([ + { type: "web", query: query, limit: 5 }, + { type: "kb", query: query, collections: ["docs"], limit: 10 }, + { type: "db", query: query, models: ["product"], limit: 10 }, + ]); + + // Merge results + const context = mergeSearchResults(webResult, kbResult, dbResult); + + return { + messages: [ + { + role: "system", + content: context, + }, + ], + uses: { search: "disabled" }, // Disable auto search + }; +} +``` + +#### Example 5: Custom Citation Format + +```typescript +function Create(ctx, messages, options) { + const query = messages[messages.length - 1].content; + const result = ctx.search.Web(query, { limit: 5 }); + + // Build custom citation prompt + const refs = result.items + .map((item, i) => `[${i + 1}] ${item.title} - ${item.url}`) + .join("\n"); + + return { + messages: [ + { + role: "system", + content: `Use [N] to cite. References:\n${refs}`, + }, + ], + uses: { search: "disabled" }, // Disable auto search + citation: { autoInjectPrompt: false }, // Override citation config + }; +} +``` + +## Configuration + +Configuration follows a three-layer hierarchy (later overrides earlier): + +1. **System Built-in Defaults** - Hardcoded sensible defaults +2. **Global Configuration** - `agent/agent.yml` (uses) + `agent/search.yao` (search options) +3. **Assistant Configuration** - `assistants//package.yao` (uses + search options) + +### Uses Configuration + +Processing tools are configured in `agent/agent.yml` under `uses`: + +```yaml +# agent/agent.yml +uses: + default: "yaobots" + title: "workers.system.title" + vision: "workers.system.vision" + fetch: "workers.system.fetch" + + # Search processing tools (NLP) + keyword: "builtin" # "builtin", "workers.nlp.keyword", "mcp:my-server.extract_keywords" + querydsl: "builtin" # "builtin", "workers.nlp.querydsl", "mcp:my-server.generate_dsl" + rerank: "builtin" # "builtin", "workers.rerank", "mcp:my-server.rerank" + + # Search handlers + web: "builtin" # "builtin", "workers.search.web", "mcp:my-server.web_search" + # Note: kb & db always use builtin (access internal data) + # Note: embedding & entity follow KB collection config +``` + +Tool format: `"builtin"`, `""` (Agent), `"mcp:."` (MCP Tool) + +**Web Search Modes:** + +| Mode | Example | Description | +| --------- | ---------------------------- | -------------------------------------------------------------------------- | +| `builtin` | `"builtin"` | Use built-in providers (Tavily, Serper) | +| Agent | `"workers.search.web"` | AI-powered search: understand intent → optimize query → search → summarize | +| MCP | `"mcp:my-server.web_search"` | External search tool via MCP protocol | + +**Why Agent for Web Search (AI Search)?** + +When `uses.web` is set to an assistant ID, the search flow becomes: + +``` +User Query: "What's the best laptop for programming in 2024?" + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Agent (workers.search.web) │ +│ 1. Understand intent: laptop recommendations for coding │ +│ 2. Generate optimized queries: │ +│ - "best programming laptop 2024 review" │ +│ - "developer laptop comparison 2024" │ +│ 3. Execute multiple searches │ +│ 4. Analyze & deduplicate results │ +│ 5. Return structured, relevant results │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +High-quality, intent-aware search results +``` + +### System Built-in Defaults (`defaults/defaults.go`) + +These are the hardcoded defaults, used by `agent/load.go` when loading configuration: + +```go +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 +func GetWeight(cfg *types.Config, source types.SourceType) float64 { + if cfg == nil || cfg.Weights == nil { + switch source { + case types.SourceUser: + return 1.0 + case types.SourceHook: + return 0.8 + default: + return 0.6 + } + } + switch source { + case types.SourceUser: + return cfg.Weights.User + case types.SourceHook: + return cfg.Weights.Hook + case types.SourceAuto: + return cfg.Weights.Auto + default: + return 0.6 + } +} +``` + +### Configuration Loading (in `agent/load.go`) + +Configuration loading follows the existing pattern in `agent/load.go`: + +```go +// agent/load.go + +import ( + searchDefaults "github.com/yaoapp/yao/agent/search/defaults" + searchTypes "github.com/yaoapp/yao/agent/search/types" +) + +var searchConfig *searchTypes.Config + +// initSearchConfig initialize the search configuration from agent/search.yao +func initSearchConfig() error { + // Start with system defaults + searchConfig = searchDefaults.SystemDefaults + + path := filepath.Join("agent", "search.yao") + if exists, _ := application.App.Exists(path); !exists { + return nil // Use defaults + } + + // Read and merge with defaults + bytes, err := application.App.Read(path) + if err != nil { + return err + } + + var cfg searchTypes.Config + err = application.Parse("search.yao", bytes, &cfg) + if err != nil { + return err + } + + // Merge: defaults < global config + searchConfig = mergeSearchConfig(searchDefaults.SystemDefaults, &cfg) + return nil +} + +// GetSearchConfig returns the global search configuration +func GetSearchConfig() *searchTypes.Config { + return searchConfig +} +``` + +### Assistant-level Config Merge (in `agent/assistant/load.go`) + +Assistant-specific search config is merged in `assistant/load.go`: + +```go +// agent/assistant/load.go + +// GetMergedSearchConfig returns merged search config for this assistant +func (ast *Assistant) GetMergedSearchConfig() *searchTypes.Config { + globalCfg := agent.GetSearchConfig() + if ast.Search == nil { + return globalCfg + } + // Merge: global < assistant + return mergeSearchConfig(globalCfg, ast.Search.ToConfig()) +} +``` + +### Global Configuration + +`agent/search.yao` - Override system defaults for all assistants: + +```jsonc +{ + // Web search settings + "web": { + "provider": "tavily", // "tavily", "serper" (builtin providers only) + "api_key_env": "TAVILY_API_KEY", + "max_results": 10 + }, + + // Knowledge base search settings + "kb": { + "threshold": 0.7, // Similarity threshold + "graph": false // Enable GraphRAG association + }, + + // Database search settings + "db": { + "max_results": 20 + }, + + // Keyword extraction options (uses.keyword) + "keyword": { + "max_keywords": 10, + "language": "auto" // "auto", "en", "zh", etc. + }, + + // QueryDSL generation options (uses.querydsl) + "querydsl": { + "strict": false // Strict mode: fail if generation fails + }, + + // Rerank options (uses.rerank) + "rerank": { + "top_n": 10 // Return top N results after reranking + }, + + // Citation format for LLM references + "citation": { + "format": "#ref:{id}", + "auto_inject_prompt": true // Auto-inject citation instructions to system prompt + }, + + // Source weighting for result merging + "weights": { + "user": 1.0, // User-provided DataContent (highest priority) + "hook": 0.8, // Hook ctx.search.*() results + "auto": 0.6 // Auto search results + }, + + // Search behavior options + "options": { + "skip_threshold": 5 // Skip auto search if user provides >= N results + } +} +``` + +### Assistant Configuration + +`assistants//package.yao` - Override for specific assistant: + +```jsonc +{ + "name": "My Assistant", + "connector": "openai", + + // Overrides global uses (agent/agent.yml) + "uses": { + "search": "builtin", // "builtin", "disabled", "", "mcp:." + "web": "builtin", // "builtin", "", "mcp:." + "keyword": "workers.nlp.keyword", // Use LLM for keyword extraction + "querydsl": "workers.nlp.querydsl", // Use LLM for QueryDSL generation + "rerank": "mcp:my-server.rerank" // Use MCP tool for reranking + }, + + // Search configuration (overrides agent/search.yao) + "search": { + // Overrides global web settings + "web": { + "provider": "tavily", + "max_results": 5 + }, + + // Overrides global kb settings + "kb": { + "collections": ["docs", "faq"], // Specific collections to search + "threshold": 0.7, + "graph": true + }, + + // Overrides global db settings + "db": { + "models": ["product", "order"], // Uses db.models if not set + "max_results": 20 + }, + + // Overrides global keyword options + "keyword": { + "max_keywords": 5 + }, + + // Overrides global querydsl options + "querydsl": { + "strict": true + }, + + // Overrides global rerank options + "rerank": { + "top_n": 5 + }, + + // Overrides global citation settings + "citation": { + "format": "#ref:{id}", + "auto_inject_prompt": true + } + }, + + // Knowledge base collections available to this assistant + "kb": { + "collections": ["docs", "faq"] + }, + + // Database models available to this assistant + "db": { + "models": ["product", "order", "customer"] + } +} +``` + +## Execution Flow + +### Search Flow + +## Execution Modes + +### Stream() Execution with Search + +``` +Stream(ctx, messages, options) + │ + ├── 1. Initialize + │ + ├── 2. Create Hook (optional) + │ └── Can call ctx.search.* and return search results + │ + ├── 3. Auto Search Decision + │ ├── 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 + │ + ├── 4. LLM Call (with search context if any) + │ + ├── 5. Next Hook (optional) + │ + └── 6. Output (response may contain #ref:xxx citations) +``` + +### Control via Uses.Search + +Search is controlled via the `Uses` mechanism, following the merge hierarchy: + +``` +Global (agent/agent.yml) → Assistant (package.yao) → CreateHook (return uses) → Request (options.uses) +``` + +| Uses.Search | Behavior | +| ----------------------- | ------------------------------------ | +| `"builtin"` | Use builtin auto search | +| `"disabled"` | Disable auto search | +| `""` | Delegate to an assistant (AI Search) | +| `"mcp:."` | Use MCP tool for search | +| `undefined` | Follow upper layer config (default) | + +**Go:** + +```go +// Use builtin auto search +uses := &context.Uses{Search: "builtin"} + +// Disable auto search +uses := &context.Uses{Search: "disabled"} + +// Delegate to AI Search assistant +uses := &context.Uses{Search: "workers.search.ai"} + +// Follow assistant config (default) +uses := &context.Uses{Search: ""} // or nil +``` + +**API Request:** + +```json +{ + "messages": [...], + "uses": { + "search": "builtin" + } +} +``` + +### Hook-Controlled Search + +Search is controlled via the `Uses` mechanism, same as Vision/Audio. The merge hierarchy is: + +``` +Global (agent/agent.yml) → Assistant (package.yao) → CreateHook (return uses) +``` + +When you need custom search logic, handle it in Create Hook and return `uses.search` to control auto search: + +```typescript +function Create(ctx, messages, options) { + const query = messages[messages.length - 1].content; + + // Custom logic: only search for certain queries + if (needsSearch(query)) { + const result = ctx.search.Web(query, { limit: 5 }); + return { + messages: [{ role: "system", content: formatContext(result) }], + uses: { search: "disabled" }, // Disable auto search (hook handled it) + }; + } + + // Let auto search handle it (follow assistant config) + return { messages: [] }; +} +``` + +**Uses.Search Values:** + +| Value | Behavior | +| ----------------------- | ------------------------------------ | +| `"builtin"` | Use builtin auto search | +| `"disabled"` | Disable auto search | +| `""` | Delegate to an assistant (AI Search) | +| `"mcp:."` | Use MCP tool for search | +| `undefined` | Follow upper layer config (default) | + +**Uses Merge Hierarchy:** + +``` +┌─────────────────────────────────────────────────────────────┐ +│ 1. Global Config (agent/agent.yml) │ +│ uses: │ +│ search: "builtin" │ +└─────────────────────────────────────────────────────────────┘ + ↓ merge +┌─────────────────────────────────────────────────────────────┐ +│ 2. Assistant Config (assistants//package.yao) │ +│ uses: │ +│ search: "workers.search.web" # Override to AI Search │ +└─────────────────────────────────────────────────────────────┘ + ↓ merge +┌─────────────────────────────────────────────────────────────┐ +│ 3. CreateHook Return │ +│ return { │ +│ uses: { search: "disabled" } # Hook handled it │ +│ } │ +└─────────────────────────────────────────────────────────────┘ +``` + +> **Note**: The `Uses` struct in `context/types_llm.go` already has a `Search` field. +> The value `"disabled"` is a special value to disable auto search when hook handles it. + +## Search Flow + +``` +Request → Trace Start → Query Process → Search → Rerank → Citations → Output → Return +``` + +### Query Processing + +| Type | Process | Tool Config | +| ---- | ----------------------------------------------------- | -------------------- | +| Web | Extract keywords → Build query | `uses.keyword` | +| KB | Get collection's embedding model → Generate embedding | KB collection config | +| DB | Parse query → Build QueryDSL → Execute on models | `uses.querydsl` | + +#### Processing Methods + +Configure via `uses.*` in `agent/agent.yml`: + +| Format | Description | Use Case | +| --------------------- | ----------------------------------------- | ------------------------------ | +| `builtin` | Rule-based, template-driven (no LLM call) | Fast, low cost, simple queries | +| `` | Delegate to an assistant (Agent) | LLM-based, custom logic | +| `mcp:.` | Call MCP tool | External services integration | + +#### Keyword Extraction (`nlp/keyword.go`) + +Configure via `uses.keyword`: + +```go +// nlp/keyword.go +package nlp + +import ( + "strings" + + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/types" +) + +// KeywordExtractor extracts keywords from user query +type KeywordExtractor struct { + usesKeyword string // "builtin", "", "mcp:." + config *types.KeywordConfig +} + +// NewKeywordExtractor creates a keyword extractor +func NewKeywordExtractor(usesKeyword string, cfg *types.KeywordConfig) *KeywordExtractor { + return &KeywordExtractor{usesKeyword: usesKeyword, config: cfg} +} + +// Extract extracts keywords from content +func (e *KeywordExtractor) Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) { + switch { + case e.usesKeyword == "builtin" || e.usesKeyword == "": + return e.builtinExtract(content, opts) + case strings.HasPrefix(e.usesKeyword, "mcp:"): + return e.mcpExtract(ctx, content, opts) + default: + return e.agentExtract(ctx, content, opts) + } +} +``` + +``` +"I want to find the best wireless headphones under $100" + ↓ builtin: simple tokenization + stopword removal + ↓ agent: LLM extracts ["wireless headphones", "under $100", "best"] +→ Keywords: ["wireless headphones", "under $100", "best"] +``` + +#### Embedding (KB Collection Config) + +Embedding is **not** part of the `nlp/` package. It follows KB collection's own configuration: + +- Each KB collection defines its own embedding provider and model +- The KB handler (`handlers/kb/`) calls the collection's embedding API directly +- Entity types for GraphRAG are also defined per collection + +```go +// handlers/kb/handler.go +func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { + // 1. Get collection config (embedding provider, model) + collection := h.getCollection(req.Collections[0]) + + // 2. Generate embedding using collection's config + vector, err := collection.Embed(ctx, req.Query) + + // 3. Vector search + // ... +} +``` + +#### QueryDSL Generation (`nlp/querydsl.go`) + +Configure via `uses.querydsl`: + +```go +// nlp/querydsl.go +package nlp + +import ( + "strings" + + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/types" +) + +// QueryDSLGenerator generates QueryDSL from natural language +type QueryDSLGenerator struct { + usesQueryDSL string // "builtin", "", "mcp:." + config *types.QueryDSLConfig +} + +// NewQueryDSLGenerator creates a QueryDSL generator +func NewQueryDSLGenerator(usesQueryDSL string, cfg *types.QueryDSLConfig) *QueryDSLGenerator { + return &QueryDSLGenerator{usesQueryDSL: usesQueryDSL, config: cfg} +} + +// Generate converts natural language to QueryDSL +func (g *QueryDSLGenerator) Generate(ctx *context.Context, query string, schemas []*types.ModelSchema) (*types.QueryDSL, error) { + switch { + case g.usesQueryDSL == "builtin" || g.usesQueryDSL == "": + return g.builtinGenerate(query, schemas) + case strings.HasPrefix(g.usesQueryDSL, "mcp:"): + return g.mcpGenerate(ctx, query, schemas) + default: + return g.agentGenerate(ctx, query, schemas) + } +} +``` + +``` +"Products cheaper than $100 from Apple" + ↓ builtin: template matching against model schema + ↓ agent: LLM generates DSL from NL + schema +→ QueryDSL: {"wheres": [{"column": "price", "op": "<", "value": 100}, {"column": "brand", "value": "Apple"}]} +``` + +## Handlers & Providers + +All handler implementations are in `search/handlers/` directory. + +### Web Search (`handlers/web/`) + +Web search supports three modes via `uses.web`: + +| Mode | Value | Description | +| ------- | ---------------------------- | ------------------------------------------- | +| Builtin | `"builtin"` | Direct API calls to Tavily/Serper | +| Agent | `"workers.search.web"` | AI-powered search with intent understanding | +| MCP | `"mcp:my-server.web_search"` | External search tool via MCP | + +```go +// handlers/web/handler.go +package web + +import ( + "strings" + + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/types" +) + +// Handler implements web search +type Handler struct { + usesWeb string // "builtin", "", "mcp:." + config *types.WebConfig +} + +// NewHandler creates a new web search handler +func NewHandler(usesWeb string, cfg *types.WebConfig) *Handler { + return &Handler{usesWeb: usesWeb, config: cfg} +} + +// Search executes web search based on uses.web mode +func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { + switch { + case h.usesWeb == "builtin" || h.usesWeb == "": + return h.builtinSearch(ctx, req) + case strings.HasPrefix(h.usesWeb, "mcp:"): + return h.mcpSearch(ctx, req) + default: + // Agent mode: delegate to assistant for AI-powered search + return h.agentSearch(ctx, req) + } +} + +// builtinSearch uses Tavily/Serper directly +func (h *Handler) builtinSearch(ctx *context.Context, req *types.Request) (*types.Result, error) { + var provider Provider + switch h.config.Provider { + case "tavily": + provider = NewTavilyProvider(h.config) + case "serper": + provider = NewSerperProvider(h.config) + } + return provider.Search(ctx, req) +} + +// agentSearch delegates to an assistant for AI-powered search +func (h *Handler) agentSearch(ctx *context.Context, req *types.Request) (*types.Result, error) { + // 1. Call assistant with search request + // 2. Assistant understands intent, generates optimized queries + // 3. Assistant executes searches (may call builtin internally) + // 4. Assistant analyzes and returns structured results + return nil, nil +} + +// mcpSearch calls external MCP tool +func (h *Handler) mcpSearch(ctx *context.Context, req *types.Request) (*types.Result, error) { + // Parse "mcp:server.tool" + mcpRef := strings.TrimPrefix(h.usesWeb, "mcp:") + parts := strings.SplitN(mcpRef, ".", 2) + if len(parts) != 2 { + return nil, fmt.Errorf("invalid MCP format, expected 'mcp:server.tool', got '%s'", h.usesWeb) + } + serverID, toolName := parts[0], parts[1] + // Call MCP tool + return nil, nil +} +``` + +**Built-in Providers (when `uses.web = "builtin"`):** + +| Provider | File | Notes | +| -------- | ----------- | ------------------------------- | +| Tavily | `tavily.go` | Recommended for AI applications | +| Serper | `serper.go` | Google search API | + +**Agent Mode (AI Search):** + +When `uses.web` is set to an assistant ID (e.g., `"workers.search.web"`), the assistant can: + +1. **Understand user intent** - Parse complex queries, identify what user really wants +2. **Generate multiple queries** - Create optimized search terms for better coverage +3. **Multi-source search** - Search multiple providers or sources +4. **Result analysis** - Deduplicate, rank, and summarize results +5. **Context-aware** - Use conversation context to improve search relevance + +``` +User Query: "What's the best laptop for programming in 2024?" + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Agent (workers.search.web) │ +│ 1. Understand intent: laptop recommendations for coding │ +│ 2. Generate optimized queries: │ +│ - "best programming laptop 2024 review" │ +│ - "developer laptop comparison 2024" │ +│ 3. Execute multiple searches via builtin providers │ +│ 4. Analyze & deduplicate results │ +│ 5. Return structured, relevant results │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +High-quality, intent-aware search results +``` + +**Example AI Search Assistant:** + +```typescript +// assistants/workers/search/web/src/index.ts +function Create(ctx, messages, options) { + const userQuery = messages[messages.length - 1].content; + + // 1. Analyze intent (this assistant has access to LLM) + const intent = analyzeIntent(ctx, userQuery); + + // 2. Generate optimized queries + const queries = generateQueries(intent); + + // 3. Execute searches using builtin provider + const allResults = []; + for (const q of queries) { + const result = ctx.search.Web(q, { + provider: "tavily", // Use builtin provider + limit: 5, + }); + allResults.push(...result.items); + } + + // 4. Merge, deduplicate, and rank results + const merged = mergeAndRank(allResults, intent); + + return { + type: "search_result", + items: merged, + }; +} +``` + +### Knowledge Base (`handlers/kb/`) + +```go +// handlers/kb/handler.go +package kb + +import ( + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/interfaces" + "github.com/yaoapp/yao/agent/search/types" +) + +// Handler implements KB search +type Handler struct { + config *types.KBConfig +} + +// NewHandler creates a new KB search handler +func NewHandler(cfg *types.KBConfig) *Handler { + return &Handler{config: cfg} +} + +// Search executes vector search and optional graph association +func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { + // 1. Generate embedding via query processor + // 2. Vector search in collections + // 3. Optional: Graph association (if req.Graph) + // 4. Return results + return nil, nil +} +``` + +| File | Description | +| ------------ | ---------------------------------- | +| `handler.go` | Main KB handler implementation | +| `vector.go` | Vector similarity search | +| `graph.go` | Graph-based association (GraphRAG) | + +### Database Search (`handlers/db/`) + +```go +// handlers/db/handler.go +package db + +import ( + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/interfaces" + "github.com/yaoapp/yao/agent/search/types" +) + +// Handler implements DB search +type Handler struct { + usesQueryDSL string // "builtin", "", "mcp:." + config *types.DBConfig +} + +// NewHandler creates a new DB search handler +func NewHandler(usesQueryDSL string, cfg *types.DBConfig) *Handler { + return &Handler{usesQueryDSL: usesQueryDSL, config: cfg} +} + +// Search converts NL to QueryDSL and executes +func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { + // 1. Get model schemas + // 2. Generate QueryDSL via query processor + // 3. Execute queries on models + // 4. Return results + return nil, nil +} +``` + +| File | Description | +| ------------ | ------------------------------ | +| `handler.go` | Main DB handler implementation | +| `query.go` | QueryDSL builder utilities | +| `schema.go` | Model schema introspection | + +Integrates with Yao's Model/QueryDSL system: + +- Natural language → QueryDSL conversion (via LLM) +- Model schema introspection for query building +- Support for: + - Global models (`models/*.mod.yao`) + - Assistant-specific models (`assistants/{id}/models/*.mod.yao` → `agents.{id}.*`) +- Permission-aware queries (respects `__yao_*` permission fields) + +### Reranking (`rerank/`) + +```go +// rerank/rerank.go +package rerank + +import ( + "github.com/yaoapp/yao/agent/search/interfaces" + "github.com/yaoapp/yao/agent/search/types" +) + +// NewReranker creates a reranker based on uses.rerank config +func NewReranker(usesRerank string) interfaces.Reranker { + switch { + case usesRerank == "builtin" || usesRerank == "": + return NewBuiltinReranker() + case strings.HasPrefix(usesRerank, "mcp:"): + // Parse "mcp:server.tool" + mcpRef := strings.TrimPrefix(usesRerank, "mcp:") + parts := strings.SplitN(mcpRef, ".", 2) + if len(parts) != 2 { + // Invalid format, fallback to builtin + return NewBuiltinReranker() + } + return NewMCPReranker(parts[0], parts[1]) // serverID, toolName + default: + // Assume it's an assistant ID + return NewAgentReranker(usesRerank) + } +} +``` + +| File | Description | +| ------------ | -------------------------------- | +| `rerank.go` | Factory and common logic | +| `builtin.go` | Simple score sorting (default) | +| `agent.go` | Delegate to an assistant (Agent) | +| `mcp.go` | Call MCP tool for reranking | + +Configure via `uses.rerank` in `agent/agent.yml`: + +| Value | Notes | +| ---------------------- | -------------------------------- | +| `builtin` | Simple score sorting (default) | +| `workers.rerank` | Delegate to an assistant (Agent) | +| `mcp:my-server.rerank` | Call MCP tool for reranking | + +## Error Handling + +Search errors don't block the agent flow. Errors are returned in `Result.Error`: + +```typescript +const result = ctx.search.Web(query); +if (result.error) { + // Handle gracefully or fallback + console.warn("Search failed:", result.error); +} +``` + +## Configuration Priority + +Configuration is merged with later layers overriding earlier ones: + +1. **System Built-in** - Hardcoded defaults (lowest priority) +2. **Global-level** - `agent/agent.yml` (uses) + `agent/search.yao` (search options) +3. **Assistant-level** - `assistants//package.yao` (uses + search) +4. **Hook-level** - CreateHook return `uses.search` value +5. **Request-level** - `options.uses.search` in Stream() call (highest priority) + - `"builtin"`: Use builtin auto search + - `"disabled"`: Disable auto search + - `""`: Delegate to AI Search assistant + - `"mcp:."`: Use MCP tool for search + +## DB Search Details + +### Query Processing Flow + +``` +Natural Language Query + │ + ▼ +┌─────────────────────────────────┐ +│ Get Model Schemas │ ← Introspect models from db.models config +│ (fields, types, relations) │ +└─────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────┐ +│ LLM: Generate QueryDSL │ ← Convert NL to Yao QueryDSL +│ (select, wheres, orders) │ +└─────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────┐ +│ Execute Query on Each Model │ ← model.Find() with QueryDSL +└─────────────────────────────────┘ + │ + ▼ + Results +``` + +### Model ID Formats + +| Format | Example | Description | +| ------ | -------------------- | --------------------------------------------------------------------- | +| Global | `product` | Global model from `models/product.mod.yao` | +| System | `__yao.user` | Yao system model | +| Agent | `agents.mybot.order` | Assistant-specific model from `assistants/mybot/models/order.mod.yao` | + +### QueryDSL Generation Prompt + +The DB handler uses LLM to convert natural language to QueryDSL: + +``` +Given the following model schemas: +- product: { id, name, price, category, status, created_at } +- order: { id, product_id, quantity, total, customer_id, status } + +User query: "find all active products under $100 in electronics category" + +Generate Yao QueryDSL: +{ + "model": "product", + "wheres": [ + { "field": "status", "op": "=", "value": "active" }, + { "field": "price", "op": "<", "value": 100 }, + { "field": "category", "op": "=", "value": "electronics" } + ], + "orders": [{ "field": "price", "order": "asc" }], + "limit": 10 +} +``` + +## Content Module Integration + +User messages may contain `type="data"` ContentParts with data source references. The `content` module processes these before LLM call. + +### DataSource Types (from `context/types.go`) + +```go +const ( + DataSourceModel DataSourceType = "model" // DB model query + DataSourceKBCollection DataSourceType = "kb_collection" // KB collection search + DataSourceKBDocument DataSourceType = "kb_document" // KB document retrieval + DataSourceTable DataSourceType = "table" // Direct table query + DataSourceAPI DataSourceType = "api" // External API + DataSourceMCPResource DataSourceType = "mcp_resource" // MCP resource +) +``` + +### Message with Data Reference + +User only specifies data source IDs. Filters are generated by Search module from natural language. + +```json +{ + "role": "user", + "content": [ + { "type": "text", "text": "Show me products under $100" }, + { + "type": "data", + "data": { + "sources": [ + { "type": "model", "name": "product" }, + { "type": "kb_collection", "name": "product-docs" } + ] + } + } + ] +} +``` + +The Search module will: + +1. Extract query from text: "products under $100" +2. For `model:product` → Generate QueryDSL: `{ "wheres": [{ "field": "price", "op": "<", "value": 100 }] }` +3. For `kb_collection:product-docs` → Vector search with query embedding + +### Source Weighting & LLM Context + +Search results carry `source` and `weight` fields, which are used to build weighted context for LLM. + +**Source Types:** + +| Source | Weight | Description | +| ------ | ------ | -------------------------------- | +| `user` | 1.0 | Explicitly referenced in message | +| `hook` | 0.8 | Called in Create/Next hook | +| `auto` | 0.6 | Triggered by assistant config | + +**ResultItem with Weight:** + +```go +type ResultItem struct { + CitationID string `json:"citation_id"` // "#ref:xxx" + Source string `json:"source"` // "user", "hook", "auto" + Weight float64 `json:"weight"` // 1.0, 0.8, 0.6 + Score float64 `json:"score"` // Relevance score + // ... other fields +} +``` + +### Unified Context Protocol + +All data sources (Content module, Hook, Auto-Search) produce the same `Reference` structure. The final LLM input uses a unified `` format. + +**Reference (Internal Structure):** + +```go +// Reference is the unified structure for all data sources +type Reference struct { + ID string `json:"id"` // Unique citation ID: "ref_001", "ref_002" + Type string `json:"type"` // "web", "kb", "db" + Source string `json:"source"` // "user", "hook", "auto" + Weight float64 `json:"weight"` // 1.0, 0.8, 0.6 + 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 +} +``` + +**Data Flow:** + +```mermaid +flowchart TD + subgraph Sources ["Data Sources"] + CM["Content Module
(db:xxx kb:xxx)"] + HS["Hook Search
ctx.search.*()"] + AS["Auto Search
(assistant config)"] + end + + CM -->|"source=user
weight=1.0"| REF + HS -->|"source=hook
weight=0.8"| REF + AS -->|"source=auto
weight=0.6"| REF + + REF["[]Reference
(Unified Structure)"] + REF --> MERGE["Merge & Deduplicate
Rerank by score × weight"] + MERGE --> BUILD["Build <references> XML"] + BUILD --> LLM["LLM Input"] +``` + +**LLM References Format:** + +```xml + + +Product: iPhone 15 Pro +Price: $999 +Category: Electronics + + +The iPhone 15 Pro features the A17 Pro chip with improved performance... +URL: https://example.com/iphone-review + + +Apple announced the iPhone 15 series in September 2023... +URL: https://news.example.com/apple-iphone-15 + + +``` + +**LLM System Prompt (auto-injected):** + +``` +You have access to reference data in tags. Each 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: +[{id}] + +Example: According to the product data[ref_001], the price is $999. +``` + +**Citation Output Format:** + +LLM outputs citations as HTML links that can be parsed and rendered by frontend: + +```html + +The iPhone 15 Pro[ref_001] +features the A17 Pro chip[ref_002]. +``` + +**Citation Link Attributes:** + +| Attribute | Description | Example | +| --------------- | ----------------------- | ----------------------- | +| `class` | Fixed class for styling | `"ref"` | +| `data-ref-id` | Reference ID | `"ref_001"` | +| `data-ref-type` | Data type | `"db"`, `"kb"`, `"web"` | +| `href` | Anchor link | `"#ref:ref_001"` | + +**Conversion Examples:** + +| Module | Input | Output Reference | +| ------- | ---------------------------------- | ---------------------------------------------- | +| Content | `db:product` (user message) | `{source:"user", weight:1.0, type:"db", ...}` | +| Content | `kb:docs` (user message) | `{source:"user", weight:1.0, type:"kb", ...}` | +| Hook | `ctx.search.Web(query)` | `{source:"hook", weight:0.8, type:"web", ...}` | +| Hook | `ctx.search.KB(query)` | `{source:"hook", weight:0.8, type:"kb", ...}` | +| Hook | `ctx.search.DB(query)` | `{source:"hook", weight:0.8, type:"db", ...}` | +| Auto | Assistant config `search.web=true` | `{source:"auto", weight:0.6, type:"web", ...}` | +| Auto | Assistant config `search.kb=true` | `{source:"auto", weight:0.6, type:"kb", ...}` | + +### Processing Flow + +``` +Stream() + │ + ├── 1. Collect search results from all sources + │ ├── User DataContent → source="user", weight=1.0 + │ ├── Hook ctx.search.*() → source="hook", weight=0.8 + │ └── Auto search → source="auto", weight=0.6 + │ + ├── 2. Merge, deduplicate, rerank by (score * weight) + │ + ├── 3. Build ... format + │ + └── 4. Inject references into messages for LLM +``` + +**Behavior Rules:** + +1. **User data sufficient**: If user provides enough data (≥ skip_threshold), skip auto search +2. **Deduplication**: Same record from different sources → keep highest weight version +3. **Final ranking**: Sort by `score * weight` after reranking + +**Configuration:** + +Global defaults (`agent/search.yao`): + +```jsonc +{ + "weights": { + "user": 1.0, // User-provided DataContent + "hook": 0.8, // Hook ctx.search.*() results + "auto": 0.6 // Auto search results + }, + "options": { + "skip_threshold": 5 // Skip auto search if user provides >= N results + } +} +``` + +Assistant-level override (`assistants//package.yao`): + +```jsonc +{ + "search": { + "weights": { + "user": 1.0, + "hook": 0.9, // Higher weight for hook results + "auto": 0.5 // Lower weight for auto results + }, + "options": { + "skip_threshold": 10 // Need more user results to skip auto search + } + } +} +``` + +**System Auto-Processing:** + +The weighting and context building is handled automatically by the system: + +``` +Stream() + │ + ├── 1. Parse user message for DataContent sources + │ └── If found → Mark as source="user", weight=1.0 + │ + ├── 2. Create Hook (optional) + │ └── If hook calls ctx.search.*() → Mark as priority=2, weight=0.8 + │ + ├── 3. Auto Search Decision + │ ├── Count user-provided results + │ ├── IF user_results >= skip_auto_if_user_results → SKIP auto search + │ └── ELSE → Execute auto search with priority=3, weight=0.6 + │ + ├── 4. Merge & Rerank (automatic) + │ ├── Collect all results with their weights + │ ├── Deduplicate (keep highest priority) + │ └── Calculate finalScore = baseScore * weight + │ + └── 5. Inject to LLM context +``` + +Users don't need to handle weights in hooks - the system manages this automatically. + +### Processing Flow in content.Vision() + +``` +content.Vision() + ├── type="text" → Pass through + ├── type="image_url" → Image processing + ├── type="file" → File processing + └── type="data" → processDataContent() + ├── DataSourceModel → Query via model.Find() → Format as text + ├── DataSourceKBCollection → search.KB() → Format as text + ├── DataSourceKBDocument → Retrieve document → Format as text + └── DataSourceMCPResource → MCP resource read → Format as text +``` + +### Implementation Location + +The `processDataContent()` function in `content/content.go` should: + +1. **For `model` type**: Call search module's DB handler or direct model query +2. **For `kb_collection` type**: Call search module's KB handler +3. **For `kb_document` type**: Retrieve specific document from KB +4. **For `mcp_resource` type**: Read MCP resource + +This allows the search module to be reused for both: + +- **Auto Search**: Triggered when `Uses.Search != "disabled"` +- **Data ContentPart**: User explicitly references data sources in message + +## Related Files + +### Internal Dependencies + +- `agent/search/types/` - All type definitions (no circular dependencies) +- `agent/search/interfaces/` - All interface definitions +- `agent/search/defaults/` - System default configuration values +- `agent/search/handlers/` - Handler implementations (web, kb, db) +- `agent/search/rerank/` - Reranker implementations +- `agent/search/nlp/` - NLP implementations (keyword, querydsl) + +### External Dependencies + +- `agent/context/jsapi.go` - JSAPI base implementation +- `agent/context/types.go` - DataSource, DataContent types +- `agent/context/types_llm.go` - Uses configuration (Search field) +- `agent/assistant/types.go` - SearchOption definition +- `agent/store/types/types.go` - KnowledgeBase, Database config +- `agent/output/message/types.go` - Output message types +- `agent/content/content.go` - Content processing (Vision function) +- `model/model.go` - Yao Model loading (global, system, assistant models) + +## See Also + +- `agent/context/JSAPI.md` - Full JSAPI documentation +- `agent/context/RESOURCE_MANAGEMENT.md` - Context lifecycle and resource management +- `agent/output/README.md` - Output system documentation +- `agent/store/CHAT_STORAGE_DESIGN.md` - Chat storage design