From d3865d782bff8d119823dfda879bfdb749b565b1 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 13 Dec 2025 14:41:03 +0800 Subject: [PATCH] Initialize Search JSAPI Factory and Update Documentation - 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. --- agent/assistant/assistant.go | 4 + agent/context/jsapi_search.go | 35 +++ agent/search/DESIGN.md | 178 +++++++++------ agent/search/interfaces/nlp.go | 4 +- agent/search/jsapi.go | 116 ++++++++++ agent/search/nlp/keyword/agent.go | 175 +++++++++++++++ agent/search/nlp/keyword/agent_test.go | 120 ++++++++++ agent/search/nlp/keyword/builtin.go | 243 +++++++++++++++++++++ agent/search/nlp/keyword/builtin_test.go | 137 ++++++++++++ agent/search/nlp/keyword/extractor.go | 106 +++++++++ agent/search/nlp/keyword/extractor_test.go | 63 ++++++ agent/search/nlp/keyword/mcp.go | 123 +++++++++++ agent/search/nlp/keyword/mcp_test.go | 155 +++++++++++++ 13 files changed, 1394 insertions(+), 65 deletions(-) create mode 100644 agent/context/jsapi_search.go create mode 100644 agent/search/jsapi.go create mode 100644 agent/search/nlp/keyword/agent.go create mode 100644 agent/search/nlp/keyword/agent_test.go create mode 100644 agent/search/nlp/keyword/builtin.go create mode 100644 agent/search/nlp/keyword/builtin_test.go create mode 100644 agent/search/nlp/keyword/extractor.go create mode 100644 agent/search/nlp/keyword/extractor_test.go create mode 100644 agent/search/nlp/keyword/mcp.go create mode 100644 agent/search/nlp/keyword/mcp_test.go diff --git a/agent/assistant/assistant.go b/agent/assistant/assistant.go index 9ef9f589..26a4f146 100644 --- a/agent/assistant/assistant.go +++ b/agent/assistant/assistant.go @@ -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 diff --git a/agent/context/jsapi_search.go b/agent/context/jsapi_search.go new file mode 100644 index 00000000..e186c483 --- /dev/null +++ b/agent/context/jsapi_search.go @@ -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) +} diff --git a/agent/search/DESIGN.md b/agent/search/DESIGN.md index a9812236..c97d650c 100644 --- a/agent/search/DESIGN.md +++ b/agent/search/DESIGN.md @@ -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`: | `` | Delegate to an assistant (Agent) | LLM-based, custom logic | | `mcp:.` | 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", "", "mcp:." +// Extractor extracts keywords from text +type Extractor 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} -} +// 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", "", "mcp:." +// 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 diff --git a/agent/search/interfaces/nlp.go b/agent/search/interfaces/nlp.go index b4a207a9..7bc0f873 100644 --- a/agent/search/interfaces/nlp.go +++ b/agent/search/interfaces/nlp.go @@ -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 diff --git a/agent/search/jsapi.go b/agent/search/jsapi.go new file mode 100644 index 00000000..1799a084 --- /dev/null +++ b/agent/search/jsapi.go @@ -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) + } +} diff --git a/agent/search/nlp/keyword/agent.go b/agent/search/nlp/keyword/agent.go new file mode 100644 index 00000000..22f72c08 --- /dev/null +++ b/agent/search/nlp/keyword/agent.go @@ -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 +} diff --git a/agent/search/nlp/keyword/agent_test.go b/agent/search/nlp/keyword/agent_test.go new file mode 100644 index 00000000..650147af --- /dev/null +++ b/agent/search/nlp/keyword/agent_test.go @@ -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 +} diff --git a/agent/search/nlp/keyword/builtin.go b/agent/search/nlp/keyword/builtin.go new file mode 100644 index 00000000..d649c09c --- /dev/null +++ b/agent/search/nlp/keyword/builtin.go @@ -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, +} diff --git a/agent/search/nlp/keyword/builtin_test.go b/agent/search/nlp/keyword/builtin_test.go new file mode 100644 index 00000000..17b31b5b --- /dev/null +++ b/agent/search/nlp/keyword/builtin_test.go @@ -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) + } +} diff --git a/agent/search/nlp/keyword/extractor.go b/agent/search/nlp/keyword/extractor.go new file mode 100644 index 00000000..7a5c6ba9 --- /dev/null +++ b/agent/search/nlp/keyword/extractor.go @@ -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) +// - "": Delegate to an LLM-powered assistant for high-quality extraction +// - "mcp:.": 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", "", "mcp:." + 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:." +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) +} diff --git a/agent/search/nlp/keyword/extractor_test.go b/agent/search/nlp/keyword/extractor_test.go new file mode 100644 index 00000000..ce29596f --- /dev/null +++ b/agent/search/nlp/keyword/extractor_test.go @@ -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) +} diff --git a/agent/search/nlp/keyword/mcp.go b/agent/search/nlp/keyword/mcp.go new file mode 100644 index 00000000..d196cd60 --- /dev/null +++ b/agent/search/nlp/keyword/mcp.go @@ -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 +} diff --git a/agent/search/nlp/keyword/mcp_test.go b/agent/search/nlp/keyword/mcp_test.go new file mode 100644 index 00000000..77a3f993 --- /dev/null +++ b/agent/search/nlp/keyword/mcp_test.go @@ -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 +}