opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Initialize Search JSAPI Factory and Update Documentation

Browse source 
- Added initialization for the Search JSAPI factory in the Assistant module to streamline search operations.
- Updated the DESIGN.md to reflect the new architecture of the search module, including detailed descriptions of the keyword extraction and QueryDSL generation processes.
- Enhanced the interfaces for keyword extraction to require context, improving flexibility for different extraction modes.
- Revised directory structures in the documentation to clarify the organization of search-related components, ensuring comprehensive guidance for developers.
This commit is contained in:
Max 2025-12-13 14:41:03 +08:00
parent b3cf5a09d9
commit d3865d782b
13 changed files with 1394 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
agent/assistant/assistant.go
View file

@ -8,6 +8,7 @@ import (
"github.com/yaoapp/yao/agent/caller"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/search"
searchTypes "github.com/yaoapp/yao/agent/search/types"
store "github.com/yaoapp/yao/agent/store/types"
sui "github.com/yaoapp/yao/sui/core"
@ -23,6 +24,9 @@ func init() {
// Return a wrapper that implements AgentCaller interface
return &agentCallerWrapper{ast: ast}, nil
}
// Initialize Search JSAPI factory
search.SetJSAPIFactory()
}
// agentCallerWrapper wraps Assistant to implement AgentCaller interface

35
agent/context/jsapi_search.go Normal file
View file

@ -0,0 +1,35 @@
package context
// SearchAPI defines the search JSAPI interface for ctx.search.*
// This interface is defined here to avoid circular dependency between context and search packages.
// The actual implementation is in agent/search/jsapi.go
type SearchAPI interface {
// Web executes web search
// Returns *types.Result or error information
Web(query string, opts map[string]interface{}) interface{}
// KB executes knowledge base search
// Returns *types.Result or error information
KB(query string, opts map[string]interface{}) interface{}
// DB executes database search
// Returns *types.Result or error information
DB(query string, opts map[string]interface{}) interface{}
// Parallel executes multiple searches in parallel
// Returns []*types.Result
Parallel(requests []interface{}) []interface{}
}
// SearchAPIFactory is a function type that creates a SearchAPI for a context
// This is set by the search package during initialization
var SearchAPIFactory func(ctx *Context) SearchAPI
// Search returns the search API for this context
// Returns nil if SearchAPIFactory is not set
func (ctx *Context) Search() SearchAPI {
if SearchAPIFactory == nil {
return nil
}
return SearchAPIFactory(ctx)
}

178
agent/search/DESIGN.md
View file

@ -151,9 +151,16 @@ agent/search/
│ └── 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
│ ├── keyword/ # Keyword extraction (Handler + Registry pattern)
│ │ ├── extractor.go # Main extractor (mode dispatch)
│ │ ├── builtin.go # Builtin frequency-based extraction
│ │ ├── agent.go # Agent mode (LLM-powered)
│ │ └── mcp.go # MCP mode (external service)
│ └── querydsl/ # QueryDSL generation for DB search
│ ├── generator.go # Main generator (mode dispatch)
│ ├── builtin.go # Builtin template-based generation
│ ├── agent.go # Agent mode (LLM-powered)
│ └── mcp.go # MCP mode (external service)
│ # Note: Embedding follows KB collection config, not in this package
│
├── handlers/ # Search handler implementations
@ -860,6 +867,32 @@ const (
The Search module is exposed via `ctx.search` object in hook scripts.
### Architecture
To avoid circular dependency between `context` and `search` packages:
```
agent/context/jsapi_search.go agent/search/jsapi.go
┌─────────────────────────┐ ┌─────────────────────────┐
│ SearchAPI interface │◄──────│ JSAPI struct │
│ SearchAPIFactory var │ │ (implements SearchAPI) │
│ ctx.Search() method │ │ SetJSAPIFactory() │
└─────────────────────────┘ └─────────────────────────┘
▲ │
│ │
└──────────────────────────────────┘
Factory registration
(in assistant/init)
```
**Key Files:**
| File | Description |
| ----------------------------- | ------------------------------ |
| `context/jsapi_search.go` | SearchAPI interface definition |
| `search/jsapi.go` | JSAPI implementation |
| `assistant/assistant.go:init` | Factory registration |
### API Methods
```typescript
@ -1582,50 +1615,64 @@ Configure via `uses.*` in `agent/agent.yml`:
| `<assistant-id>` | Delegate to an assistant (Agent) | LLM-based, custom logic |
| `mcp:<server>.<tool>` | Call MCP tool | External services integration |
#### Keyword Extraction (`nlp/keyword.go`)
#### Keyword Extraction (`nlp/keyword/`)
Configure via `uses.keyword`:
Configure via `uses.keyword`. The keyword extraction module follows the Handler + Registry pattern with three modes:
| Mode | Value | Description |
| ------- | ---------------------------- | --------------------------------------------- |
| Builtin | `"builtin"` | Frequency-based extraction (no external deps) |
| Agent | `"workers.nlp.keyword"` | LLM-powered semantic extraction |
| MCP | `"mcp:nlp.extract_keywords"` | External service via MCP |
**Directory Structure:**
```
nlp/keyword/
├── extractor.go # Main entry point (mode dispatch)
├── builtin.go # Builtin: frequency-based, stopword filtering
├── agent.go # Agent: delegate to LLM assistant
└── mcp.go # MCP: call external tool
```
**Usage:**
```go
// nlp/keyword.go
package nlp
// nlp/keyword/extractor.go
package keyword
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", "<assistant-id>", "mcp:<server>.<tool>"
// Extractor extracts keywords from text
type Extractor struct {
usesKeyword string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
config *types.KeywordConfig
}
// NewKeywordExtractor creates a keyword extractor
func NewKeywordExtractor(usesKeyword string, cfg *types.KeywordConfig) *KeywordExtractor {
return &KeywordExtractor{usesKeyword: usesKeyword, config: cfg}
}
// NewExtractor creates a new keyword extractor
func NewExtractor(usesKeyword string, cfg *types.KeywordConfig) *Extractor
// 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)
}
}
// Extract extracts keywords based on configured mode
func (e *Extractor) Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error)
```
**Builtin Implementation:**
The builtin extractor uses simple frequency-based extraction with no external dependencies:
- Tokenization (handles English and Chinese)
- Stop word filtering (common English and Chinese stop words)
- Frequency counting and ranking
- Returns top N keywords by frequency
> **Note**: For production use cases requiring high accuracy (semantic understanding, phrase extraction), use Agent or MCP mode.
**Example:**
```
"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"]
↓ builtin: tokenization + stopword removal + frequency ranking
→ ["wireless", "headphones", "find", "best"]
↓ agent: LLM semantic extraction
→ ["wireless headphones", "under $100", "best"]
```
#### Embedding (KB Collection Config)
@ -1650,51 +1697,54 @@ func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Resul
}
```
#### QueryDSL Generation (`nlp/querydsl.go`)
#### QueryDSL Generation (`nlp/querydsl/`)
Configure via `uses.querydsl`:
Configure via `uses.querydsl`. The QueryDSL generation module follows the same pattern as keyword extraction:
| Mode | Value | Description |
| ------- | ----------------------------- | ------------------------------------------- |
| Builtin | `"builtin"` | Template-based generation from model schema |
| Agent | `"workers.nlp.querydsl"` | LLM-powered semantic query generation |
| MCP | `"mcp:nlp.generate_querydsl"` | External service via MCP |
**Directory Structure:**
```
nlp/querydsl/
├── generator.go # Main entry point (mode dispatch)
├── builtin.go # Builtin: template-based generation
├── agent.go # Agent: delegate to LLM assistant
└── mcp.go # MCP: call external tool
```
**Usage:**
```go
// nlp/querydsl.go
package nlp
// nlp/querydsl/generator.go
package querydsl
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", "<assistant-id>", "mcp:<server>.<tool>"
// Generator generates QueryDSL from natural language
type Generator struct {
usesQueryDSL string
config *types.QueryDSLConfig
}
// NewQueryDSLGenerator creates a QueryDSL generator
func NewQueryDSLGenerator(usesQueryDSL string, cfg *types.QueryDSLConfig) *QueryDSLGenerator {
return &QueryDSLGenerator{usesQueryDSL: usesQueryDSL, config: cfg}
}
// NewGenerator creates a new QueryDSL generator
func NewGenerator(usesQueryDSL string, cfg *types.QueryDSLConfig) *Generator
// Generate converts natural language to QueryDSL
// Uses GOU types directly: model.Model and gou.QueryDSL
func (g *QueryDSLGenerator) Generate(query string, models []*model.Model) (*gou.QueryDSL, error) {
switch {
case g.usesQueryDSL == "builtin" || g.usesQueryDSL == "":
return g.builtinGenerate(query, models)
case strings.HasPrefix(g.usesQueryDSL, "mcp:"):
return g.mcpGenerate(query, models)
default:
return g.agentGenerate(query, models)
}
}
func (g *Generator) Generate(query string, models []*model.Model) (*gou.QueryDSL, error)
```
**Example:**
```
"Products cheaper than $100 from Apple"
↓ builtin: template matching against model schema
→ QueryDSL with simple keyword matching
↓ agent: LLM generates DSL from NL + schema
→ QueryDSL: {"wheres": [{"column": "price", "op": "<", "value": 100}, {"column": "brand", "value": "Apple"}]}
→ QueryDSL: {"wheres": [{"column": "price", "op": "<", "value": 100}, {"column": "brand", "value": "Apple"}]}
```
## Handlers & Providers

4
agent/search/interfaces/nlp.go
View file

@ -3,13 +3,15 @@ package interfaces
import (
"github.com/yaoapp/gou/model"
"github.com/yaoapp/gou/query/gou"
"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(content string, opts *types.KeywordOptions) ([]string, error)
// ctx is required for Agent and MCP modes, can be nil for builtin mode
Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error)
}
// QueryDSLGenerator generates QueryDSL for DB search

116
agent/search/jsapi.go Normal file
View file

@ -0,0 +1,116 @@
package search
import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// JSAPI implements context.SearchAPI interface
// Provides ctx.search.Web(), ctx.search.KB(), ctx.search.DB(), ctx.search.Parallel()
type JSAPI struct {
ctx *context.Context
config *types.Config
uses *Uses
}
// NewJSAPI creates a new search JSAPI instance
func NewJSAPI(ctx *context.Context, config *types.Config, uses *Uses) *JSAPI {
return &JSAPI{
ctx: ctx,
config: config,
uses: uses,
}
}
// Web executes web search
// Options:
// - limit: int - max results (default: 10)
// - sites: []string - restrict to specific sites
// - time_range: string - "day", "week", "month", "year"
// - rerank: map[string]interface{} - rerank options
func (api *JSAPI) Web(query string, opts map[string]interface{}) interface{} {
// TODO: Implement web search
// 1. Build Request from query and opts
// 2. Call web handler
// 3. Return Result or error
return &types.Result{
Type: types.SearchTypeWeb,
Query: query,
Error: "not implemented",
}
}
// KB executes knowledge base search
// Options:
// - collections: []string - collection IDs
// - threshold: float64 - similarity threshold (0-1)
// - limit: int - max results
// - graph: bool - enable graph association
// - rerank: map[string]interface{} - rerank options
func (api *JSAPI) KB(query string, opts map[string]interface{}) interface{} {
// TODO: Implement KB search
// 1. Build Request from query and opts
// 2. Call KB handler
// 3. Return Result or error
return &types.Result{
Type: types.SearchTypeKB,
Query: query,
Error: "not implemented",
}
}
// DB executes database search
// Options:
// - models: []string - model IDs
// - wheres: []map[string]interface{} - pre-defined filters (GOU QueryDSL Where format)
// - orders: []map[string]interface{} - sort orders (GOU QueryDSL Order format)
// - select: []string - fields to return
// - limit: int - max results
// - rerank: map[string]interface{} - rerank options
func (api *JSAPI) DB(query string, opts map[string]interface{}) interface{} {
// TODO: Implement DB search
// 1. Build Request from query and opts
// 2. Call DB handler
// 3. Return Result or error
return &types.Result{
Type: types.SearchTypeDB,
Query: query,
Error: "not implemented",
}
}
// Parallel executes multiple searches in parallel
// Each request should have:
// - type: string - "web", "kb", or "db"
// - query: string - search query
// - ... other type-specific options
func (api *JSAPI) Parallel(requests []interface{}) []interface{} {
// TODO: Implement parallel search
// 1. Parse requests into []Request
// 2. Call SearchMultiple
// 3. Return []Result
results := make([]interface{}, len(requests))
for i := range requests {
results[i] = &types.Result{
Error: "not implemented",
}
}
return results
}
// init registers the JSAPI factory with context package
func init() {
// Note: The actual factory is set by assistant package during initialization
// This avoids circular dependency: context -> search -> context
// See: assistant/assistant.go init()
}
// SetJSAPIFactory sets the factory function for creating SearchAPI instances
// Called by assistant package during initialization
func SetJSAPIFactory() {
context.SearchAPIFactory = func(ctx *context.Context) context.SearchAPI {
// Get config and uses from context or use defaults
// TODO: Get actual config from assistant
return NewJSAPI(ctx, nil, nil)
}
}

175
agent/search/nlp/keyword/agent.go Normal file
View file

