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.
This commit is contained in:
Max
parent
ac0fa681ea
commit
b8c5829eb0
9 changed files with 188 additions and 709 deletions
|
|
@ -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{
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
// - "<assistant-id>": Delegate to an LLM-powered assistant for high-quality extraction
|
||||
// - "builtin" or "": Uses __yao.keyword system agent (LLM-powered)
|
||||
// - "<assistant-id>": Delegate to a custom LLM-powered assistant
|
||||
// - "mcp:<server>.<tool>": 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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
// - "<assistant-id>": Delegate to an LLM-powered assistant for high-quality generation
|
||||
// - "builtin" or "": Uses __yao.querydsl system agent (LLM-powered)
|
||||
// - "<assistant-id>": Delegate to a custom LLM-powered assistant
|
||||
// - "mcp:<server>.<tool>": 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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue