From b8c5829eb015d723ee280852d83668722be03141 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 17 Dec 2025 17:55:58 +0800 Subject: [PATCH] Refactor Keyword Extraction and QueryDSL Generation for Context Requirement - Updated the keyword extraction and QueryDSL generation processes to require a context parameter, enhancing the robustness of the extraction methods. - Replaced the previous frequency-based extraction with a system agent approach, utilizing the __yao.keyword and __yao.querydsl agents for improved accuracy and context awareness. - Removed obsolete builtin extraction implementations and tests, streamlining the codebase. - Enhanced test cases to validate the new context requirements, ensuring proper error handling when context is not provided. - Updated documentation to reflect changes in the extraction methods and their dependencies on context. --- agent/assistant/search.go | 22 +- agent/search/nlp/keyword/builtin.go | 243 -------------------- agent/search/nlp/keyword/builtin_test.go | 137 ----------- agent/search/nlp/keyword/extractor.go | 34 ++- agent/search/nlp/keyword/extractor_test.go | 64 +++--- agent/search/nlp/keyword/mcp_test.go | 10 +- agent/search/nlp/querydsl/builtin.go | 124 ---------- agent/search/nlp/querydsl/generator.go | 34 ++- agent/search/nlp/querydsl/generator_test.go | 229 +++++++++--------- 9 files changed, 188 insertions(+), 709 deletions(-) delete mode 100644 agent/search/nlp/keyword/builtin.go delete mode 100644 agent/search/nlp/keyword/builtin_test.go delete mode 100644 agent/search/nlp/querydsl/builtin.go diff --git a/agent/assistant/search.go b/agent/assistant/search.go index 7aae5518..3685eba7 100644 --- a/agent/assistant/search.go +++ b/agent/assistant/search.go @@ -56,19 +56,16 @@ func (ast *Assistant) shouldAutoSearch(ctx *context.Context, messages []context. // checkSearchIntent uses __yao.needsearch agent to determine if search is needed // Returns true if search is needed, false otherwise func (ast *Assistant) checkSearchIntent(ctx *context.Context, messages []context.Message) bool { - // Get the last user message - var userQuery string - for i := len(messages) - 1; i >= 0; i-- { - if messages[i].Role == "user" { - if content, ok := messages[i].Content.(string); ok { - userQuery = content - break - } + // Filter out system messages and pass full conversation context + var intentMessages []context.Message + for _, msg := range messages { + if msg.Role != "system" { + intentMessages = append(intentMessages, msg) } } - if userQuery == "" { - return true // No user message, proceed with search + if len(intentMessages) == 0 { + return true // No messages, proceed with search } // Try to get __yao.needsearch agent @@ -81,11 +78,6 @@ func (ast *Assistant) checkSearchIntent(ctx *context.Context, messages []context // === Output: Send loading message === loadingID := ast.sendIntentLoading(ctx) - // Build messages for intent detection - intentMessages := []context.Message{ - {Role: "user", Content: userQuery}, - } - // Call the needsearch agent (Stack will auto-track) // IMPORTANT: Skip search to prevent infinite loop, skip output to prevent JSON showing in UI opts := &context.Options{ diff --git a/agent/search/nlp/keyword/builtin.go b/agent/search/nlp/keyword/builtin.go deleted file mode 100644 index d649c09c..00000000 --- a/agent/search/nlp/keyword/builtin.go +++ /dev/null @@ -1,243 +0,0 @@ -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 deleted file mode 100644 index 17b31b5b..00000000 --- a/agent/search/nlp/keyword/builtin_test.go +++ /dev/null @@ -1,137 +0,0 @@ -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 index 7a5c6ba9..3a36cb32 100644 --- a/agent/search/nlp/keyword/extractor.go +++ b/agent/search/nlp/keyword/extractor.go @@ -1,19 +1,21 @@ // 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 +// - "builtin" or "": Uses __yao.keyword system agent (LLM-powered) +// - "": Delegate to a custom LLM-powered assistant // - "mcp:.": Call external MCP tool -// -// For production use cases requiring high accuracy, use Agent or MCP mode. package keyword import ( + "fmt" "strings" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/types" ) +// SystemKeywordAgent is the default system agent for keyword extraction +const SystemKeywordAgent = "__yao.keyword" + // Extractor extracts keywords from text // Mode is determined by uses.keyword configuration type Extractor struct { @@ -39,12 +41,13 @@ func (e *Extractor) Extract(ctx *context.Context, content string, opts *types.Ke switch { case e.usesKeyword == "builtin" || e.usesKeyword == "": - return e.builtinExtract(content, mergedOpts) + // Use system keyword agent + return e.agentExtract(ctx, content, SystemKeywordAgent, 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) + return e.agentExtract(ctx, content, e.usesKeyword, mergedOpts) } } @@ -78,18 +81,13 @@ func (e *Extractor) mergeOptions(opts *types.KeywordOptions) *types.KeywordOptio 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) +func (e *Extractor) agentExtract(ctx *context.Context, content string, agentID string, opts *types.KeywordOptions) ([]string, error) { + if ctx == nil { + return nil, fmt.Errorf("context is required for keyword extraction") + } + provider := NewAgentProvider(agentID) return provider.Extract(ctx, content, opts) } @@ -99,8 +97,8 @@ func (e *Extractor) mcpExtract(ctx *context.Context, content string, opts *types 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) + // Fallback to system agent on invalid MCP format + return e.agentExtract(ctx, content, SystemKeywordAgent, e.mergeOptions(nil)) } return provider.Extract(ctx, content, opts) } diff --git a/agent/search/nlp/keyword/extractor_test.go b/agent/search/nlp/keyword/extractor_test.go index ce29596f..ba91eb93 100644 --- a/agent/search/nlp/keyword/extractor_test.go +++ b/agent/search/nlp/keyword/extractor_test.go @@ -8,56 +8,48 @@ import ( "github.com/yaoapp/yao/agent/search/types" ) -func TestExtractor_BuiltinMode(t *testing.T) { - // Test builtin mode (no external dependencies) +func TestExtractor_BuiltinMode_RequiresContext(t *testing.T) { + // Test builtin mode requires context (now uses __yao.keyword agent) 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) + // Without context, should return error + _, err := extractor.Extract(nil, "How to build a search engine with Elasticsearch?", nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") } -func TestExtractor_EmptyUsesKeyword(t *testing.T) { - // Empty uses.keyword should default to builtin +func TestExtractor_EmptyUsesKeyword_RequiresContext(t *testing.T) { + // Empty uses.keyword should default to __yao.keyword agent extractor := keyword.NewExtractor("", nil) - keywords, err := extractor.Extract(nil, "Machine learning algorithms", nil) - assert.NoError(t, err) - assert.NotEmpty(t, keywords) + // Without context, should return error + _, err := extractor.Extract(nil, "Machine learning algorithms", nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") } -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, - }) +func TestExtractor_AgentMode_RequiresContext(t *testing.T) { + // Custom agent mode requires context + extractor := keyword.NewExtractor("custom.keyword.agent", nil) - 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) + _, err := extractor.Extract(nil, "Test query", nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") } -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 +func TestExtractor_MCPMode_InvalidFormat(t *testing.T) { + // Invalid MCP format should fallback to system agent (which requires context) extractor := keyword.NewExtractor("mcp:invalid", nil) - keywords, err := extractor.Extract(nil, "Test query", nil) - assert.NoError(t, err) - assert.NotEmpty(t, keywords) + _, err := extractor.Extract(nil, "Test query", nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") +} + +func TestExtractor_SystemKeywordAgentConstant(t *testing.T) { + // Verify the system keyword agent constant + assert.Equal(t, "__yao.keyword", keyword.SystemKeywordAgent) } diff --git a/agent/search/nlp/keyword/mcp_test.go b/agent/search/nlp/keyword/mcp_test.go index 77a3f993..40f40e3e 100644 --- a/agent/search/nlp/keyword/mcp_test.go +++ b/agent/search/nlp/keyword/mcp_test.go @@ -72,13 +72,13 @@ func TestMCPProviderWithCustomOptions(t *testing.T) { } func TestMCPProviderInvalidFormat(t *testing.T) { - // Test invalid MCP format fallback to builtin + // Test invalid MCP format fallback to system agent (requires context) 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") + // Should fallback to system agent which requires context + _, err := extractor.Extract(nil, "test content for keyword extraction", nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") } func TestMCPProviderServerNotFound(t *testing.T) { diff --git a/agent/search/nlp/querydsl/builtin.go b/agent/search/nlp/querydsl/builtin.go deleted file mode 100644 index 76e5cdc6..00000000 --- a/agent/search/nlp/querydsl/builtin.go +++ /dev/null @@ -1,124 +0,0 @@ -package querydsl - -import ( - "github.com/yaoapp/gou/model" - "github.com/yaoapp/gou/query/gou" -) - -// BuiltinGenerator implements template-based QueryDSL generation -// This is a placeholder implementation that returns a basic QueryDSL. -// -// TODO: Implement actual template-based generation: -// - Parse natural language query -// - Match against model schema -// - Generate appropriate where clauses -// - Handle common query patterns (search, filter, sort) -// -// For production use cases requiring high accuracy, use Agent or MCP mode. -type BuiltinGenerator struct{} - -// NewBuiltinGenerator creates a new builtin QueryDSL generator -func NewBuiltinGenerator() *BuiltinGenerator { - return &BuiltinGenerator{} -} - -// Generate generates QueryDSL from natural language -// Currently returns a placeholder QueryDSL that searches all searchable fields -func (g *BuiltinGenerator) Generate(input *Input) (*Result, error) { - if input == nil || input.Query == "" { - return &Result{ - Warnings: []string{"empty query, returning empty DSL"}, - }, nil - } - - // Build a basic QueryDSL - dsl := &gou.QueryDSL{} - - // Set limit - limit := input.Limit - if limit <= 0 { - limit = 20 - } - dsl.Limit = limit - - // Apply pre-defined wheres if provided - if len(input.Wheres) > 0 { - dsl.Wheres = input.Wheres - } - - // Apply orders if provided - if len(input.Orders) > 0 { - dsl.Orders = input.Orders - } - - // Load models and try to generate basic search conditions - // Use the first model as the primary table, others can be joined - if len(input.ModelIDs) > 0 { - primaryModelID := input.ModelIDs[0] - - // Check if model exists before selecting - if !model.Exists(primaryModelID) { - return &Result{ - DSL: dsl, - Explain: "Generated basic QueryDSL (model not found)", - Warnings: []string{ - "model '" + primaryModelID + "' not found, returning basic DSL without search conditions", - }, - }, nil - } - - primaryModel := model.Select(primaryModelID) - if primaryModel != nil && len(primaryModel.MetaData.Columns) > 0 { - // Find searchable text columns (string/text types with index) - var searchableColumns []string - for _, col := range primaryModel.MetaData.Columns { - // Use Index as a proxy for searchable, and check for text types - if col.Index && (col.Type == "string" || col.Type == "text" || col.Type == "longText") { - searchableColumns = append(searchableColumns, col.Name) - } - } - - // If we have searchable columns and no pre-defined wheres, add a basic search - if len(searchableColumns) > 0 && len(input.Wheres) == 0 { - // Build OR conditions for searchable columns - orWheres := make([]gou.Where, 0, len(searchableColumns)) - for _, col := range searchableColumns { - orWheres = append(orWheres, gou.Where{ - Condition: gou.Condition{ - Field: &gou.Expression{Field: col}, - OP: "match", - Value: input.Query, - }, - }) - } - - // Wrap in OR group if multiple columns - if len(orWheres) > 1 { - // Mark all but the first as OR conditions - for i := 1; i < len(orWheres); i++ { - orWheres[i].OR = true - } - dsl.Wheres = []gou.Where{ - { - Wheres: orWheres, - }, - } - } else if len(orWheres) == 1 { - dsl.Wheres = orWheres - } - } - } - - // TODO: For multi-model queries, generate joins based on model relations - // This requires analyzing the relations between models and generating - // appropriate JOIN clauses in the QueryDSL - } - - return &Result{ - DSL: dsl, - Explain: "Generated basic search QueryDSL using builtin template (placeholder implementation)", - Warnings: []string{ - "builtin generator is a placeholder, consider using Agent or MCP mode for production", - }, - }, nil -} diff --git a/agent/search/nlp/querydsl/generator.go b/agent/search/nlp/querydsl/generator.go index b36c6285..0b531afe 100644 --- a/agent/search/nlp/querydsl/generator.go +++ b/agent/search/nlp/querydsl/generator.go @@ -1,13 +1,12 @@ // Package querydsl provides QueryDSL generation from natural language for DB search // Supports three modes via uses.querydsl configuration: -// - "builtin": Template-based generation (no external dependencies) -// - "": Delegate to an LLM-powered assistant for high-quality generation +// - "builtin" or "": Uses __yao.querydsl system agent (LLM-powered) +// - "": Delegate to a custom LLM-powered assistant // - "mcp:.": Call external MCP tool -// -// For production use cases requiring high accuracy, use Agent or MCP mode. package querydsl import ( + "fmt" "strings" "github.com/yaoapp/gou/query/gou" @@ -15,6 +14,9 @@ import ( "github.com/yaoapp/yao/agent/search/types" ) +// SystemQueryDSLAgent is the default system agent for QueryDSL generation +const SystemQueryDSLAgent = "__yao.querydsl" + // Generator generates QueryDSL from natural language // Mode is determined by uses.querydsl configuration type Generator struct { @@ -40,12 +42,13 @@ func (g *Generator) Generate(ctx *context.Context, input *Input) (*Result, error switch { case g.usesQueryDSL == "builtin" || g.usesQueryDSL == "": - result, err = g.builtinGenerate(input) + // Use system querydsl agent + result, err = g.agentGenerate(ctx, input, SystemQueryDSLAgent) case strings.HasPrefix(g.usesQueryDSL, "mcp:"): result, err = g.mcpGenerate(ctx, input) default: // Assume it's an assistant ID for Agent mode - result, err = g.agentGenerate(ctx, input) + result, err = g.agentGenerate(ctx, input, g.usesQueryDSL) } if err != nil { @@ -60,18 +63,13 @@ func (g *Generator) Generate(ctx *context.Context, input *Input) (*Result, error return result, nil } -// builtinGenerate uses template-based generation -// This is a lightweight implementation with no external dependencies. -// For better results, use Agent or MCP mode. -func (g *Generator) builtinGenerate(input *Input) (*Result, error) { - generator := NewBuiltinGenerator() - return generator.Generate(input) -} - // agentGenerate delegates to an LLM-powered assistant // The assistant can understand context and generate semantically correct QueryDSL -func (g *Generator) agentGenerate(ctx *context.Context, input *Input) (*Result, error) { - provider := NewAgentProvider(g.usesQueryDSL) +func (g *Generator) agentGenerate(ctx *context.Context, input *Input, agentID string) (*Result, error) { + if ctx == nil { + return nil, fmt.Errorf("context is required for QueryDSL generation") + } + provider := NewAgentProvider(agentID) return provider.Generate(ctx, input) } @@ -81,8 +79,8 @@ func (g *Generator) mcpGenerate(ctx *context.Context, input *Input) (*Result, er mcpRef := strings.TrimPrefix(g.usesQueryDSL, "mcp:") provider, err := NewMCPProvider(mcpRef) if err != nil { - // Fallback to builtin on invalid MCP format - return g.builtinGenerate(input) + // Fallback to system agent on invalid MCP format + return g.agentGenerate(ctx, input, SystemQueryDSLAgent) } return provider.Generate(ctx, input) } diff --git a/agent/search/nlp/querydsl/generator_test.go b/agent/search/nlp/querydsl/generator_test.go index 0f2b5736..a966d4b6 100644 --- a/agent/search/nlp/querydsl/generator_test.go +++ b/agent/search/nlp/querydsl/generator_test.go @@ -46,27 +46,24 @@ func TestNewGenerator(t *testing.T) { } } -func TestGenerator_Generate_Builtin(t *testing.T) { +func TestGenerator_Generate_Builtin_RequiresContext(t *testing.T) { + // Builtin mode now uses __yao.querydsl agent which requires context gen := NewGenerator("builtin", nil) - // Note: In real usage, models are loaded internally via model.Select() - // For this test, we just verify the basic flow works without models input := &Input{ Query: "find all active users", ModelIDs: []string{"user"}, Limit: 10, } - result, err := gen.Generate(nil, input) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.NotNil(t, result.DSL) - assert.NotEmpty(t, result.Explain) - assert.NotEmpty(t, result.Warnings) + // Without context, should return error + _, err := gen.Generate(nil, input) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") } -func TestGenerator_Generate_EmptyMode(t *testing.T) { - // Empty mode should default to builtin +func TestGenerator_Generate_EmptyMode_RequiresContext(t *testing.T) { + // Empty mode defaults to __yao.querydsl agent which requires context gen := NewGenerator("", nil) input := &Input{ @@ -75,116 +72,45 @@ func TestGenerator_Generate_EmptyMode(t *testing.T) { Limit: 5, } - result, err := gen.Generate(nil, input) - assert.NoError(t, err) - assert.NotNil(t, result) + // Without context, should return error + _, err := gen.Generate(nil, input) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") } -func TestBuiltinGenerator_Generate(t *testing.T) { - gen := NewBuiltinGenerator() +func TestGenerator_Generate_AgentMode_RequiresContext(t *testing.T) { + // Custom agent mode requires context + gen := NewGenerator("custom.querydsl.agent", nil) - t.Run("empty query", func(t *testing.T) { - result, err := gen.Generate(&Input{}) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.Nil(t, result.DSL) - assert.Contains(t, result.Warnings, "empty query, returning empty DSL") - }) + input := &Input{ + Query: "find users", + ModelIDs: []string{"user"}, + Limit: 10, + } - t.Run("nil input", func(t *testing.T) { - result, err := gen.Generate(nil) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.Nil(t, result.DSL) - }) + _, err := gen.Generate(nil, input) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") +} - t.Run("basic query without models loaded", func(t *testing.T) { - // Models are loaded internally via model.Select() - // When model is not found, it still generates basic DSL - result, err := gen.Generate(&Input{ - Query: "find users", - ModelIDs: []string{"user"}, - Limit: 10, - }) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.NotNil(t, result.DSL) - assert.Equal(t, 10, result.DSL.Limit) - }) +func TestGenerator_Generate_MCPMode_InvalidFormat(t *testing.T) { + // Invalid MCP format should fallback to system agent (which requires context) + gen := NewGenerator("mcp:invalid", nil) - t.Run("query with pre-defined wheres", func(t *testing.T) { - preWheres := []gou.Where{ - { - Condition: gou.Condition{ - Field: &gou.Expression{Field: "status"}, - OP: "=", - Value: "active", - }, - }, - } - result, err := gen.Generate(&Input{ - Query: "find users", - ModelIDs: []string{"user"}, - Wheres: preWheres, - Limit: 10, - }) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.NotNil(t, result.DSL) - // Should use pre-defined wheres - assert.Equal(t, preWheres, result.DSL.Wheres) - }) + input := &Input{ + Query: "find users", + ModelIDs: []string{"user"}, + Limit: 10, + } - t.Run("query with orders", func(t *testing.T) { - orders := gou.Orders{ - {Field: &gou.Expression{Field: "created_at"}, Sort: "desc"}, - } - result, err := gen.Generate(&Input{ - Query: "find users", - ModelIDs: []string{"user"}, - Orders: orders, - Limit: 10, - }) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.NotNil(t, result.DSL) - assert.Equal(t, orders, result.DSL.Orders) - }) + _, err := gen.Generate(nil, input) + assert.Error(t, err) + assert.Contains(t, err.Error(), "context is required") +} - t.Run("query with allowed fields", func(t *testing.T) { - result, err := gen.Generate(&Input{ - Query: "find users", - ModelIDs: []string{"user"}, - AllowedFields: []string{"id", "name", "email"}, - Limit: 10, - }) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.NotNil(t, result.DSL) - }) - - t.Run("default limit", func(t *testing.T) { - result, err := gen.Generate(&Input{ - Query: "find users", - ModelIDs: []string{"user"}, - }) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.NotNil(t, result.DSL) - assert.Equal(t, 20, result.DSL.Limit) - }) - - t.Run("multi-model query", func(t *testing.T) { - // Models are loaded internally via model.Select() - result, err := gen.Generate(&Input{ - Query: "find user orders", - ModelIDs: []string{"user", "order"}, - Limit: 10, - }) - assert.NoError(t, err) - assert.NotNil(t, result) - assert.NotNil(t, result.DSL) - }) +func TestSystemQueryDSLAgentConstant(t *testing.T) { + // Verify the system querydsl agent constant + assert.Equal(t, "__yao.querydsl", SystemQueryDSLAgent) } func TestResult(t *testing.T) { @@ -201,3 +127,80 @@ func TestResult(t *testing.T) { assert.NotEmpty(t, result.Explain) assert.Len(t, result.Warnings, 1) } + +func TestGenerator_ValidateFields(t *testing.T) { + gen := NewGenerator("", nil) + + t.Run("validate select fields", func(t *testing.T) { + result := &Result{ + DSL: &gou.QueryDSL{ + Select: []gou.Expression{ + {Field: "id"}, + {Field: "name"}, + {Field: "secret_field"}, + }, + }, + } + allowedFields := []string{"id", "name", "email"} + + validated := gen.validateFields(result, allowedFields) + assert.NotNil(t, validated) + assert.Len(t, validated.DSL.Select, 2) + assert.Contains(t, validated.Warnings[0], "secret_field") + }) + + t.Run("validate where fields", func(t *testing.T) { + result := &Result{ + DSL: &gou.QueryDSL{ + Wheres: []gou.Where{ + { + Condition: gou.Condition{ + Field: &gou.Expression{Field: "status"}, + OP: "=", + Value: "active", + }, + }, + { + Condition: gou.Condition{ + Field: &gou.Expression{Field: "secret"}, + OP: "=", + Value: "hidden", + }, + }, + }, + }, + } + allowedFields := []string{"status", "name"} + + validated := gen.validateFields(result, allowedFields) + assert.NotNil(t, validated) + assert.Len(t, validated.DSL.Wheres, 1) + assert.Contains(t, validated.Warnings[0], "secret") + }) + + t.Run("validate order fields", func(t *testing.T) { + result := &Result{ + DSL: &gou.QueryDSL{ + Orders: gou.Orders{ + {Field: &gou.Expression{Field: "created_at"}, Sort: "desc"}, + {Field: &gou.Expression{Field: "secret_sort"}, Sort: "asc"}, + }, + }, + } + allowedFields := []string{"created_at", "updated_at"} + + validated := gen.validateFields(result, allowedFields) + assert.NotNil(t, validated) + assert.Len(t, validated.DSL.Orders, 1) + assert.Contains(t, validated.Warnings[0], "secret_sort") + }) + + t.Run("nil DSL", func(t *testing.T) { + result := &Result{DSL: nil} + allowedFields := []string{"id", "name"} + + validated := gen.validateFields(result, allowedFields) + assert.NotNil(t, validated) + assert.Nil(t, validated.DSL) + }) +}