@ -0,0 +1,175 @@
package keyword
import (
"encoding/json"
"fmt"
"github.com/yaoapp/yao/agent/caller"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// AgentProvider delegates keyword extraction to an LLM-powered assistant
// The assistant can understand context and extract semantically relevant keywords
type AgentProvider struct {
agentID string // Assistant ID to delegate to
}
// NewAgentProvider creates a new agent-based keyword extractor
func NewAgentProvider(agentID string) *AgentProvider {
return &AgentProvider{
agentID: agentID,
}
}
// Extract extracts keywords by calling the target agent
// The agent receives the content and returns extracted keywords
func (p *AgentProvider) Extract(ctx *agentContext.Context, content string, opts *types.KeywordOptions) ([]string, error) {
if ctx == nil {
return nil, fmt.Errorf("context is required for agent keyword extraction")
}
// Check if AgentGetterFunc is initialized
if caller.AgentGetterFunc == nil {
return nil, fmt.Errorf("AgentGetterFunc not initialized")
}
// Get the agent
agent, err := caller.AgentGetterFunc(p.agentID)
if err != nil {
return nil, fmt.Errorf("failed to get agent %s: %w", p.agentID, err)
}
// Build the request message
requestData := map[string]interface{}{
"content": content,
"max_keywords": opts.MaxKeywords,
"language": opts.Language,
}
requestJSON, _ := json.Marshal(requestData)
// Create message for the agent
messages := []agentContext.Message{
{
Role: "user",
Content: string(requestJSON),
},
}
// Call the agent with skip options (no history, no output)
options := &agentContext.Options{
Skip: &agentContext.Skip{
History: true,
Output: true,
},
}
result, err := agent.Stream(ctx, messages, options)
if err != nil {
return nil, fmt.Errorf("agent call failed: %w", err)
}
// Debug: log the result type and value
// fmt.Printf("DEBUG Agent result type: %T, value: %+v\n", result, result)
// Parse the result
return p.parseResult(result)
}
// parseResult extracts keywords from the agent's response
// The agent should return data in NextHookResponse format: { data: { keywords: [...] } }
// The Stream() response wraps this in: { next: { data: { keywords: [...] } } }
func (p *AgentProvider) parseResult(result interface{}) ([]string, error) {
if result == nil {
return []string{}, nil
}
// Try to convert to map first (most common case)
var data map[string]interface{}
switch v := result.(type) {
case map[string]interface{}:
data = v
case string:
// Try to parse as JSON
if err := json.Unmarshal([]byte(v), &data); err != nil {
// Not a JSON object, try as array
var keywords []string
if err := json.Unmarshal([]byte(v), &keywords); err == nil {
return keywords, nil
}
// Return as single keyword
return []string{v}, nil
}
case []string:
return v, nil
case []interface{}:
keywords := make([]string, 0, len(v))
for _, item := range v {
if s, ok := item.(string); ok {
keywords = append(keywords, s)
}
}
return keywords, nil
default:
// Try to marshal and unmarshal
jsonBytes, err := json.Marshal(result)
if err != nil {
return []string{}, nil
}
if err := json.Unmarshal(jsonBytes, &data); err != nil {
return []string{}, nil
}
}
// Check for "next" field (custom hook data from NextHookResponse)
// Stream() returns: { next: { data: { keywords: [...] } } }
if next, hasNext := data["next"]; hasNext && next != nil {
if nextMap, ok := next.(map[string]interface{}); ok {
data = nextMap
} else if nextStr, ok := next.(string); ok {
if err := json.Unmarshal([]byte(nextStr), &data); err != nil {
return []string{}, nil
}
}
}
// Extract keywords from data
// Try common field names: "keywords", "data", "data.keywords"
if kw, ok := data["keywords"]; ok {
return p.extractKeywordsFromValue(kw)
}
if d, ok := data["data"]; ok {
if dm, ok := d.(map[string]interface{}); ok {
if kw, ok := dm["keywords"]; ok {
return p.extractKeywordsFromValue(kw)
}
}
return p.extractKeywordsFromValue(d)
}
return []string{}, nil
}
// extractKeywordsFromValue extracts string array from various types
func (p *AgentProvider) extractKeywordsFromValue(v interface{}) ([]string, error) {
switch kw := v.(type) {
case []string:
return kw, nil
case []interface{}:
keywords := make([]string, 0, len(kw))
for _, item := range kw {
if s, ok := item.(string); ok {
keywords = append(keywords, s)
}
}
return keywords, nil
case string:
var keywords []string
if err := json.Unmarshal([]byte(kw), &keywords); err == nil {
return keywords, nil
}
return []string{kw}, nil
}
return []string{}, nil
}

120
agent/search/nlp/keyword/agent_test.go Normal file
View file

@ -0,0 +1,120 @@
package keyword_test
import (
"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/search/nlp/keyword"
searchTypes "github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
)
func TestAgentProviderWithAssistantConfig(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Load the keyword-agent assistant that will provide keywords
ast, err := assistant.Get("tests.keyword-agent")
require.NoError(t, err)
require.NotNil(t, ast)
// Create test context
ctx := newTestContext(t)
// Create extractor with agent mode
extractor := keyword.NewExtractor("tests.keyword-agent", &searchTypes.KeywordConfig{
MaxKeywords: 5,
Language: "auto",
})
// Test extraction
content := "Machine learning and deep learning are subfields of artificial intelligence"
keywords, err := extractor.Extract(ctx, content, nil)
require.NoError(t, err)
assert.NotEmpty(t, keywords, "Agent should return keywords")
assert.LessOrEqual(t, len(keywords), 5, "Should respect max_keywords")
// Verify keywords are relevant
t.Logf("Extracted keywords: %v", keywords)
}
func TestAgentProviderWithCustomOptions(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newTestContext(t)
// Create extractor with agent mode
extractor := keyword.NewExtractor("tests.keyword-agent", &searchTypes.KeywordConfig{
MaxKeywords: 10,
})
// Test with runtime options override
content := "Python programming language for data science and web development"
keywords, err := extractor.Extract(ctx, content, &searchTypes.KeywordOptions{
MaxKeywords: 3, // Override to 3
})
require.NoError(t, err)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 3, "Should respect runtime max_keywords override")
t.Logf("Extracted keywords (max 3): %v", keywords)
}
func TestAgentProviderWithoutContext(t *testing.T) {
// Test that agent mode requires context
extractor := keyword.NewExtractor("tests.keyword-agent", nil)
_, err := extractor.Extract(nil, "test content", nil)
assert.Error(t, err, "Agent mode should require context")
assert.Contains(t, err.Error(), "context is required")
}
func TestAgentProviderAgentNotFound(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newTestContext(t)
// Create extractor with non-existent agent
extractor := keyword.NewExtractor("non-existent-agent", nil)
_, err := extractor.Extract(ctx, "test content", nil)
assert.Error(t, err, "Should error for non-existent agent")
assert.Contains(t, err.Error(), "failed to get agent")
}
// newTestContext creates a test context with required fields
func newTestContext(t *testing.T) *context.Context {
t.Helper()
authorized := &oauthTypes.AuthorizedInfo{
UserID: "test-user",
}
chatID := "test-chat-keyword"
ctx := context.New(t.Context(), authorized, chatID)
return ctx
}

243
agent/search/nlp/keyword/builtin.go Normal file
View file

@ -0,0 +1,243 @@
package keyword
import (
"regexp"
"sort"
"strings"
"unicode"
)
// BuiltinExtractor implements simple frequency-based keyword extraction
// This is a lightweight implementation with no external dependencies.
//
// Algorithm:
// 1. Tokenize text (split by whitespace and punctuation)
// 2. Normalize (lowercase, trim)
// 3. Filter stop words and short words
// 4. Count word frequency
// 5. Return top N words by frequency
//
// Limitations:
// - No semantic understanding
// - No phrase extraction (single words only)
// - Basic Chinese support (splits by punctuation, no proper segmentation)
//
// For better results, use Agent or MCP mode with LLM-based extraction.
type BuiltinExtractor struct {
stopWords map[string]bool
minLength int // minimum word length to consider
}
// Result represents an extracted keyword with its score
type Result struct {
Word string `json:"word"`
Score float64 `json:"score"` // frequency-based score (0-1)
}
// NewBuiltinExtractor creates a new builtin keyword extractor
func NewBuiltinExtractor() *BuiltinExtractor {
return &BuiltinExtractor{
stopWords: defaultStopWords,
minLength: 2,
}
}
// Extract extracts keywords from text using frequency-based algorithm
func (e *BuiltinExtractor) Extract(text string, limit int) []Result {
if text == "" || limit <= 0 {
return []Result{}
}
// Step 1: Tokenize
tokens := e.tokenize(text)
// Step 2 & 3: Normalize and filter
var words []string
for _, token := range tokens {
word := e.normalize(token)
if e.shouldKeep(word) {
words = append(words, word)
}
}
if len(words) == 0 {
return []Result{}
}
// Step 4: Count frequency
freq := make(map[string]int)
for _, word := range words {
freq[word]++
}
// Step 5: Sort by frequency and return top N
type wordFreq struct {
word string
freq int
}
var sorted []wordFreq
for word, count := range freq {
sorted = append(sorted, wordFreq{word, count})
}
sort.Slice(sorted, func(i, j int) bool {
return sorted[i].freq > sorted[j].freq
})
// Calculate max frequency for normalization
maxFreq := 1
if len(sorted) > 0 {
maxFreq = sorted[0].freq
}
// Build result with normalized scores
result := make([]Result, 0, limit)
for i := 0; i < len(sorted) && i < limit; i++ {
result = append(result, Result{
Word: sorted[i].word,
Score: float64(sorted[i].freq) / float64(maxFreq),
})
}
return result
}
// ExtractAsStrings is a convenience method that returns just the keyword strings
func (e *BuiltinExtractor) ExtractAsStrings(text string, limit int) []string {
results := e.Extract(text, limit)
words := make([]string, len(results))
for i, r := range results {
words[i] = r.Word
}
return words
}
// tokenize splits text into tokens
// Handles both English (space-separated) and Chinese (character-based with punctuation splits)
func (e *BuiltinExtractor) tokenize(text string) []string {
// Split by whitespace and common punctuation
splitter := regexp.MustCompile(`[\s\p{P}\p{S}]+`)
tokens := splitter.Split(text, -1)
// Further split mixed Chinese/English text
var result []string
for _, token := range tokens {
if token == "" {
continue
}
// Split Chinese characters as individual tokens (basic approach)
// For proper Chinese segmentation, use Agent/MCP mode
subTokens := e.splitMixedText(token)
result = append(result, subTokens...)
}
return result
}
// splitMixedText handles mixed Chinese/English text
// Chinese characters are grouped together, English words stay as-is
func (e *BuiltinExtractor) splitMixedText(text string) []string {
var result []string
var current strings.Builder
var lastType int // 0=none, 1=chinese, 2=other
for _, r := range text {
currentType := 0
if unicode.Is(unicode.Han, r) {
currentType = 1
} else if unicode.IsLetter(r) || unicode.IsDigit(r) {
currentType = 2
}
if currentType == 0 {
// Non-word character, flush current
if current.Len() > 0 {
result = append(result, current.String())
current.Reset()
}
lastType = 0
continue
}
if lastType != 0 && lastType != currentType {
// Type changed, flush current
if current.Len() > 0 {
result = append(result, current.String())
current.Reset()
}
}
current.WriteRune(r)
lastType = currentType
}
// Flush remaining
if current.Len() > 0 {
result = append(result, current.String())
}
return result
}
// normalize converts word to lowercase and trims whitespace
func (e *BuiltinExtractor) normalize(word string) string {
return strings.ToLower(strings.TrimSpace(word))
}
// shouldKeep checks if a word should be kept (not a stop word, meets length requirement)
func (e *BuiltinExtractor) shouldKeep(word string) bool {
if len(word) < e.minLength {
return false
}
if e.stopWords[word] {
return false
}
// Keep if it contains at least one letter or Chinese character
for _, r := range word {
if unicode.IsLetter(r) {
return true
}
}
return false
}
// defaultStopWords contains common stop words for English and Chinese
// This is a minimal set to keep the implementation lightweight.
// For comprehensive stop word filtering, use Agent/MCP mode.
var defaultStopWords = map[string]bool{
// English stop words (most common ~100)
"a": true, "an": true, "the": true, "and": true, "or": true, "but": true,
"is": true, "are": true, "was": true, "were": true, "be": true, "been": true, "being": true,
"have": true, "has": true, "had": true, "do": true, "does": true, "did": true,
"will": true, "would": true, "could": true, "should": true, "may": true, "might": true,
"must": true, "shall": true, "can": true, "need": true, "dare": true,
"i": true, "you": true, "he": true, "she": true, "it": true, "we": true, "they": true,
"me": true, "him": true, "her": true, "us": true, "them": true,
"my": true, "your": true, "his": true, "its": true, "our": true, "their": true,
"mine": true, "yours": true, "hers": true, "ours": true, "theirs": true,
"this": true, "that": true, "these": true, "those": true,
"what": true, "which": true, "who": true, "whom": true, "whose": true,
"where": true, "when": true, "why": true, "how": true,
"all": true, "each": true, "every": true, "both": true, "few": true, "more": true,
"most": true, "other": true, "some": true, "such": true, "no": true, "not": true,
"only": true, "same": true, "so": true, "than": true, "too": true, "very": true,
"just": true, "also": true, "now": true, "here": true, "there": true,
"in": true, "on": true, "at": true, "by": true, "for": true, "with": true,
"about": true, "against": true, "between": true, "into": true, "through": true,
"during": true, "before": true, "after": true, "above": true, "below": true,
"to": true, "from": true, "up": true, "down": true, "out": true, "off": true,
"over": true, "under": true, "again": true, "further": true, "then": true, "once": true,
"as": true, "if": true, "because": true, "until": true, "while": true,
// Chinese stop words (most common ~50)
"的": true, "了": true, "和": true, "是": true, "就": true,
"都": true, "而": true, "及": true, "与": true, "着": true,
"或": true, "一个": true, "没有": true, "我们": true, "你们": true,
"他们": true, "它们": true, "这个": true, "那个": true, "这些": true,
"那些": true, "这里": true, "那里": true, "什么": true, "怎么": true,
"为什么": true, "哪里": true, "谁": true, "哪个": true, "多少": true,
"在": true, "有": true, "个": true, "中": true, "为": true,
"以": true, "于": true, "上": true, "下": true, "不": true,
"也": true, "很": true, "到": true, "说": true, "要": true,
"会": true, "可以": true, "这": true, "那": true, "但": true,
"如果": true, "因为": true, "所以": true, "虽然": true, "但是": true,
}

137
agent/search/nlp/keyword/builtin_test.go Normal file
View file

@ -0,0 +1,137 @@
package keyword
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestBuiltinExtractor_Extract(t *testing.T) {
extractor := NewBuiltinExtractor()
tests := []struct {
name string
text string
limit int
minCount int // minimum expected keywords
}{
{
name: "English text",
text: "The quick brown fox jumps over the lazy dog. The fox is very quick.",
limit: 5,
minCount: 3, // fox, quick, etc.
},
{
name: "Chinese text",
text: "人工智能技术正在快速发展，机器学习和深度学习是人工智能的核心技术",
limit: 5,
minCount: 2,
},
{
name: "Mixed text",
text: "AI人工智能 machine learning 机器学习 deep learning 深度学习",
limit: 10,
minCount: 3,
},
{
name: "Empty text",
text: "",
limit: 5,
minCount: 0,
},
{
name: "Only stop words",
text: "the a an is are was were",
limit: 5,
minCount: 0,
},
{
name: "Technical query",
text: "How to implement a search engine with Elasticsearch and Redis caching?",
limit: 5,
minCount: 3, // search, engine, elasticsearch, redis, caching
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
results := extractor.Extract(tt.text, tt.limit)
assert.GreaterOrEqual(t, len(results), tt.minCount, "Expected at least %d keywords", tt.minCount)
assert.LessOrEqual(t, len(results), tt.limit, "Should not exceed limit")
// Check scores are valid
for _, r := range results {
assert.NotEmpty(t, r.Word)
assert.GreaterOrEqual(t, r.Score, 0.0)
assert.LessOrEqual(t, r.Score, 1.0)
}
})
}
}
func TestBuiltinExtractor_ExtractAsStrings(t *testing.T) {
extractor := NewBuiltinExtractor()
text := "Machine learning and deep learning are subfields of artificial intelligence"
keywords := extractor.ExtractAsStrings(text, 5)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 5)
// Check that common ML terms are extracted
keywordSet := make(map[string]bool)
for _, k := range keywords {
keywordSet[k] = true
}
assert.True(t, keywordSet["learning"] || keywordSet["machine"] || keywordSet["artificial"],
"Expected at least one relevant keyword")
}
func TestBuiltinExtractor_StopWords(t *testing.T) {
extractor := NewBuiltinExtractor()
// Test that stop words are filtered
text := "the quick brown fox is very lazy"
results := extractor.Extract(text, 10)
for _, r := range results {
assert.NotEqual(t, "the", r.Word)
assert.NotEqual(t, "is", r.Word)
assert.NotEqual(t, "very", r.Word)
}
}
func TestBuiltinExtractor_Frequency(t *testing.T) {
extractor := NewBuiltinExtractor()
// Word "search" appears 3 times, should rank higher
text := "search engine optimization, search ranking, search results"
results := extractor.Extract(text, 3)
assert.NotEmpty(t, results)
// "search" should be the top keyword
assert.Equal(t, "search", results[0].Word)
assert.Equal(t, 1.0, results[0].Score) // highest frequency = 1.0
}
func TestBuiltinExtractor_ZeroLimit(t *testing.T) {
extractor := NewBuiltinExtractor()
results := extractor.Extract("some text here", 0)
assert.Empty(t, results)
}
func TestBuiltinExtractor_ChineseStopWords(t *testing.T) {
extractor := NewBuiltinExtractor()
// Test that Chinese stop words are filtered
text := "这是一个关于人工智能的文章"
results := extractor.Extract(text, 10)
for _, r := range results {
assert.NotEqual(t, "这", r.Word)
assert.NotEqual(t, "是", r.Word)
assert.NotEqual(t, "一个", r.Word)
assert.NotEqual(t, "的", r.Word)
}
}

106
agent/search/nlp/keyword/extractor.go Normal file
View file

@ -0,0 +1,106 @@
// Package keyword provides keyword extraction for web search optimization
// Supports three modes via uses.keyword configuration:
// - "builtin": Simple frequency-based extraction (no external dependencies)
// - "<assistant-id>": Delegate to an LLM-powered assistant for high-quality extraction
// - "mcp:<server>.<tool>": Call external MCP tool
//
// For production use cases requiring high accuracy, use Agent or MCP mode.
package keyword
import (
"strings"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// Extractor extracts keywords from text
// Mode is determined by uses.keyword configuration
type Extractor struct {
usesKeyword string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
config *types.KeywordConfig // Keyword extraction options
}
// NewExtractor creates a new keyword extractor
// usesKeyword: value from uses.keyword config
// cfg: keyword extraction options from search config
func NewExtractor(usesKeyword string, cfg *types.KeywordConfig) *Extractor {
return &Extractor{
usesKeyword: usesKeyword,
config: cfg,
}
}
// Extract extracts keywords from content based on configured mode
// Returns a list of keywords optimized for search queries
func (e *Extractor) Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) {
// Merge options with config defaults
mergedOpts := e.mergeOptions(opts)
switch {
case e.usesKeyword == "builtin" || e.usesKeyword == "":
return e.builtinExtract(content, mergedOpts)
case strings.HasPrefix(e.usesKeyword, "mcp:"):
return e.mcpExtract(ctx, content, mergedOpts)
default:
// Assume it's an assistant ID for Agent mode
return e.agentExtract(ctx, content, mergedOpts)
}
}
// mergeOptions merges runtime options with config defaults
func (e *Extractor) mergeOptions(opts *types.KeywordOptions) *types.KeywordOptions {
result := &types.KeywordOptions{
MaxKeywords: 10, // default
Language: "auto", // default
}
// Apply config defaults
if e.config != nil {
if e.config.MaxKeywords > 0 {
result.MaxKeywords = e.config.MaxKeywords
}
if e.config.Language != "" {
result.Language = e.config.Language
}
}
// Apply runtime options (highest priority)
if opts != nil {
if opts.MaxKeywords > 0 {
result.MaxKeywords = opts.MaxKeywords
}
if opts.Language != "" {
result.Language = opts.Language
}
}
return result
}
// builtinExtract uses simple frequency-based extraction
// This is a lightweight implementation with no external dependencies.
// For better results, use Agent or MCP mode.
func (e *Extractor) builtinExtract(content string, opts *types.KeywordOptions) ([]string, error) {
extractor := NewBuiltinExtractor()
return extractor.ExtractAsStrings(content, opts.MaxKeywords), nil
}
// agentExtract delegates to an LLM-powered assistant
// The assistant can understand context and extract semantically relevant keywords
func (e *Extractor) agentExtract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) {
provider := NewAgentProvider(e.usesKeyword)
return provider.Extract(ctx, content, opts)
}
// mcpExtract calls an external MCP tool
// Format: "mcp:<server>.<tool>"
func (e *Extractor) mcpExtract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) {
mcpRef := strings.TrimPrefix(e.usesKeyword, "mcp:")
provider, err := NewMCPProvider(mcpRef)
if err != nil {
// Fallback to builtin on invalid MCP format
return e.builtinExtract(content, opts)
}
return provider.Extract(ctx, content, opts)
}

63
agent/search/nlp/keyword/extractor_test.go Normal file
View file

@ -0,0 +1,63 @@
package keyword_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/search/nlp/keyword"
"github.com/yaoapp/yao/agent/search/types"
)
func TestExtractor_BuiltinMode(t *testing.T) {
// Test builtin mode (no external dependencies)
extractor := keyword.NewExtractor("builtin", &types.KeywordConfig{
MaxKeywords: 5,
Language: "auto",
})
keywords, err := extractor.Extract(nil, "How to build a search engine with Elasticsearch?", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 5)
}
func TestExtractor_EmptyUsesKeyword(t *testing.T) {
// Empty uses.keyword should default to builtin
extractor := keyword.NewExtractor("", nil)
keywords, err := extractor.Extract(nil, "Machine learning algorithms", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords)
}
func TestExtractor_RuntimeOptionsOverride(t *testing.T) {
// Config has max_keywords=10, but runtime opts override to 3
extractor := keyword.NewExtractor("builtin", &types.KeywordConfig{
MaxKeywords: 10,
})
keywords, err := extractor.Extract(nil, "one two three four five six seven eight nine ten", &types.KeywordOptions{
MaxKeywords: 3,
})
assert.NoError(t, err)
assert.LessOrEqual(t, len(keywords), 3)
}
func TestExtractor_ConfigDefaults(t *testing.T) {
// No config, should use defaults
extractor := keyword.NewExtractor("builtin", nil)
keywords, err := extractor.Extract(nil, "Test query for keyword extraction", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 10) // default max_keywords is 10
}
func TestExtractor_InvalidMCPFormat(t *testing.T) {
// Invalid MCP format should fallback to builtin
extractor := keyword.NewExtractor("mcp:invalid", nil)
keywords, err := extractor.Extract(nil, "Test query", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords)
}

123
agent/search/nlp/keyword/mcp.go Normal file
View file

@ -0,0 +1,123 @@
package keyword
import (
"encoding/json"
"fmt"
"strings"
"github.com/yaoapp/gou/mcp"
gouMCPTypes "github.com/yaoapp/gou/mcp/types"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types"
)
// MCPProvider delegates keyword extraction to an MCP tool
type MCPProvider struct {
serverID string // MCP server ID
toolName string // Tool name to call
}
// NewMCPProvider creates a new MCP-based keyword extractor
// mcpRef format: "server.tool" (e.g., "nlp.extract_keywords")
func NewMCPProvider(mcpRef string) (*MCPProvider, error) {
parts := strings.SplitN(mcpRef, ".", 2)
if len(parts) != 2 {
return nil, fmt.Errorf("invalid MCP format, expected 'server.tool', got '%s'", mcpRef)
}
return &MCPProvider{
serverID: parts[0],
toolName: parts[1],
}, nil
}
// Extract extracts keywords by calling the MCP tool
func (p *MCPProvider) Extract(ctx *agentContext.Context, content string, opts *types.KeywordOptions) ([]string, error) {
// Get MCP client
client, err := mcp.Select(p.serverID)
if err != nil {
return nil, fmt.Errorf("MCP server '%s' not found: %w", p.serverID, err)
}
// Build arguments for the MCP tool
arguments := map[string]interface{}{
"content": content,
"max_keywords": opts.MaxKeywords,
"language": opts.Language,
}
// Call the MCP tool (ctx embeds context.Context)
callResult, err := client.CallTool(ctx, p.toolName, arguments)
if err != nil {
return nil, fmt.Errorf("MCP tool call failed: %w", err)
}
// Parse the result
return p.parseResult(callResult)
}
// parseResult extracts keywords from the MCP tool response
func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse) ([]string, error) {
if result == nil {
return []string{}, nil
}
// Check for errors in result
if result.IsError {
errMsg := "MCP tool returned error"
if len(result.Content) > 0 && result.Content[0].Text != "" {
errMsg = result.Content[0].Text
}
return nil, fmt.Errorf("%s", errMsg)
}
// Parse content - expect JSON data with "keywords" field
if len(result.Content) == 0 {
return []string{}, nil
}
// Try to extract keywords from content
for _, content := range result.Content {
// Check text content type
if content.Type == gouMCPTypes.ToolContentTypeText && content.Text != "" {
// Try to parse as JSON
var data map[string]interface{}
if err := json.Unmarshal([]byte(content.Text), &data); err == nil {
// Look for "keywords" field
if kw, ok := data["keywords"]; ok {
return p.extractKeywordsFromValue(kw)
}
}
// Try to parse as direct array
var keywords []string
if err := json.Unmarshal([]byte(content.Text), &keywords); err == nil {
return keywords, nil
}
}
}
return []string{}, nil
}
// extractKeywordsFromValue extracts string array from various types
func (p *MCPProvider) extractKeywordsFromValue(v interface{}) ([]string, error) {
switch kw := v.(type) {
case []string:
return kw, nil
case []interface{}:
keywords := make([]string, 0, len(kw))
for _, item := range kw {
if s, ok := item.(string); ok {
keywords = append(keywords, s)
}
}
return keywords, nil
case string:
var keywords []string
if err := json.Unmarshal([]byte(kw), &keywords); err == nil {
return keywords, nil
}
return []string{kw}, nil
}
return []string{}, nil
}

155
agent/search/nlp/keyword/mcp_test.go Normal file
View file

@ -0,0 +1,155 @@
package keyword_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/nlp/keyword"
searchTypes "github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
)
func TestMCPProviderWithAssistantConfig(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with MCP mode
extractor := keyword.NewExtractor("mcp:search.extract_keywords", &searchTypes.KeywordConfig{
MaxKeywords: 5,
Language: "auto",
})
// Test extraction
content := "Machine learning and deep learning are subfields of artificial intelligence"
keywords, err := extractor.Extract(ctx, content, nil)
require.NoError(t, err)
assert.NotEmpty(t, keywords, "MCP should return keywords")
assert.LessOrEqual(t, len(keywords), 5, "Should respect max_keywords")
t.Logf("Extracted keywords via MCP: %v", keywords)
}
func TestMCPProviderWithCustomOptions(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with MCP mode
extractor := keyword.NewExtractor("mcp:search.extract_keywords", &searchTypes.KeywordConfig{
MaxKeywords: 10,
})
// Test with runtime options override
content := "Python programming language for data science and web development"
keywords, err := extractor.Extract(ctx, content, &searchTypes.KeywordOptions{
MaxKeywords: 3, // Override to 3
})
require.NoError(t, err)
assert.NotEmpty(t, keywords)
assert.LessOrEqual(t, len(keywords), 3, "Should respect runtime max_keywords override")
t.Logf("Extracted keywords via MCP (max 3): %v", keywords)
}
func TestMCPProviderInvalidFormat(t *testing.T) {
// Test invalid MCP format fallback to builtin
extractor := keyword.NewExtractor("mcp:invalid", nil)
// Should fallback to builtin (no error)
keywords, err := extractor.Extract(nil, "test content for keyword extraction", nil)
assert.NoError(t, err)
assert.NotEmpty(t, keywords, "Should fallback to builtin and extract keywords")
}
func TestMCPProviderServerNotFound(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with non-existent MCP server
extractor := keyword.NewExtractor("mcp:nonexistent.extract_keywords", &searchTypes.KeywordConfig{})
_, err := extractor.Extract(ctx, "test content", nil)
assert.Error(t, err, "Should error for non-existent MCP server")
assert.Contains(t, err.Error(), "not found")
}
func TestMCPProviderToolNotFound(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with non-existent tool
extractor := keyword.NewExtractor("mcp:search.nonexistent_tool", &searchTypes.KeywordConfig{})
_, err := extractor.Extract(ctx, "test content", nil)
assert.Error(t, err, "Should error for non-existent MCP tool")
}
func TestMCPProviderEmptyContent(t *testing.T) {
// Skip if running short tests
if testing.Short() {
t.Skip("Skipping integration test")
}
// Initialize test environment
testutils.Prepare(t)
defer testutils.Clean(t)
// Create test context
ctx := newMCPTestContext(t)
// Create extractor with MCP mode
extractor := keyword.NewExtractor("mcp:search.extract_keywords", nil)
// Test with empty content - MCP tool should return error
_, err := extractor.Extract(ctx, "", nil)
assert.Error(t, err, "Should error for empty content")
}
// newMCPTestContext creates a test context for MCP tests
func newMCPTestContext(t *testing.T) *context.Context {
t.Helper()
authorized := &oauthTypes.AuthorizedInfo{
UserID: "test-user",
}
chatID := "test-chat-mcp-keyword"
ctx := context.New(t.Context(), authorized, chatID)
return ctx
}