Merge pull request #1387 from trheyi/main
Enhance Search Functionality with Keyword Extraction and Intent Detection
This commit is contained in:
commit
248579bc65
58 changed files with 3834 additions and 1659 deletions
|
|
@ -16,7 +16,7 @@ import (
|
|||
|
||||
// Stream stream the agent
|
||||
// handler is optional, if not provided, a default handler will be used
|
||||
func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Message, options ...*context.Options) (interface{}, error) {
|
||||
func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Message, options ...*context.Options) (*context.Response, error) {
|
||||
|
||||
// Update logger with assistant ID and start logging
|
||||
ctx.Logger.SetAssistantID(ast.ID)
|
||||
|
|
@ -205,8 +205,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
|
|||
// ================================================
|
||||
// Execute Auto Search (if enabled)
|
||||
// ================================================
|
||||
if ast.shouldAutoSearch(ctx, completionMessages, createResponse, opts) {
|
||||
refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse, opts)
|
||||
if intent := ast.shouldAutoSearch(ctx, completionMessages, createResponse, opts); intent != nil {
|
||||
refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse, intent, opts)
|
||||
if refCtx != nil && len(refCtx.References) > 0 {
|
||||
completionMessages = ast.injectSearchContext(completionMessages, refCtx)
|
||||
}
|
||||
|
|
@ -376,7 +376,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
|
|||
// ================================================
|
||||
// Execute Next Hook and Process Response
|
||||
// ================================================
|
||||
var finalResponse interface{}
|
||||
var finalResponse *context.Response
|
||||
var nextResponse *context.NextHookResponse = nil
|
||||
|
||||
if ast.HookScript != nil {
|
||||
|
|
|
|||
|
|
@ -56,19 +56,18 @@ func TestAgentNextStandard(t *testing.T) {
|
|||
assert.NoError(t, err)
|
||||
assert.NotNil(t, response)
|
||||
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion)
|
||||
assert.Nil(t, resp.Next)
|
||||
assert.NotNil(t, response.Completion)
|
||||
assert.Nil(t, response.Next)
|
||||
|
||||
// Verify response structure
|
||||
assert.Equal(t, "tests.realworld-next", resp.AssistantID)
|
||||
assert.NotEmpty(t, resp.ContextID)
|
||||
assert.NotEmpty(t, resp.RequestID)
|
||||
assert.NotEmpty(t, resp.TraceID)
|
||||
assert.NotEmpty(t, resp.ChatID)
|
||||
assert.Equal(t, "tests.realworld-next", response.AssistantID)
|
||||
assert.NotEmpty(t, response.ContextID)
|
||||
assert.NotEmpty(t, response.RequestID)
|
||||
assert.NotEmpty(t, response.TraceID)
|
||||
assert.NotEmpty(t, response.ChatID)
|
||||
|
||||
// Verify completion has content
|
||||
assert.NotNil(t, resp.Completion.Content)
|
||||
assert.NotNil(t, response.Completion.Content)
|
||||
|
||||
t.Log("✓ Standard response test passed")
|
||||
}
|
||||
|
|
@ -94,19 +93,18 @@ func TestAgentNextCustomData(t *testing.T) {
|
|||
assert.NoError(t, err)
|
||||
assert.NotNil(t, response)
|
||||
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion)
|
||||
assert.NotNil(t, resp.Next)
|
||||
assert.NotNil(t, response.Completion)
|
||||
assert.NotNil(t, response.Next)
|
||||
|
||||
// Verify response structure
|
||||
assert.Equal(t, "tests.realworld-next", resp.AssistantID)
|
||||
assert.NotEmpty(t, resp.ContextID)
|
||||
assert.NotEmpty(t, resp.RequestID)
|
||||
assert.NotEmpty(t, resp.TraceID)
|
||||
assert.Equal(t, "tests.realworld-next", response.AssistantID)
|
||||
assert.NotEmpty(t, response.ContextID)
|
||||
assert.NotEmpty(t, response.RequestID)
|
||||
assert.NotEmpty(t, response.TraceID)
|
||||
|
||||
// Verify custom data structure (from scenarioCustomData)
|
||||
// resp.Next contains the "data" field value from NextHookResponse
|
||||
nextData, ok := resp.Next.(map[string]interface{})
|
||||
// response.Next contains the "data" field value from NextHookResponse
|
||||
nextData, ok := response.Next.(map[string]interface{})
|
||||
assert.True(t, ok, "Next should be a map")
|
||||
assert.Equal(t, "custom_response", nextData["type"])
|
||||
assert.Equal(t, "This is a custom response from Next Hook", nextData["message"])
|
||||
|
|
@ -137,21 +135,19 @@ func TestAgentNextDelegate(t *testing.T) {
|
|||
assert.NoError(t, err)
|
||||
assert.NotNil(t, response)
|
||||
|
||||
resp := response.(*context.Response)
|
||||
|
||||
// Verify response structure
|
||||
assert.NotEmpty(t, resp.AssistantID)
|
||||
assert.NotEmpty(t, resp.ContextID)
|
||||
assert.NotEmpty(t, resp.RequestID)
|
||||
assert.NotEmpty(t, resp.TraceID)
|
||||
assert.NotEmpty(t, response.AssistantID)
|
||||
assert.NotEmpty(t, response.ContextID)
|
||||
assert.NotEmpty(t, response.RequestID)
|
||||
assert.NotEmpty(t, response.TraceID)
|
||||
|
||||
// Verify completion (delegated agent should have returned completion)
|
||||
assert.NotNil(t, resp.Completion)
|
||||
assert.NotNil(t, resp.Completion.Content)
|
||||
assert.NotNil(t, response.Completion)
|
||||
assert.NotNil(t, response.Completion.Content)
|
||||
|
||||
// Next should be from the delegated agent
|
||||
// If delegated agent also has Next hook, it will be present
|
||||
t.Logf("✓ Delegation test passed (delegated to: %s)", resp.AssistantID)
|
||||
t.Logf("✓ Delegation test passed (delegated to: %s)", response.AssistantID)
|
||||
}
|
||||
|
||||
// TestAgentNextConditional tests agent with conditional logic in Next Hook
|
||||
|
|
@ -177,18 +173,17 @@ func TestAgentNextConditional(t *testing.T) {
|
|||
assert.NoError(t, err)
|
||||
assert.NotNil(t, response)
|
||||
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Next)
|
||||
assert.NotNil(t, response.Next)
|
||||
|
||||
// Verify response structure
|
||||
assert.Equal(t, "tests.realworld-next", resp.AssistantID)
|
||||
assert.NotEmpty(t, resp.ContextID)
|
||||
assert.NotEmpty(t, resp.RequestID)
|
||||
assert.NotEmpty(t, resp.TraceID)
|
||||
assert.Equal(t, "tests.realworld-next", response.AssistantID)
|
||||
assert.NotEmpty(t, response.ContextID)
|
||||
assert.NotEmpty(t, response.RequestID)
|
||||
assert.NotEmpty(t, response.TraceID)
|
||||
|
||||
// Verify conditional response structure (from scenarioConditional)
|
||||
// resp.Next contains the "data" field value from NextHookResponse
|
||||
nextData, ok := resp.Next.(map[string]interface{})
|
||||
// response.Next contains the "data" field value from NextHookResponse
|
||||
nextData, ok := response.Next.(map[string]interface{})
|
||||
assert.True(t, ok, "Next should be a map")
|
||||
assert.Equal(t, "Conditional analysis complete", nextData["message"])
|
||||
assert.Contains(t, nextData, "action")
|
||||
|
|
@ -224,19 +219,18 @@ func TestAgentWithoutNextHook(t *testing.T) {
|
|||
assert.NoError(t, err)
|
||||
assert.NotNil(t, response)
|
||||
|
||||
resp := response.(*context.Response)
|
||||
assert.Nil(t, resp.Next)
|
||||
assert.Nil(t, response.Next)
|
||||
|
||||
// Verify response structure
|
||||
assert.Equal(t, "tests.create", resp.AssistantID)
|
||||
assert.NotEmpty(t, resp.ContextID)
|
||||
assert.NotEmpty(t, resp.RequestID)
|
||||
assert.NotEmpty(t, resp.TraceID)
|
||||
assert.NotEmpty(t, resp.ChatID)
|
||||
assert.Equal(t, "tests.create", response.AssistantID)
|
||||
assert.NotEmpty(t, response.ContextID)
|
||||
assert.NotEmpty(t, response.RequestID)
|
||||
assert.NotEmpty(t, response.TraceID)
|
||||
assert.NotEmpty(t, response.ChatID)
|
||||
|
||||
// Verify completion
|
||||
assert.NotNil(t, resp.Completion)
|
||||
assert.NotNil(t, resp.Completion.Content)
|
||||
assert.NotNil(t, response.Completion)
|
||||
assert.NotNil(t, response.Completion.Content)
|
||||
|
||||
t.Log("✓ No Next Hook test passed")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -51,7 +51,7 @@ type agentCallerWrapper struct {
|
|||
ast *Assistant
|
||||
}
|
||||
|
||||
func (w *agentCallerWrapper) Stream(ctx *agentContext.Context, messages []agentContext.Message, options ...*agentContext.Options) (interface{}, error) {
|
||||
func (w *agentCallerWrapper) Stream(ctx *agentContext.Context, messages []agentContext.Message, options ...*agentContext.Options) (*agentContext.Response, error) {
|
||||
return w.ast.Stream(ctx, messages, options...)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import (
|
|||
)
|
||||
|
||||
// processNextResponse processes the Next hook's response and handles agent delegation or custom data
|
||||
func (ast *Assistant) processNextResponse(npc *NextProcessContext) (interface{}, error) {
|
||||
func (ast *Assistant) processNextResponse(npc *NextProcessContext) (*agentContext.Response, error) {
|
||||
// If no Next hook response, return standard response
|
||||
if npc.NextResponse == nil {
|
||||
return ast.buildStandardResponse(npc), nil
|
||||
|
|
@ -42,7 +42,7 @@ func (ast *Assistant) handleDelegation(
|
|||
ctx *agentContext.Context,
|
||||
delegate *agentContext.DelegateConfig,
|
||||
streamHandler func(message.StreamChunkType, []byte) int,
|
||||
) (interface{}, error) {
|
||||
) (*agentContext.Response, error) {
|
||||
// Load the target assistant
|
||||
targetAssistant, err := Get(delegate.AgentID)
|
||||
if err != nil {
|
||||
|
|
@ -62,7 +62,7 @@ func (ast *Assistant) handleDelegation(
|
|||
}
|
||||
|
||||
// buildStandardResponse builds the standard agent response when no custom Next hook processing is needed
|
||||
func (ast *Assistant) buildStandardResponse(npc *NextProcessContext) interface{} {
|
||||
func (ast *Assistant) buildStandardResponse(npc *NextProcessContext) *agentContext.Response {
|
||||
return &agentContext.Response{
|
||||
ContextID: npc.Context.ID,
|
||||
RequestID: npc.Context.RequestID(),
|
||||
|
|
|
|||
|
|
@ -17,16 +17,32 @@ import (
|
|||
)
|
||||
|
||||
// shouldAutoSearch determines if auto search should be executed
|
||||
// Returns false if:
|
||||
// Returns nil if search should be skipped, otherwise returns SearchIntent with types to search
|
||||
// Search is skipped if:
|
||||
// - opts.Skip.Search is true
|
||||
// - createResponse.Search is false
|
||||
// - uses.search is "disabled"
|
||||
// - assistant has no search configuration
|
||||
// - needsearch intent detection returns false
|
||||
func (ast *Assistant) shouldAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, opts *context.Options) bool {
|
||||
func (ast *Assistant) shouldAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, opts *context.Options) *SearchIntent {
|
||||
// Check if search is skipped via options
|
||||
if opts != nil && opts.Skip != nil && opts.Skip.Search {
|
||||
ctx.Logger.Debug("Auto search skipped by opts.Skip.Search")
|
||||
return false
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check createResponse.Search field (highest priority from Create hook)
|
||||
// Supports: bool | SearchIntent | nil
|
||||
if createResponse != nil && createResponse.Search != nil {
|
||||
intent := parseSearchField(createResponse.Search)
|
||||
if intent != nil {
|
||||
if !intent.NeedSearch {
|
||||
ctx.Logger.Info("Auto search disabled by createResponse.Search")
|
||||
return nil
|
||||
}
|
||||
ctx.Logger.Info("Auto search controlled by createResponse.Search: types=%v", intent.SearchTypes)
|
||||
return intent
|
||||
}
|
||||
}
|
||||
|
||||
// Get merged uses configuration
|
||||
|
|
@ -35,57 +51,123 @@ func (ast *Assistant) shouldAutoSearch(ctx *context.Context, messages []context.
|
|||
// Check if search is explicitly disabled
|
||||
if uses != nil && uses.Search == "disabled" {
|
||||
ctx.Logger.Info("Auto search disabled by uses.search=disabled")
|
||||
return false
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if assistant has search configuration
|
||||
if ast.Search == nil && (uses == nil || uses.Search == "") {
|
||||
return false
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check search intent using __yao.needsearch agent
|
||||
if !ast.checkSearchIntent(ctx, messages) {
|
||||
intent := ast.checkSearchIntent(ctx, messages)
|
||||
if intent == nil || !intent.NeedSearch {
|
||||
ctx.Logger.Info("Auto search skipped: intent detection returned false")
|
||||
return false
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if search is enabled (builtin, agent, mcp, or empty means builtin)
|
||||
return true
|
||||
return intent
|
||||
}
|
||||
|
||||
// parseSearchField parses the Search field from HookCreateResponse
|
||||
// Supports: bool | SearchIntent | map[string]any | nil
|
||||
func parseSearchField(search any) *SearchIntent {
|
||||
if search == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch v := search.(type) {
|
||||
case bool:
|
||||
// bool: true = enable all, false = disable all
|
||||
if v {
|
||||
return &SearchIntent{
|
||||
NeedSearch: true,
|
||||
SearchTypes: []string{"web", "kb", "db"},
|
||||
Confidence: 1.0,
|
||||
Reason: "enabled by hook",
|
||||
}
|
||||
}
|
||||
return &SearchIntent{
|
||||
NeedSearch: false,
|
||||
SearchTypes: []string{},
|
||||
Confidence: 1.0,
|
||||
Reason: "disabled by hook",
|
||||
}
|
||||
|
||||
case *SearchIntent:
|
||||
// SearchIntent is alias for context.SearchIntent, so this covers both
|
||||
return v
|
||||
|
||||
case SearchIntent:
|
||||
return &v
|
||||
|
||||
case map[string]any:
|
||||
// Parse from map (e.g., from JSON)
|
||||
intent := &SearchIntent{
|
||||
NeedSearch: false,
|
||||
SearchTypes: []string{},
|
||||
Confidence: 0.5,
|
||||
}
|
||||
|
||||
if needSearch, ok := v["need_search"].(bool); ok {
|
||||
intent.NeedSearch = needSearch
|
||||
}
|
||||
|
||||
if types, ok := v["search_types"].([]any); ok {
|
||||
for _, t := range types {
|
||||
if typeStr, ok := t.(string); ok {
|
||||
intent.SearchTypes = append(intent.SearchTypes, typeStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if confidence, ok := v["confidence"].(float64); ok {
|
||||
intent.Confidence = confidence
|
||||
}
|
||||
|
||||
if reason, ok := v["reason"].(string); ok {
|
||||
intent.Reason = reason
|
||||
}
|
||||
|
||||
return intent
|
||||
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
// Returns SearchIntent with search types and confidence
|
||||
func (ast *Assistant) checkSearchIntent(ctx *context.Context, messages []context.Message) *SearchIntent {
|
||||
// Default intent: no search needed (fallback when agent unavailable or fails)
|
||||
defaultIntent := &SearchIntent{
|
||||
NeedSearch: false,
|
||||
SearchTypes: []string{},
|
||||
Confidence: 0,
|
||||
}
|
||||
|
||||
// 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 defaultIntent // No messages, skip search
|
||||
}
|
||||
|
||||
// Try to get __yao.needsearch agent
|
||||
needsearchAst, err := Get("__yao.needsearch")
|
||||
if err != nil {
|
||||
ctx.Logger.Debug("__yao.needsearch agent not available: %v, proceeding with search", err)
|
||||
return true // Agent not available, proceed with search
|
||||
ctx.Logger.Debug("__yao.needsearch agent not available: %v, skipping search", err)
|
||||
return defaultIntent // Agent not available, skip search
|
||||
}
|
||||
|
||||
// === 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{
|
||||
|
|
@ -98,58 +180,108 @@ func (ast *Assistant) checkSearchIntent(ctx *context.Context, messages []context
|
|||
|
||||
result, err := needsearchAst.Stream(ctx, intentMessages, opts)
|
||||
if err != nil {
|
||||
ctx.Logger.Debug("__yao.needsearch failed: %v, proceeding with search", err)
|
||||
// === Output: Send done (error case, proceed with search) ===
|
||||
ast.sendIntentDone(ctx, loadingID, true, "")
|
||||
return true // On error, proceed with search
|
||||
ctx.Logger.Debug("__yao.needsearch failed: %v, skipping search", err)
|
||||
// === Output: Send done (error case, skip search) ===
|
||||
ast.sendIntentDone(ctx, loadingID, false, "")
|
||||
return defaultIntent // On error, skip search
|
||||
}
|
||||
|
||||
// Parse the result
|
||||
// Next hook returns {data: {need_search: bool, search_types: [], confidence: float}}
|
||||
if response, ok := result.(*context.Response); ok {
|
||||
// First try to get from Next hook response
|
||||
if response.Next != nil {
|
||||
if nextData, ok := response.Next.(map[string]interface{}); ok {
|
||||
// Check for data field (from Next hook's {data: result})
|
||||
var intentData map[string]interface{}
|
||||
if data, ok := nextData["data"].(map[string]interface{}); ok {
|
||||
intentData = data
|
||||
} else {
|
||||
intentData = nextData
|
||||
}
|
||||
|
||||
if needSearch, ok := intentData["need_search"].(bool); ok {
|
||||
reason, _ := intentData["reason"].(string)
|
||||
ctx.Logger.Debug("Search intent (from Next): need_search=%v, reason=%s", needSearch, reason)
|
||||
ast.sendIntentDone(ctx, loadingID, needSearch, reason)
|
||||
return needSearch
|
||||
}
|
||||
// First try to get from Next hook response
|
||||
if result.Next != nil {
|
||||
if nextData, ok := result.Next.(map[string]interface{}); ok {
|
||||
// Check for data field (from Next hook's {data: result})
|
||||
var intentData map[string]interface{}
|
||||
if data, ok := nextData["data"].(map[string]interface{}); ok {
|
||||
intentData = data
|
||||
} else {
|
||||
intentData = nextData
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: parse from Completion.Content if Next hook didn't process
|
||||
if response.Completion != nil {
|
||||
content, ok := response.Completion.Content.(string)
|
||||
if !ok || content == "" {
|
||||
ast.sendIntentDone(ctx, loadingID, true, "")
|
||||
return true
|
||||
intent := parseSearchIntent(intentData)
|
||||
if intent != nil {
|
||||
ctx.Logger.Debug("Search intent (from Next): need_search=%v, types=%v, confidence=%.2f, reason=%s",
|
||||
intent.NeedSearch, intent.SearchTypes, intent.Confidence, intent.Reason)
|
||||
ast.sendIntentDone(ctx, loadingID, intent.NeedSearch, intent.Reason)
|
||||
return intent
|
||||
}
|
||||
needSearch, reason := parseNeedSearchFromContent(content)
|
||||
ctx.Logger.Debug("Search intent (from Content): need_search=%v, reason=%s", needSearch, reason)
|
||||
ast.sendIntentDone(ctx, loadingID, needSearch, reason)
|
||||
return needSearch
|
||||
}
|
||||
}
|
||||
|
||||
// Default: proceed with search if we can't parse the result
|
||||
// Fallback: parse from Completion.Content if Next hook didn't process
|
||||
if result.Completion != nil {
|
||||
content, ok := result.Completion.Content.(string)
|
||||
if !ok || content == "" {
|
||||
ast.sendIntentDone(ctx, loadingID, false, "")
|
||||
return defaultIntent
|
||||
}
|
||||
intent := parseSearchIntentFromContent(content)
|
||||
ctx.Logger.Debug("Search intent (from Content): need_search=%v, types=%v, confidence=%.2f, reason=%s",
|
||||
intent.NeedSearch, intent.SearchTypes, intent.Confidence, intent.Reason)
|
||||
ast.sendIntentDone(ctx, loadingID, intent.NeedSearch, intent.Reason)
|
||||
return intent
|
||||
}
|
||||
|
||||
// Default: skip search if we can't parse the result
|
||||
// === Output: Send done (default case) ===
|
||||
ast.sendIntentDone(ctx, loadingID, true, "")
|
||||
return true
|
||||
ast.sendIntentDone(ctx, loadingID, false, "")
|
||||
return defaultIntent
|
||||
}
|
||||
|
||||
// parseNeedSearchFromContent parses need_search result from LLM completion content
|
||||
// parseSearchIntent parses SearchIntent from intent data map
|
||||
func parseSearchIntent(intentData map[string]interface{}) *SearchIntent {
|
||||
if intentData == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
needSearch, ok := intentData["need_search"].(bool)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
intent := &SearchIntent{
|
||||
NeedSearch: needSearch,
|
||||
SearchTypes: []string{},
|
||||
Confidence: 0.5, // Default confidence
|
||||
}
|
||||
|
||||
// Parse search_types
|
||||
if types, ok := intentData["search_types"].([]interface{}); ok {
|
||||
for _, t := range types {
|
||||
if typeStr, ok := t.(string); ok {
|
||||
// Validate type
|
||||
typeStr = strings.ToLower(typeStr)
|
||||
if typeStr == "web" || typeStr == "kb" || typeStr == "db" {
|
||||
intent.SearchTypes = append(intent.SearchTypes, typeStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Parse confidence
|
||||
if confidence, ok := intentData["confidence"].(float64); ok {
|
||||
intent.Confidence = confidence
|
||||
}
|
||||
|
||||
// Parse reason
|
||||
if reason, ok := intentData["reason"].(string); ok {
|
||||
intent.Reason = reason
|
||||
}
|
||||
|
||||
return intent
|
||||
}
|
||||
|
||||
// parseSearchIntentFromContent parses SearchIntent from LLM completion content
|
||||
// Handles JSON wrapped in markdown code blocks
|
||||
func parseNeedSearchFromContent(content string) (bool, string) {
|
||||
func parseSearchIntentFromContent(content string) *SearchIntent {
|
||||
// Default intent: no search needed
|
||||
defaultIntent := &SearchIntent{
|
||||
NeedSearch: false,
|
||||
SearchTypes: []string{},
|
||||
Confidence: 0,
|
||||
}
|
||||
|
||||
// Remove markdown code block if present
|
||||
content = strings.TrimSpace(content)
|
||||
if strings.HasPrefix(content, "```json") {
|
||||
|
|
@ -165,17 +297,16 @@ func parseNeedSearchFromContent(content string) (bool, string) {
|
|||
// Try to parse JSON
|
||||
var result map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(content), &result); err != nil {
|
||||
// Failed to parse, default to search
|
||||
return true, ""
|
||||
// Failed to parse, default to no search
|
||||
return defaultIntent
|
||||
}
|
||||
|
||||
needSearch, ok := result["need_search"].(bool)
|
||||
if !ok {
|
||||
return true, ""
|
||||
intent := parseSearchIntent(result)
|
||||
if intent == nil {
|
||||
return defaultIntent
|
||||
}
|
||||
|
||||
reason, _ := result["reason"].(string)
|
||||
return needSearch, reason
|
||||
return intent
|
||||
}
|
||||
|
||||
// sendIntentLoading sends the initial intent detection loading message
|
||||
|
|
@ -271,10 +402,11 @@ func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResp
|
|||
return uses
|
||||
}
|
||||
|
||||
// executeAutoSearch executes auto search based on configuration
|
||||
// executeAutoSearch executes auto search based on configuration and intent
|
||||
// Returns ReferenceContext with results and formatted context
|
||||
// intent specifies which search types to execute (from needsearch agent)
|
||||
// opts is optional, used to check Skip.Keyword
|
||||
func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, opts ...*context.Options) *searchTypes.ReferenceContext {
|
||||
func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, intent *SearchIntent, opts ...*context.Options) *searchTypes.ReferenceContext {
|
||||
ctx.Logger.Phase("Search")
|
||||
defer ctx.Logger.PhaseComplete("Search")
|
||||
|
||||
|
|
@ -311,33 +443,23 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context
|
|||
skipKeyword = opts[0].Skip.Keyword
|
||||
}
|
||||
|
||||
// Extract keywords for web search if:
|
||||
// 1. uses.keyword is configured (not empty)
|
||||
// 2. Skip.Keyword is not true
|
||||
// 3. Web search is enabled
|
||||
var extractedKeywords []string
|
||||
webSearchEnabled := searchConfig != nil && searchConfig.Web != nil
|
||||
if webSearchEnabled && !skipKeyword && searchUses.Keyword != "" {
|
||||
extractor := keyword.NewExtractor(searchUses.Keyword, searchConfig.Keyword)
|
||||
keywords, err := extractor.Extract(ctx, query, nil)
|
||||
if err != nil {
|
||||
ctx.Logger.Warn("Keyword extraction failed, using original query: %v", err)
|
||||
} else if len(keywords) > 0 {
|
||||
extractedKeywords = keywords
|
||||
// Use extracted keywords as the search query for web search
|
||||
optimizedQuery := strings.Join(keywords, " ")
|
||||
ctx.Logger.Info("Extracted keywords for web search: %s -> %s", truncateString(query, 30), optimizedQuery)
|
||||
query = optimizedQuery
|
||||
}
|
||||
// Build search requests based on configuration and intent
|
||||
// Keyword extraction is done inside buildSearchRequests for web search
|
||||
buildOpts := &buildSearchRequestsOptions{
|
||||
skipKeyword: skipKeyword,
|
||||
usesKeyword: searchUses.Keyword,
|
||||
}
|
||||
|
||||
// Build search requests based on configuration
|
||||
requests := ast.buildSearchRequests(query, searchConfig)
|
||||
requests, extractedKeywords := ast.buildSearchRequests(ctx, query, searchConfig, intent, buildOpts)
|
||||
if len(requests) == 0 {
|
||||
ctx.Logger.Info("No search requests to execute")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Update query if keywords were extracted (for web search)
|
||||
if len(extractedKeywords) > 0 {
|
||||
query = keywordsToQuery(extractedKeywords)
|
||||
}
|
||||
|
||||
// === Output: Send loading message ===
|
||||
loadingID := ast.sendSearchLoading(ctx)
|
||||
|
||||
|
|
@ -435,6 +557,52 @@ func (ast *Assistant) sendSearchLoading(ctx *context.Context) string {
|
|||
return msgID
|
||||
}
|
||||
|
||||
// sendKeywordLoading sends the keyword extraction loading message
|
||||
// Returns the message ID for later replacement
|
||||
func (ast *Assistant) sendKeywordLoading(ctx *context.Context) string {
|
||||
loadingMsg := i18n.T(ctx.Locale, "search.keyword.loading")
|
||||
|
||||
msg := &message.Message{
|
||||
Type: "loading",
|
||||
Props: map[string]any{
|
||||
"message": loadingMsg,
|
||||
},
|
||||
}
|
||||
|
||||
// Send and get message ID
|
||||
msgID, err := ctx.SendStream(msg)
|
||||
if err != nil {
|
||||
ctx.Logger.Warn("Failed to send keyword loading message: %v", err)
|
||||
return ""
|
||||
}
|
||||
|
||||
return msgID
|
||||
}
|
||||
|
||||
// sendKeywordDone replaces keyword loading with done message
|
||||
func (ast *Assistant) sendKeywordDone(ctx *context.Context, loadingID string, success bool) {
|
||||
if loadingID == "" {
|
||||
return
|
||||
}
|
||||
|
||||
resultMsg := i18n.T(ctx.Locale, "search.keyword.done")
|
||||
|
||||
msg := &message.Message{
|
||||
MessageID: loadingID,
|
||||
Delta: true,
|
||||
DeltaAction: message.DeltaReplace,
|
||||
Type: "loading",
|
||||
Props: map[string]any{
|
||||
"message": resultMsg,
|
||||
"done": true,
|
||||
},
|
||||
}
|
||||
|
||||
if err := ctx.Send(msg); err != nil {
|
||||
ctx.Logger.Warn("Failed to send keyword done message: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// sendSearchResult replaces loading with result message (without done flag)
|
||||
func (ast *Assistant) sendSearchResult(ctx *context.Context, loadingID string, count int) {
|
||||
if loadingID == "" {
|
||||
|
|
@ -564,22 +732,67 @@ func (ast *Assistant) completeSearchTrace(node traceTypes.Node, resultCount int,
|
|||
})
|
||||
}
|
||||
|
||||
// buildSearchRequests builds search requests based on assistant configuration
|
||||
func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Config) []*searchTypes.Request {
|
||||
var requests []*searchTypes.Request
|
||||
// buildSearchRequestsOptions contains options for building search requests
|
||||
type buildSearchRequestsOptions struct {
|
||||
skipKeyword bool // Skip keyword extraction
|
||||
usesKeyword string // Keyword extractor config: "builtin", "<assistant-id>", "mcp:<server>.<tool>"
|
||||
}
|
||||
|
||||
// buildSearchRequests builds search requests based on assistant configuration and intent
|
||||
// intent specifies which search types to execute (from needsearch agent)
|
||||
// Returns requests and extracted keywords (if any)
|
||||
func (ast *Assistant) buildSearchRequests(ctx *context.Context, query string, config *searchTypes.Config, intent *SearchIntent, opts *buildSearchRequestsOptions) ([]*searchTypes.Request, []searchTypes.Keyword) {
|
||||
var requests []*searchTypes.Request
|
||||
var extractedKeywords []searchTypes.Keyword
|
||||
|
||||
// Helper to check if a search type is allowed by intent
|
||||
isTypeAllowed := func(searchType string) bool {
|
||||
if intent == nil || len(intent.SearchTypes) == 0 {
|
||||
return true // No intent or empty types means all types allowed
|
||||
}
|
||||
for _, t := range intent.SearchTypes {
|
||||
if t == searchType {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Web search - check if web search is configured and allowed by intent
|
||||
if config != nil && config.Web != nil && isTypeAllowed("web") {
|
||||
webQuery := query
|
||||
|
||||
// Extract keywords for web search if configured
|
||||
if opts != nil && !opts.skipKeyword && opts.usesKeyword != "" {
|
||||
// === Output: Send keyword extraction loading ===
|
||||
keywordLoadingID := ast.sendKeywordLoading(ctx)
|
||||
|
||||
extractor := keyword.NewExtractor(opts.usesKeyword, config.Keyword)
|
||||
keywords, err := extractor.Extract(ctx, query, nil)
|
||||
if err != nil {
|
||||
ctx.Logger.Warn("Keyword extraction failed, using original query: %v", err)
|
||||
ast.sendKeywordDone(ctx, keywordLoadingID, false)
|
||||
} else if len(keywords) > 0 {
|
||||
extractedKeywords = keywords
|
||||
// Use extracted keywords as the search query for web search
|
||||
webQuery = keywordsToQuery(keywords)
|
||||
ctx.Logger.Info("Extracted keywords for web search: %s -> %s", truncateString(query, 30), webQuery)
|
||||
ast.sendKeywordDone(ctx, keywordLoadingID, true)
|
||||
} else {
|
||||
ast.sendKeywordDone(ctx, keywordLoadingID, true)
|
||||
}
|
||||
}
|
||||
|
||||
// Web search - check if web search is configured
|
||||
if config != nil && config.Web != nil {
|
||||
requests = append(requests, &searchTypes.Request{
|
||||
Type: searchTypes.SearchTypeWeb,
|
||||
Query: query,
|
||||
Query: webQuery,
|
||||
Source: searchTypes.SourceAuto,
|
||||
Limit: config.Web.MaxResults,
|
||||
})
|
||||
}
|
||||
|
||||
// KB search - check if KB is configured
|
||||
if ast.KB != nil && len(ast.KB.Collections) > 0 {
|
||||
// KB search - check if KB is configured and allowed by intent
|
||||
if ast.KB != nil && len(ast.KB.Collections) > 0 && isTypeAllowed("kb") {
|
||||
limit := 10
|
||||
threshold := 0.7
|
||||
if config != nil && config.KB != nil {
|
||||
|
|
@ -589,7 +802,7 @@ func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Conf
|
|||
}
|
||||
requests = append(requests, &searchTypes.Request{
|
||||
Type: searchTypes.SearchTypeKB,
|
||||
Query: query,
|
||||
Query: query, // KB uses original query for semantic search
|
||||
Source: searchTypes.SourceAuto,
|
||||
Limit: limit,
|
||||
Collections: ast.KB.Collections,
|
||||
|
|
@ -598,22 +811,22 @@ func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Conf
|
|||
})
|
||||
}
|
||||
|
||||
// DB search - check if DB is configured
|
||||
if ast.DB != nil && len(ast.DB.Models) > 0 {
|
||||
// DB search - check if DB is configured and allowed by intent
|
||||
if ast.DB != nil && len(ast.DB.Models) > 0 && isTypeAllowed("db") {
|
||||
limit := 20
|
||||
if config != nil && config.DB != nil && config.DB.MaxResults > 0 {
|
||||
limit = config.DB.MaxResults
|
||||
}
|
||||
requests = append(requests, &searchTypes.Request{
|
||||
Type: searchTypes.SearchTypeDB,
|
||||
Query: query,
|
||||
Query: query, // DB uses original query for QueryDSL generation
|
||||
Source: searchTypes.SourceAuto,
|
||||
Limit: limit,
|
||||
Models: ast.DB.Models,
|
||||
})
|
||||
}
|
||||
|
||||
return requests
|
||||
return requests, extractedKeywords
|
||||
}
|
||||
|
||||
// injectSearchContext injects search results into messages
|
||||
|
|
@ -709,7 +922,7 @@ func truncateString(s string, maxLen int) string {
|
|||
// SearchExecutionResult holds all data from search execution for storage
|
||||
type SearchExecutionResult struct {
|
||||
Query string // Original query (before keyword optimization)
|
||||
Keywords []string // Extracted keywords
|
||||
Keywords []searchTypes.Keyword // Extracted keywords with weights
|
||||
Config map[string]any // Search config used
|
||||
RefCtx *searchTypes.ReferenceContext // Reference context with results
|
||||
Duration int64 // Search duration in ms
|
||||
|
|
@ -717,6 +930,54 @@ type SearchExecutionResult struct {
|
|||
SearchType string // "auto", "web", "kb", "db"
|
||||
}
|
||||
|
||||
// keywordsToQuery converts keywords with weights to a search query string
|
||||
// Keywords are sorted by weight (descending) and joined with spaces
|
||||
func keywordsToQuery(keywords []searchTypes.Keyword) string {
|
||||
if len(keywords) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Sort by weight descending (higher weight first)
|
||||
sorted := make([]searchTypes.Keyword, len(keywords))
|
||||
copy(sorted, keywords)
|
||||
for i := 0; i < len(sorted)-1; i++ {
|
||||
for j := i + 1; j < len(sorted); j++ {
|
||||
if sorted[j].W > sorted[i].W {
|
||||
sorted[i], sorted[j] = sorted[j], sorted[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Join keywords
|
||||
parts := make([]string, len(sorted))
|
||||
for i, kw := range sorted {
|
||||
parts[i] = kw.K
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
|
||||
// keywordsToStrings converts keywords to string slice for storage
|
||||
func keywordsToStrings(keywords []searchTypes.Keyword) []string {
|
||||
if len(keywords) == 0 {
|
||||
return nil
|
||||
}
|
||||
result := make([]string, len(keywords))
|
||||
for i, kw := range keywords {
|
||||
result[i] = kw.K
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// containsSearchType checks if a search type is in the list
|
||||
func containsSearchType(types []string, searchType string) bool {
|
||||
for _, t := range types {
|
||||
if t == searchType {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// saveSearch saves search results to storage
|
||||
// Called after search execution completes (success or failure)
|
||||
func (ast *Assistant) saveSearch(ctx *context.Context, execResult *SearchExecutionResult) {
|
||||
|
|
@ -732,7 +993,7 @@ func (ast *Assistant) saveSearch(ctx *context.Context, execResult *SearchExecuti
|
|||
RequestID: ctx.RequestID(),
|
||||
ChatID: ctx.ChatID,
|
||||
Query: execResult.Query,
|
||||
Keywords: execResult.Keywords,
|
||||
Keywords: keywordsToStrings(execResult.Keywords),
|
||||
Config: execResult.Config,
|
||||
Source: execResult.SearchType,
|
||||
Duration: execResult.Duration,
|
||||
|
|
|
|||
|
|
@ -77,8 +77,7 @@ func TestSearchAutoDisabled(t *testing.T) {
|
|||
require.NoError(t, err)
|
||||
require.NotNil(t, response)
|
||||
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion, "should have completion")
|
||||
assert.NotNil(t, response.Completion, "should have completion")
|
||||
t.Logf("✓ Stream executed without search (disabled)")
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -118,8 +118,7 @@ func TestSearchAutoFull(t *testing.T) {
|
|||
}
|
||||
|
||||
require.NotNil(t, response)
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion, "should have completion")
|
||||
assert.NotNil(t, response.Completion, "should have completion")
|
||||
t.Logf("✓ Stream executed with full search config (Web + KB + DB)")
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -101,8 +101,7 @@ func TestSearchAutoHookDisable(t *testing.T) {
|
|||
require.NoError(t, err)
|
||||
require.NotNil(t, response)
|
||||
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion, "should have completion")
|
||||
assert.NotNil(t, response.Completion, "should have completion")
|
||||
t.Logf("✓ Stream executed with hook disabling search")
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -87,8 +87,7 @@ func TestSearchAutoKeyword(t *testing.T) {
|
|||
}
|
||||
|
||||
require.NotNil(t, response)
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion, "should have completion")
|
||||
assert.NotNil(t, response.Completion, "should have completion")
|
||||
t.Logf("✓ Stream with keyword extraction executed successfully")
|
||||
})
|
||||
|
||||
|
|
@ -127,8 +126,7 @@ func TestSearchAutoKeyword(t *testing.T) {
|
|||
}
|
||||
|
||||
require.NotNil(t, response)
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion, "should have completion")
|
||||
assert.NotNil(t, response.Completion, "should have completion")
|
||||
t.Logf("✓ Stream with Skip.Keyword executed successfully")
|
||||
})
|
||||
}
|
||||
|
|
@ -179,8 +177,7 @@ func TestSearchAutoKeywordNotConfigured(t *testing.T) {
|
|||
}
|
||||
|
||||
require.NotNil(t, response)
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion, "should have completion")
|
||||
assert.NotNil(t, response.Completion, "should have completion")
|
||||
t.Logf("✓ Stream without keyword config executed successfully")
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -96,8 +96,7 @@ func TestSearchAutoWeb(t *testing.T) {
|
|||
}
|
||||
|
||||
require.NotNil(t, response)
|
||||
resp := response.(*context.Response)
|
||||
assert.NotNil(t, resp.Completion, "should have completion")
|
||||
assert.NotNil(t, response.Completion, "should have completion")
|
||||
t.Logf("✓ Stream executed successfully with auto search")
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -85,6 +85,10 @@ type NextProcessContext struct {
|
|||
CreateResponse *chatctx.HookCreateResponse // Create hook response
|
||||
}
|
||||
|
||||
// SearchIntent is an alias for context.SearchIntent
|
||||
// Used for search intent detection from __yao.needsearch agent
|
||||
type SearchIntent = chatctx.SearchIntent
|
||||
|
||||
// ParsedContent extracts the actual tool return value from MCP ToolContent array
|
||||
// According to MCP protocol:
|
||||
// - Content is []ToolContent array
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import (
|
|||
// AgentCaller interface for calling agents (to avoid circular dependency)
|
||||
// Used by content handlers (vision, audio, etc.) and search handlers (agent mode)
|
||||
type AgentCaller interface {
|
||||
Stream(ctx *agentContext.Context, messages []agentContext.Message, options ...*agentContext.Options) (interface{}, error)
|
||||
Stream(ctx *agentContext.Context, messages []agentContext.Message, options ...*agentContext.Options) (*agentContext.Response, error)
|
||||
}
|
||||
|
||||
// AgentGetterFunc is a function type that gets an agent by ID
|
||||
|
|
|
|||
|
|
@ -42,10 +42,10 @@ func CallAgent(ctx *agentContext.Context, agentID string, message agentContext.M
|
|||
|
||||
// Extract text from agent response
|
||||
// Two formats are supported:
|
||||
// 1. Custom Hook response (from Next hook)
|
||||
// 2. Standard Agent Stream response (LLM completion)
|
||||
// 1. Custom Hook response (from Next hook) - response.Next
|
||||
// 2. Standard Agent Stream response (LLM completion) - response.Completion
|
||||
|
||||
return extractTextFromAgentResponse(response)
|
||||
return extractTextFromResponse(response)
|
||||
}
|
||||
|
||||
// CallAgentWithFileInfo calls an agent to process content with file metadata
|
||||
|
|
@ -120,116 +120,63 @@ func CallAgentWithFileInfo(ctx *agentContext.Context, agentID string, message ag
|
|||
return CallAgent(ctx, agentID, message)
|
||||
}
|
||||
|
||||
// extractTextFromAgentResponse extracts text from agent response
|
||||
// Handles two main response formats from agent.Stream():
|
||||
//
|
||||
// 1. Standard Response (No Next Hook or Next Hook returns nil):
|
||||
// Structure: { completion: { content: "text" | [...ContentPart] } }
|
||||
// Action: Extract text from completion.content field
|
||||
//
|
||||
// 2. Next Hook Response with Custom Data:
|
||||
// Structure: { next: <any data from Next hook> }
|
||||
// Action:
|
||||
// - If next is string → return directly
|
||||
// - If next is map/object → JSON stringify and return
|
||||
// - This preserves the complete custom data structure from the hook
|
||||
// extractTextFromResponse extracts text from agent response
|
||||
// Now that agent.Stream() returns *agentContext.Response directly,
|
||||
// we can access fields without type assertions or JSON conversion.
|
||||
//
|
||||
// Priority:
|
||||
// 1. Check for "next" field (custom hook data) → return complete data
|
||||
// 2. Check for "completion" field (standard LLM response) → extract text only
|
||||
// 3. Fallback to direct string or JSON stringify
|
||||
func extractTextFromAgentResponse(response interface{}) (string, error) {
|
||||
// 1. Check response.Next (custom hook data) → return complete data
|
||||
// 2. Check response.Completion (standard LLM response) → extract text only
|
||||
func extractTextFromResponse(response *agentContext.Response) (string, error) {
|
||||
if response == nil {
|
||||
return "", fmt.Errorf("agent returned nil response")
|
||||
}
|
||||
|
||||
// First, try to convert to map if it's a struct
|
||||
// agent.Stream() may return *agentContext.Response which needs to be converted
|
||||
var responseMap map[string]interface{}
|
||||
|
||||
// Check if it's already a map
|
||||
if rm, ok := response.(map[string]interface{}); ok {
|
||||
responseMap = rm
|
||||
} else {
|
||||
// Try to marshal and unmarshal to convert struct to map
|
||||
jsonBytes, err := jsoniter.Marshal(response)
|
||||
if err != nil {
|
||||
// If it's a plain string, return directly
|
||||
if responseStr, ok := response.(string); ok {
|
||||
return responseStr, nil
|
||||
}
|
||||
return "", fmt.Errorf("failed to serialize agent response: %w", err)
|
||||
}
|
||||
|
||||
// Unmarshal to map
|
||||
if err := jsoniter.Unmarshal(jsonBytes, &responseMap); err != nil {
|
||||
// If unmarshal fails, return the JSON string
|
||||
return string(jsonBytes), nil
|
||||
}
|
||||
}
|
||||
|
||||
// Priority 1: Check for "next" field (custom hook data)
|
||||
// If Next hook returns custom data, it's stored in the "next" field
|
||||
// Return the complete custom data structure (preserve hook's intent)
|
||||
if next, hasNext := responseMap["next"]; hasNext && next != nil {
|
||||
// Priority 1: Check Next field (custom hook data)
|
||||
// If Next hook returns custom data, return the complete structure
|
||||
if response.Next != nil {
|
||||
// If next is a string, return directly
|
||||
if nextStr, ok := next.(string); ok {
|
||||
if nextStr, ok := response.Next.(string); ok {
|
||||
return nextStr, nil
|
||||
}
|
||||
// Otherwise, JSON stringify to preserve complete structure
|
||||
jsonBytes, err := jsoniter.Marshal(next)
|
||||
jsonBytes, err := jsoniter.Marshal(response.Next)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to serialize next hook data: %w", err)
|
||||
}
|
||||
return string(jsonBytes), nil
|
||||
}
|
||||
|
||||
// Priority 2: Check for "completion" field (standard LLM response)
|
||||
// Priority 2: Check Completion field (standard LLM response)
|
||||
// Extract text content from the LLM completion
|
||||
if completion, hasCompletion := responseMap["completion"]; hasCompletion && completion != nil {
|
||||
if completionMap, ok := completion.(map[string]interface{}); ok {
|
||||
// Extract content from completion
|
||||
if content, hasContent := completionMap["content"]; hasContent {
|
||||
// Content can be string or []ContentPart (multimodal)
|
||||
switch v := content.(type) {
|
||||
case string:
|
||||
// Simple text content
|
||||
return v, nil
|
||||
case []interface{}:
|
||||
// Multimodal content array - extract all text parts
|
||||
var text string
|
||||
for _, part := range v {
|
||||
if partMap, ok := part.(map[string]interface{}); ok {
|
||||
if partType, _ := partMap["type"].(string); partType == "text" {
|
||||
if textContent, ok := partMap["text"].(string); ok {
|
||||
text += textContent
|
||||
}
|
||||
}
|
||||
if response.Completion != nil {
|
||||
// Content can be string or []ContentPart (multimodal)
|
||||
switch v := response.Completion.Content.(type) {
|
||||
case string:
|
||||
// Simple text content
|
||||
return v, nil
|
||||
case []interface{}:
|
||||
// Multimodal content array - extract all text parts
|
||||
var text string
|
||||
for _, part := range v {
|
||||
if partMap, ok := part.(map[string]interface{}); ok {
|
||||
if partType, _ := partMap["type"].(string); partType == "text" {
|
||||
if textContent, ok := partMap["text"].(string); ok {
|
||||
text += textContent
|
||||
}
|
||||
}
|
||||
if text != "" {
|
||||
return text, nil
|
||||
}
|
||||
// No text found in content parts
|
||||
return "", fmt.Errorf("no text content found in completion content parts")
|
||||
}
|
||||
}
|
||||
if text != "" {
|
||||
return text, nil
|
||||
}
|
||||
// No text found in content parts
|
||||
return "", fmt.Errorf("no text content found in completion content parts")
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: Try to find a "content" field directly (shouldn't happen normally)
|
||||
if content, hasContent := responseMap["content"]; hasContent {
|
||||
if contentStr, ok := content.(string); ok {
|
||||
return contentStr, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Last resort: JSON stringify the entire response
|
||||
jsonBytes, err := jsoniter.Marshal(response)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to serialize agent response: %w", err)
|
||||
}
|
||||
return string(jsonBytes), nil
|
||||
// No content found
|
||||
return "", fmt.Errorf("no content found in agent response")
|
||||
}
|
||||
|
||||
// CallMCPTool calls an MCP tool to process content
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ func (opts *Options) ToMap() map[string]interface{} {
|
|||
result["mode"] = opts.Mode
|
||||
}
|
||||
if opts.Search != nil {
|
||||
result["search"] = *opts.Search
|
||||
result["search"] = opts.Search
|
||||
}
|
||||
if opts.Skip != nil {
|
||||
result["skip"] = opts.Skip
|
||||
|
|
@ -25,6 +25,9 @@ func (opts *Options) ToMap() map[string]interface{} {
|
|||
if opts.DisableGlobalPrompts {
|
||||
result["disable_global_prompts"] = opts.DisableGlobalPrompts
|
||||
}
|
||||
if opts.Metadata != nil {
|
||||
result["metadata"] = opts.Metadata
|
||||
}
|
||||
|
||||
// Note: Runtime fields (Context, Writer) are not serialized (json:"-")
|
||||
// They should not be included in the map
|
||||
|
|
@ -47,8 +50,9 @@ func OptionsFromMap(m map[string]interface{}) *Options {
|
|||
if mode, ok := m["mode"].(string); ok {
|
||||
opts.Mode = mode
|
||||
}
|
||||
if search, ok := m["search"].(bool); ok {
|
||||
opts.Search = &search
|
||||
// Search supports: bool | SearchIntent | map[string]any | nil
|
||||
if search := m["search"]; search != nil {
|
||||
opts.Search = search
|
||||
}
|
||||
if skipMap, ok := m["skip"].(map[string]interface{}); ok {
|
||||
skip := &Skip{}
|
||||
|
|
@ -66,6 +70,9 @@ func OptionsFromMap(m map[string]interface{}) *Options {
|
|||
if disableGlobalPrompts, ok := m["disable_global_prompts"].(bool); ok {
|
||||
opts.DisableGlobalPrompts = disableGlobalPrompts
|
||||
}
|
||||
if metadata, ok := m["metadata"].(map[string]interface{}); ok {
|
||||
opts.Metadata = metadata
|
||||
}
|
||||
|
||||
// Note: Context and Writer are runtime fields, not restored from map
|
||||
// They should be set by the caller if needed
|
||||
|
|
|
|||
|
|
@ -266,6 +266,15 @@ type Context struct {
|
|||
Metadata map[string]interface{} `json:"metadata,omitempty"` // The metadata of the request, it will be used to pass data to the page
|
||||
}
|
||||
|
||||
// SearchIntent represents the result of search intent detection
|
||||
// Used by Create hook to specify fine-grained search behavior
|
||||
type SearchIntent struct {
|
||||
NeedSearch bool `json:"need_search"` // Whether search is needed
|
||||
SearchTypes []string `json:"search_types,omitempty"` // Types of search to perform: "web", "kb", "db"
|
||||
Confidence float64 `json:"confidence,omitempty"` // Confidence level (0-1)
|
||||
Reason string `json:"reason,omitempty"` // Reason for the decision
|
||||
}
|
||||
|
||||
// Options represents the options for the context
|
||||
type Options struct {
|
||||
|
||||
|
|
@ -284,11 +293,17 @@ type Options struct {
|
|||
// Disable global prompts, default is false
|
||||
DisableGlobalPrompts bool `json:"disable_global_prompts,omitempty"` // Temporarily disable global prompts for this request
|
||||
|
||||
// Search mode, default is true
|
||||
Search *bool `json:"search,omitempty"` // Search mode, default is true
|
||||
// Search controls search behavior, supports multiple types:
|
||||
// - bool: true = enable all search types, false = disable all search
|
||||
// - SearchIntent: fine-grained control with specific types, confidence, etc.
|
||||
// - nil: use default behavior (determined by __yao.needsearch agent)
|
||||
Search any `json:"search,omitempty"` // Search mode: bool | SearchIntent | nil
|
||||
|
||||
// Agent mode, use to select the mode of the request, default is "chat"
|
||||
Mode string `json:"mode,omitempty"` // Agent mode, use to select the mode of the request, default is "chat"
|
||||
|
||||
// Metadata for passing custom data to hooks (e.g., scenario selection)
|
||||
Metadata map[string]any `json:"metadata,omitempty"` // Custom metadata passed to Create/Next hooks
|
||||
}
|
||||
|
||||
// Stack represents the call stack node for tracing agent-to-agent calls
|
||||
|
|
@ -369,6 +384,12 @@ type HookCreateResponse struct {
|
|||
|
||||
// ForceUses controls whether to force using Uses tools even when model has native capabilities
|
||||
ForceUses *bool `json:"force_uses,omitempty"` // Force using Uses tools regardless of model capabilities
|
||||
|
||||
// Search controls search behavior, supports multiple types:
|
||||
// - bool: true = enable all search types, false = disable all search
|
||||
// - SearchIntent: fine-grained control with specific types, confidence, etc.
|
||||
// - nil: use default behavior (determined by __yao.needsearch agent)
|
||||
Search any `json:"search,omitempty"` // Search mode: bool | SearchIntent | nil
|
||||
}
|
||||
|
||||
// NextHookPayload payload for the next hook
|
||||
|
|
|
|||
|
|
@ -112,6 +112,10 @@ func init() {
|
|||
"search.intent.need_search": "Searching for references...",
|
||||
"search.intent.no_search": "No references needed",
|
||||
|
||||
// Keyword Extraction: assistant/search.go - Keyword extraction messages
|
||||
"search.keyword.loading": "Analyzing conversation...",
|
||||
"search.keyword.done": "Analysis complete",
|
||||
|
||||
// Search: assistant/search.go - Trace labels
|
||||
"search.trace.label": "Search",
|
||||
"search.trace.description": "Search the web and knowledge base for relevant information",
|
||||
|
|
@ -201,6 +205,10 @@ func init() {
|
|||
"search.intent.need_search": "正在查询相关资料...",
|
||||
"search.intent.no_search": "无需查询资料",
|
||||
|
||||
// Keyword Extraction: assistant/search.go - Keyword extraction messages
|
||||
"search.keyword.loading": "正在分析对话内容...",
|
||||
"search.keyword.done": "分析完成",
|
||||
|
||||
// Search: assistant/search.go - Trace labels
|
||||
"search.trace.label": "搜索",
|
||||
"search.trace.description": "搜索网络和知识库获取相关信息",
|
||||
|
|
@ -318,6 +326,10 @@ func init() {
|
|||
"search.intent.need_search": "正在查询相关资料...",
|
||||
"search.intent.no_search": "无需查询资料",
|
||||
|
||||
// Keyword Extraction: assistant/search.go - Keyword extraction messages
|
||||
"search.keyword.loading": "正在分析对话内容...",
|
||||
"search.keyword.done": "分析完成",
|
||||
|
||||
// Search: assistant/search.go - Trace labels
|
||||
"search.trace.label": "搜索",
|
||||
"search.trace.description": "搜索网络和知识库获取相关信息",
|
||||
|
|
|
|||
|
|
@ -637,11 +637,25 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
|
|||
errorDetected := false
|
||||
|
||||
// Wrap streamHandler to detect JSON error responses
|
||||
// Note: API error responses are raw JSON without "data: " prefix
|
||||
// Normal SSE data always starts with "data: " prefix
|
||||
wrappedHandler := func(data []byte) int {
|
||||
dataStr := string(data)
|
||||
trimmed := strings.TrimSpace(dataStr)
|
||||
|
||||
// Detect if this looks like a JSON error response (starts with "{" or contains "error")
|
||||
if strings.Contains(dataStr, `"error"`) || (strings.TrimSpace(dataStr) == "{" && !errorDetected) {
|
||||
// Skip empty lines
|
||||
if trimmed == "" {
|
||||
return http.HandlerReturnOk
|
||||
}
|
||||
|
||||
// Normal SSE data starts with "data: " - pass to streamHandler
|
||||
if strings.HasPrefix(dataStr, "data: ") {
|
||||
return streamHandler(data)
|
||||
}
|
||||
|
||||
// Detect if this looks like a JSON error response (raw JSON without "data: " prefix)
|
||||
// API errors are returned as raw JSON: {"error": {...}}
|
||||
if strings.HasPrefix(trimmed, "{") && strings.Contains(dataStr, `"error"`) {
|
||||
errorDetected = true
|
||||
}
|
||||
|
||||
|
|
@ -652,7 +666,7 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
|
|||
return http.HandlerReturnOk
|
||||
}
|
||||
|
||||
// Otherwise, use normal handler
|
||||
// Unknown format, pass to streamHandler (it will skip non-SSE data)
|
||||
return streamHandler(data)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -209,10 +209,14 @@ func initAssistant() error {
|
|||
// Set global Uses configuration
|
||||
if agentDSL.Uses != nil {
|
||||
globalUses := &context.Uses{
|
||||
Vision: agentDSL.Uses.Vision,
|
||||
Audio: agentDSL.Uses.Audio,
|
||||
Search: agentDSL.Uses.Search,
|
||||
Fetch: agentDSL.Uses.Fetch,
|
||||
Vision: agentDSL.Uses.Vision,
|
||||
Audio: agentDSL.Uses.Audio,
|
||||
Search: agentDSL.Uses.Search,
|
||||
Fetch: agentDSL.Uses.Fetch,
|
||||
Web: agentDSL.Uses.Web,
|
||||
Keyword: agentDSL.Uses.Keyword,
|
||||
QueryDSL: agentDSL.Uses.QueryDSL,
|
||||
Rerank: agentDSL.Uses.Rerank,
|
||||
}
|
||||
assistant.SetGlobalUses(globalUses)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,8 +1,15 @@
|
|||
package db
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/gou/model"
|
||||
"github.com/yaoapp/gou/query"
|
||||
"github.com/yaoapp/gou/query/gou"
|
||||
agentContext "github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/search/nlp/querydsl"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
)
|
||||
|
||||
|
|
@ -23,8 +30,13 @@ func (h *Handler) Type() types.SearchType {
|
|||
}
|
||||
|
||||
// Search converts NL to QueryDSL and executes
|
||||
// TODO: Implement actual QueryDSL generation and model query logic
|
||||
// Note: This method doesn't have context, use SearchWithContext for full functionality
|
||||
func (h *Handler) Search(req *types.Request) (*types.Result, error) {
|
||||
return h.SearchWithContext(nil, req)
|
||||
}
|
||||
|
||||
// SearchWithContext executes DB search with context (required for QueryDSL generation)
|
||||
func (h *Handler) SearchWithContext(ctx *agentContext.Context, req *types.Request) (*types.Result, error) {
|
||||
start := time.Now()
|
||||
|
||||
// Validate request
|
||||
|
|
@ -41,13 +53,13 @@ func (h *Handler) Search(req *types.Request) (*types.Result, error) {
|
|||
}
|
||||
|
||||
// Get models from request or config
|
||||
models := req.Models
|
||||
if len(models) == 0 && h.config != nil {
|
||||
models = h.config.Models
|
||||
modelIDs := req.Models
|
||||
if len(modelIDs) == 0 && h.config != nil {
|
||||
modelIDs = h.config.Models
|
||||
}
|
||||
|
||||
// If no models specified, return empty result
|
||||
if len(models) == 0 {
|
||||
if len(modelIDs) == 0 {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: req.Query,
|
||||
|
|
@ -55,6 +67,7 @@ func (h *Handler) Search(req *types.Request) (*types.Result, error) {
|
|||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(start).Milliseconds(),
|
||||
Error: "no models specified",
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
|
@ -67,27 +80,300 @@ func (h *Handler) Search(req *types.Request) (*types.Result, error) {
|
|||
maxResults = 20 // default
|
||||
}
|
||||
|
||||
// TODO: Implement actual DB search
|
||||
// 1. Get model schemas for specified models
|
||||
// 2. Generate QueryDSL from natural language query using uses.querydsl mode:
|
||||
// - "builtin": template-based generation
|
||||
// - "<assistant-id>": delegate to LLM assistant
|
||||
// - "mcp:<server>.<tool>": call external MCP tool
|
||||
// 3. Execute QueryDSL on each model
|
||||
// 4. Format results and return
|
||||
// Context is required for QueryDSL generation
|
||||
if ctx == nil {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(start).Milliseconds(),
|
||||
Error: "context is required for DB search",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// For now, return empty result (skeleton)
|
||||
result := &types.Result{
|
||||
// 1. Load all models and build combined schema
|
||||
models := make(map[string]*model.Model)
|
||||
schemas := make([]map[string]interface{}, 0, len(modelIDs))
|
||||
|
||||
for _, modelID := range modelIDs {
|
||||
mod, err := model.Get(modelID)
|
||||
if err != nil {
|
||||
continue // Skip non-existent models
|
||||
}
|
||||
models[modelID] = mod
|
||||
schemas = append(schemas, h.buildModelSchema(mod))
|
||||
}
|
||||
|
||||
if len(schemas) == 0 {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(start).Milliseconds(),
|
||||
Error: "no valid models found",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 2. Generate QueryDSL with all schemas
|
||||
generator := querydsl.NewGenerator(h.usesQueryDSL, nil)
|
||||
input := &querydsl.Input{
|
||||
Query: req.Query,
|
||||
ModelIDs: modelIDs,
|
||||
Scenario: req.Scenario, // Pass scenario: filter, aggregation, join, complex
|
||||
Limit: maxResults,
|
||||
}
|
||||
|
||||
// Build schema input: single schema or array of schemas
|
||||
var schemaInput interface{}
|
||||
if len(schemas) == 1 {
|
||||
schemaInput = schemas[0]
|
||||
} else {
|
||||
schemaInput = schemas
|
||||
}
|
||||
|
||||
input.ExtraParams = map[string]interface{}{
|
||||
"schema": schemaInput,
|
||||
}
|
||||
|
||||
result, err := generator.Generate(ctx, input)
|
||||
if err != nil {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(start).Milliseconds(),
|
||||
Error: fmt.Sprintf("QueryDSL generation failed: %v", err),
|
||||
}, nil
|
||||
}
|
||||
|
||||
if result == nil || result.DSL == nil {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(start).Milliseconds(),
|
||||
Error: "no QueryDSL generated",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 3. Merge preset conditions into generated DSL
|
||||
h.mergeDSLConditions(result.DSL, req)
|
||||
|
||||
// 4. Execute QueryDSL using gou query engine
|
||||
records, err := h.executeDSL(result.DSL)
|
||||
if err != nil {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(start).Milliseconds(),
|
||||
Error: fmt.Sprintf("query execution failed: %v", err),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 5. Determine the primary model for result formatting
|
||||
// Use the "from" table from DSL, or first model
|
||||
primaryModelID := modelIDs[0]
|
||||
if result.DSL.From != nil && result.DSL.From.Name != "" {
|
||||
// Find model by table name
|
||||
for id, mod := range models {
|
||||
if mod.MetaData.Table.Name == result.DSL.From.Name {
|
||||
primaryModelID = id
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
primaryModel := models[primaryModelID]
|
||||
if primaryModel == nil {
|
||||
primaryModel, _ = model.Get(primaryModelID) // May be nil, that's ok
|
||||
}
|
||||
|
||||
// 6. Convert records to ResultItems
|
||||
items := h.convertToResultItems(records, primaryModelID, primaryModel, req.Source)
|
||||
|
||||
// Apply limit
|
||||
if len(items) > maxResults {
|
||||
items = items[:maxResults]
|
||||
}
|
||||
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Items: items,
|
||||
Total: len(items),
|
||||
Duration: time.Since(start).Milliseconds(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// mergeDSLConditions merges preset conditions from request into generated DSL
|
||||
func (h *Handler) mergeDSLConditions(dsl *gou.QueryDSL, req *types.Request) {
|
||||
if dsl == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Store maxResults for later use
|
||||
_ = maxResults
|
||||
// Merge preset Wheres (prepend to ensure they take priority)
|
||||
if len(req.Wheres) > 0 {
|
||||
dsl.Wheres = append(req.Wheres, dsl.Wheres...)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
// Merge preset Orders (prepend to ensure they take priority)
|
||||
if len(req.Orders) > 0 {
|
||||
dsl.Orders = append(req.Orders, dsl.Orders...)
|
||||
}
|
||||
|
||||
// Merge preset Select fields
|
||||
if len(req.Select) > 0 {
|
||||
// Convert string fields to Expression
|
||||
selectExprs := make([]gou.Expression, 0, len(req.Select))
|
||||
for _, field := range req.Select {
|
||||
selectExprs = append(selectExprs, gou.Expression{Field: field})
|
||||
}
|
||||
// If DSL has no select, use preset; otherwise merge
|
||||
if len(dsl.Select) == 0 {
|
||||
dsl.Select = selectExprs
|
||||
} else {
|
||||
// Prepend preset fields
|
||||
dsl.Select = append(selectExprs, dsl.Select...)
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure limit is set
|
||||
if dsl.Limit == 0 && req.Limit > 0 {
|
||||
dsl.Limit = req.Limit
|
||||
}
|
||||
}
|
||||
|
||||
// buildModelSchema builds a simplified schema for QueryDSL generator
|
||||
func (h *Handler) buildModelSchema(mod *model.Model) map[string]interface{} {
|
||||
columns := make([]map[string]interface{}, 0, len(mod.Columns))
|
||||
for _, col := range mod.Columns {
|
||||
colInfo := map[string]interface{}{
|
||||
"name": col.Name,
|
||||
"type": col.Type,
|
||||
}
|
||||
if col.Label != "" {
|
||||
colInfo["label"] = col.Label
|
||||
}
|
||||
if col.Description != "" {
|
||||
colInfo["description"] = col.Description
|
||||
}
|
||||
columns = append(columns, colInfo)
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"name": mod.MetaData.Table.Name,
|
||||
"columns": columns,
|
||||
}
|
||||
}
|
||||
|
||||
// executeDSL executes the QueryDSL and returns records
|
||||
func (h *Handler) executeDSL(dsl interface{}) ([]map[string]interface{}, error) {
|
||||
// Get the default query engine
|
||||
engine, err := query.Select("default")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query engine not found: %w", err)
|
||||
}
|
||||
|
||||
// Marshal DSL to JSON
|
||||
dslJSON, err := json.Marshal(dsl)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal DSL: %w", err)
|
||||
}
|
||||
|
||||
// Load and execute the query
|
||||
q, err := engine.Load(json.RawMessage(dslJSON))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to load DSL: %w", err)
|
||||
}
|
||||
|
||||
// Execute query
|
||||
rawRecords := q.Get(nil)
|
||||
|
||||
// Convert to map[string]interface{}
|
||||
records := make([]map[string]interface{}, 0, len(rawRecords))
|
||||
for _, rec := range rawRecords {
|
||||
records = append(records, map[string]interface{}(rec))
|
||||
}
|
||||
|
||||
return records, nil
|
||||
}
|
||||
|
||||
// convertToResultItems converts query results to ResultItems
|
||||
func (h *Handler) convertToResultItems(records []map[string]interface{}, modelID string, mod *model.Model, source types.SourceType) []*types.ResultItem {
|
||||
items := make([]*types.ResultItem, 0, len(records))
|
||||
|
||||
primaryKey := "id"
|
||||
if mod != nil && mod.PrimaryKey != "" {
|
||||
primaryKey = mod.PrimaryKey
|
||||
}
|
||||
|
||||
for _, rec := range records {
|
||||
item := &types.ResultItem{
|
||||
Type: types.SearchTypeDB,
|
||||
Source: source,
|
||||
Model: modelID,
|
||||
Data: rec,
|
||||
}
|
||||
|
||||
// Try to extract title from common fields
|
||||
item.Title = h.extractTitle(rec, mod)
|
||||
|
||||
// Try to extract content/description
|
||||
item.Content = h.extractContent(rec, mod)
|
||||
|
||||
// Try to extract record ID
|
||||
if id, ok := rec[primaryKey]; ok {
|
||||
item.RecordID = id
|
||||
}
|
||||
|
||||
items = append(items, item)
|
||||
}
|
||||
|
||||
return items
|
||||
}
|
||||
|
||||
// extractTitle tries to extract a title from the record
|
||||
func (h *Handler) extractTitle(rec map[string]interface{}, mod *model.Model) string {
|
||||
// Common title fields
|
||||
titleFields := []string{"title", "name", "subject", "label"}
|
||||
for _, field := range titleFields {
|
||||
if val, ok := rec[field]; ok {
|
||||
if str, ok := val.(string); ok && str != "" {
|
||||
return str
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractContent tries to extract content from the record
|
||||
func (h *Handler) extractContent(rec map[string]interface{}, mod *model.Model) string {
|
||||
// Common content fields
|
||||
contentFields := []string{"content", "description", "summary", "text", "body"}
|
||||
for _, field := range contentFields {
|
||||
if val, ok := rec[field]; ok {
|
||||
if str, ok := val.(string); ok && str != "" {
|
||||
return str
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: serialize first few fields as content
|
||||
content, _ := json.Marshal(rec)
|
||||
if len(content) > 500 {
|
||||
content = content[:500]
|
||||
}
|
||||
return string(content)
|
||||
}
|
||||
|
|
|
|||
221
agent/search/handlers/db/handler_integration_test.go
Normal file
221
agent/search/handlers/db/handler_integration_test.go
Normal file
|
|
@ -0,0 +1,221 @@
|
|||
package db_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/yaoapp/gou/model"
|
||||
"github.com/yaoapp/gou/query/gou"
|
||||
"github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/search/handlers/db"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ============================================================================
|
||||
// Integration Tests - Requires database and models
|
||||
// ============================================================================
|
||||
|
||||
func TestHandler_Search_Integration(t *testing.T) {
|
||||
// Skip if running short tests
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
// Initialize test environment (loads models, database, query engine, etc.)
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Create test context
|
||||
ctx := newTestContext(t)
|
||||
|
||||
// Verify __yao.role model is loaded
|
||||
mod := model.Select("__yao.role")
|
||||
require.NotNil(t, mod, "__yao.role model should be loaded")
|
||||
|
||||
t.Run("search_role_model_with_results", func(t *testing.T) {
|
||||
// First, ensure there's at least one role in the database
|
||||
ensureTestRole(t, mod)
|
||||
|
||||
// Create handler with builtin QueryDSL generator
|
||||
h := db.NewHandler("builtin", &types.DBConfig{
|
||||
Models: []string{"__yao.role"},
|
||||
MaxResults: 10,
|
||||
})
|
||||
|
||||
req := &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "查询所有角色",
|
||||
Source: types.SourceAuto,
|
||||
Models: []string{"__yao.role"},
|
||||
Scenario: types.ScenarioFilter,
|
||||
Limit: 10,
|
||||
}
|
||||
|
||||
result, err := h.SearchWithContext(ctx, req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// Verify result structure
|
||||
assert.Equal(t, types.SearchTypeDB, result.Type)
|
||||
assert.Equal(t, "查询所有角色", result.Query)
|
||||
assert.Equal(t, types.SourceAuto, result.Source)
|
||||
assert.GreaterOrEqual(t, result.Duration, int64(0))
|
||||
|
||||
// Should have results
|
||||
if result.Error != "" {
|
||||
t.Logf("Search error: %s", result.Error)
|
||||
}
|
||||
assert.Empty(t, result.Error, "Search should not return error")
|
||||
assert.Greater(t, len(result.Items), 0, "Should have at least one result")
|
||||
assert.Equal(t, len(result.Items), result.Total)
|
||||
|
||||
// Verify result items
|
||||
for _, item := range result.Items {
|
||||
assert.Equal(t, types.SearchTypeDB, item.Type)
|
||||
assert.Equal(t, types.SourceAuto, item.Source)
|
||||
assert.Equal(t, "__yao.role", item.Model)
|
||||
assert.NotNil(t, item.Data, "Data should not be nil")
|
||||
assert.NotNil(t, item.RecordID, "RecordID should not be nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("search_with_filter_scenario", func(t *testing.T) {
|
||||
h := db.NewHandler("builtin", nil)
|
||||
|
||||
req := &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "查询系统角色",
|
||||
Source: types.SourceHook,
|
||||
Models: []string{"__yao.role"},
|
||||
Scenario: types.ScenarioFilter,
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := h.SearchWithContext(ctx, req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
assert.Equal(t, types.SearchTypeDB, result.Type)
|
||||
assert.Equal(t, types.SourceHook, result.Source)
|
||||
assert.LessOrEqual(t, len(result.Items), 5, "Should respect limit")
|
||||
})
|
||||
|
||||
t.Run("search_with_preset_wheres", func(t *testing.T) {
|
||||
h := db.NewHandler("builtin", nil)
|
||||
|
||||
req := &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "查询角色",
|
||||
Source: types.SourceAuto,
|
||||
Models: []string{"__yao.role"},
|
||||
Wheres: []gou.Where{
|
||||
{Condition: gou.Condition{Field: &gou.Expression{Field: "is_active"}, Value: true, OP: "="}},
|
||||
},
|
||||
Limit: 10,
|
||||
}
|
||||
|
||||
result, err := h.SearchWithContext(ctx, req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// All results should have is_active = true (due to preset where)
|
||||
for _, item := range result.Items {
|
||||
if data, ok := item.Data["is_active"]; ok {
|
||||
// is_active could be bool or int depending on driver
|
||||
switch v := data.(type) {
|
||||
case bool:
|
||||
assert.True(t, v)
|
||||
case int64:
|
||||
assert.Equal(t, int64(1), v)
|
||||
case float64:
|
||||
assert.Equal(t, float64(1), v)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("search_nonexistent_model_graceful", func(t *testing.T) {
|
||||
h := db.NewHandler("builtin", nil)
|
||||
|
||||
req := &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "查询文章",
|
||||
Source: types.SourceAuto,
|
||||
Models: []string{"nonexistent_model", "article", "fake_model"},
|
||||
Limit: 10,
|
||||
}
|
||||
|
||||
// Should NOT panic, should return gracefully with error
|
||||
result, err := h.SearchWithContext(ctx, req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// Should have error message about no valid models
|
||||
assert.Equal(t, types.SearchTypeDB, result.Type)
|
||||
assert.Equal(t, "no valid models found", result.Error)
|
||||
assert.Empty(t, result.Items)
|
||||
})
|
||||
|
||||
t.Run("search_mixed_models_partial_exist", func(t *testing.T) {
|
||||
h := db.NewHandler("builtin", nil)
|
||||
|
||||
req := &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "查询角色",
|
||||
Source: types.SourceAuto,
|
||||
Models: []string{"nonexistent_model", "__yao.role", "fake_model"}, // Only __yao.role exists
|
||||
Limit: 10,
|
||||
}
|
||||
|
||||
// Should NOT panic, should work with the existing model
|
||||
result, err := h.SearchWithContext(ctx, req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// Should succeed with partial models
|
||||
assert.Equal(t, types.SearchTypeDB, result.Type)
|
||||
if result.Error == "" {
|
||||
// If no error, should have results from __yao.role
|
||||
assert.GreaterOrEqual(t, len(result.Items), 0)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 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-db-search"
|
||||
ctx := context.New(t.Context(), authorized, chatID)
|
||||
return ctx
|
||||
}
|
||||
|
||||
// ensureTestRole ensures there's at least one role in the database for testing
|
||||
func ensureTestRole(t *testing.T, mod *model.Model) {
|
||||
t.Helper()
|
||||
|
||||
// Try to find existing roles
|
||||
rows, err := mod.Get(model.QueryParam{Limit: 1})
|
||||
if err == nil && len(rows) > 0 {
|
||||
return // Already have roles
|
||||
}
|
||||
|
||||
// Create a test role
|
||||
_, err = mod.Create(map[string]interface{}{
|
||||
"role_id": "test_role",
|
||||
"name": "Test Role",
|
||||
"description": "A test role for unit testing",
|
||||
"is_active": true,
|
||||
"is_system": false,
|
||||
"level": 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Logf("Note: Could not create test role: %v", err)
|
||||
}
|
||||
}
|
||||
|
|
@ -4,6 +4,8 @@ import (
|
|||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/gou/model"
|
||||
"github.com/yaoapp/gou/query/gou"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
)
|
||||
|
||||
|
|
@ -38,14 +40,13 @@ func TestHandler_Type(t *testing.T) {
|
|||
assert.Equal(t, types.SearchTypeDB, h.Type())
|
||||
}
|
||||
|
||||
func TestHandler_Search(t *testing.T) {
|
||||
func TestHandler_Search_Validation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
usesQueryDSL string
|
||||
config *types.DBConfig
|
||||
req *types.Request
|
||||
expectError string
|
||||
expectItems int
|
||||
}{
|
||||
{
|
||||
name: "empty query",
|
||||
|
|
@ -56,7 +57,6 @@ func TestHandler_Search(t *testing.T) {
|
|||
Query: "",
|
||||
},
|
||||
expectError: "query is required",
|
||||
expectItems: 0,
|
||||
},
|
||||
{
|
||||
name: "no models in request or config",
|
||||
|
|
@ -66,92 +66,20 @@ func TestHandler_Search(t *testing.T) {
|
|||
Type: types.SearchTypeDB,
|
||||
Query: "find products under $100",
|
||||
},
|
||||
expectError: "",
|
||||
expectItems: 0,
|
||||
expectError: "no models specified",
|
||||
},
|
||||
{
|
||||
name: "models from config",
|
||||
name: "context required for DB search",
|
||||
usesQueryDSL: "builtin",
|
||||
config: &types.DBConfig{
|
||||
Models: []string{"product"},
|
||||
MaxResults: 20,
|
||||
Models: []string{"product"},
|
||||
},
|
||||
req: &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "find products under $100",
|
||||
},
|
||||
expectError: "",
|
||||
expectItems: 0, // skeleton returns empty
|
||||
},
|
||||
{
|
||||
name: "models from request",
|
||||
usesQueryDSL: "builtin",
|
||||
config: nil,
|
||||
req: &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "find products under $100",
|
||||
Models: []string{"product", "order"},
|
||||
},
|
||||
expectError: "",
|
||||
expectItems: 0, // skeleton returns empty
|
||||
},
|
||||
{
|
||||
name: "with limit",
|
||||
usesQueryDSL: "builtin",
|
||||
config: &types.DBConfig{
|
||||
Models: []string{"product"},
|
||||
},
|
||||
req: &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "find products",
|
||||
Models: []string{"product"},
|
||||
Limit: 5,
|
||||
},
|
||||
expectError: "",
|
||||
expectItems: 0, // skeleton returns empty
|
||||
},
|
||||
{
|
||||
name: "with wheres",
|
||||
usesQueryDSL: "builtin",
|
||||
config: &types.DBConfig{
|
||||
Models: []string{"product"},
|
||||
},
|
||||
req: &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "find products",
|
||||
Models: []string{"product"},
|
||||
// Wheres would be set here in real usage
|
||||
},
|
||||
expectError: "",
|
||||
expectItems: 0, // skeleton returns empty
|
||||
},
|
||||
{
|
||||
name: "agent mode",
|
||||
usesQueryDSL: "workers.nlp.querydsl",
|
||||
config: &types.DBConfig{
|
||||
Models: []string{"product"},
|
||||
},
|
||||
req: &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "find products",
|
||||
Models: []string{"product"},
|
||||
},
|
||||
expectError: "",
|
||||
expectItems: 0, // skeleton returns empty
|
||||
},
|
||||
{
|
||||
name: "mcp mode",
|
||||
usesQueryDSL: "mcp:nlp.generate_querydsl",
|
||||
config: &types.DBConfig{
|
||||
Models: []string{"product"},
|
||||
},
|
||||
req: &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "find products",
|
||||
Models: []string{"product"},
|
||||
},
|
||||
expectError: "",
|
||||
expectItems: 0, // skeleton returns empty
|
||||
expectError: "context is required for DB search",
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -163,16 +91,8 @@ func TestHandler_Search(t *testing.T) {
|
|||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, types.SearchTypeDB, result.Type)
|
||||
assert.Equal(t, tt.req.Query, result.Query)
|
||||
assert.Equal(t, tt.expectItems, len(result.Items))
|
||||
|
||||
if tt.expectError != "" {
|
||||
assert.Equal(t, tt.expectError, result.Error)
|
||||
} else {
|
||||
assert.Empty(t, result.Error)
|
||||
}
|
||||
|
||||
// Duration should be set
|
||||
assert.Equal(t, tt.expectError, result.Error)
|
||||
assert.Equal(t, 0, len(result.Items))
|
||||
assert.GreaterOrEqual(t, result.Duration, int64(0))
|
||||
})
|
||||
}
|
||||
|
|
@ -195,21 +115,374 @@ func TestHandler_Search_SourcePreserved(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestHandler_Search_MaxResultsFromConfig(t *testing.T) {
|
||||
cfg := &types.DBConfig{
|
||||
Models: []string{"product"},
|
||||
MaxResults: 50,
|
||||
}
|
||||
h := NewHandler("builtin", cfg)
|
||||
func TestHandler_BuildModelSchema(t *testing.T) {
|
||||
h := NewHandler("builtin", nil)
|
||||
|
||||
req := &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "test",
|
||||
Models: []string{"product"},
|
||||
// No limit in request, should use config's MaxResults
|
||||
// Create a mock model for testing
|
||||
mod := &model.Model{
|
||||
MetaData: model.MetaData{
|
||||
Table: model.Table{
|
||||
Name: "test_products",
|
||||
},
|
||||
},
|
||||
Columns: map[string]*model.Column{
|
||||
"id": {
|
||||
Name: "id",
|
||||
Type: "ID",
|
||||
Label: "ID",
|
||||
},
|
||||
"name": {
|
||||
Name: "name",
|
||||
Type: "string",
|
||||
Label: "Name",
|
||||
Description: "Product name",
|
||||
},
|
||||
"price": {
|
||||
Name: "price",
|
||||
Type: "decimal",
|
||||
Label: "Price",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
schema := h.buildModelSchema(mod)
|
||||
|
||||
assert.NotNil(t, schema)
|
||||
assert.Equal(t, "test_products", schema["name"])
|
||||
|
||||
columns, ok := schema["columns"].([]map[string]interface{})
|
||||
assert.True(t, ok)
|
||||
assert.Len(t, columns, 3)
|
||||
|
||||
// Verify columns have required fields
|
||||
for _, col := range columns {
|
||||
assert.NotEmpty(t, col["name"])
|
||||
assert.NotEmpty(t, col["type"])
|
||||
}
|
||||
result, err := h.Search(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
// Skeleton doesn't actually use maxResults yet, but the test ensures the handler runs
|
||||
}
|
||||
|
||||
func TestHandler_BuildModelSchema_MultipleModels(t *testing.T) {
|
||||
h := NewHandler("builtin", nil)
|
||||
|
||||
// Create mock models for testing joins
|
||||
productMod := &model.Model{
|
||||
MetaData: model.MetaData{
|
||||
Table: model.Table{Name: "products"},
|
||||
},
|
||||
Columns: map[string]*model.Column{
|
||||
"id": {Name: "id", Type: "ID"},
|
||||
"name": {Name: "name", Type: "string"},
|
||||
"category_id": {Name: "category_id", Type: "integer"},
|
||||
},
|
||||
}
|
||||
|
||||
categoryMod := &model.Model{
|
||||
MetaData: model.MetaData{
|
||||
Table: model.Table{Name: "categories"},
|
||||
},
|
||||
Columns: map[string]*model.Column{
|
||||
"id": {Name: "id", Type: "ID"},
|
||||
"name": {Name: "name", Type: "string"},
|
||||
},
|
||||
}
|
||||
|
||||
productSchema := h.buildModelSchema(productMod)
|
||||
categorySchema := h.buildModelSchema(categoryMod)
|
||||
|
||||
assert.Equal(t, "products", productSchema["name"])
|
||||
assert.Equal(t, "categories", categorySchema["name"])
|
||||
|
||||
// Verify both schemas can be combined into an array
|
||||
schemas := []map[string]interface{}{productSchema, categorySchema}
|
||||
assert.Len(t, schemas, 2)
|
||||
}
|
||||
|
||||
func TestHandler_ExtractTitle(t *testing.T) {
|
||||
h := NewHandler("builtin", nil)
|
||||
mod := &model.Model{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
record map[string]interface{}
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "title field",
|
||||
record: map[string]interface{}{"title": "Test Title", "id": 1},
|
||||
expected: "Test Title",
|
||||
},
|
||||
{
|
||||
name: "name field",
|
||||
record: map[string]interface{}{"name": "Test Name", "id": 1},
|
||||
expected: "Test Name",
|
||||
},
|
||||
{
|
||||
name: "subject field",
|
||||
record: map[string]interface{}{"subject": "Test Subject", "id": 1},
|
||||
expected: "Test Subject",
|
||||
},
|
||||
{
|
||||
name: "label field",
|
||||
record: map[string]interface{}{"label": "Test Label", "id": 1},
|
||||
expected: "Test Label",
|
||||
},
|
||||
{
|
||||
name: "no title field",
|
||||
record: map[string]interface{}{"id": 1, "price": 100},
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "empty title",
|
||||
record: map[string]interface{}{"title": "", "name": "Fallback"},
|
||||
expected: "Fallback",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
title := h.extractTitle(tt.record, mod)
|
||||
assert.Equal(t, tt.expected, title)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_ExtractContent(t *testing.T) {
|
||||
h := NewHandler("builtin", nil)
|
||||
mod := &model.Model{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
record map[string]interface{}
|
||||
expectEmpty bool
|
||||
}{
|
||||
{
|
||||
name: "content field",
|
||||
record: map[string]interface{}{"content": "Test Content"},
|
||||
expectEmpty: false,
|
||||
},
|
||||
{
|
||||
name: "description field",
|
||||
record: map[string]interface{}{"description": "Test Description"},
|
||||
expectEmpty: false,
|
||||
},
|
||||
{
|
||||
name: "summary field",
|
||||
record: map[string]interface{}{"summary": "Test Summary"},
|
||||
expectEmpty: false,
|
||||
},
|
||||
{
|
||||
name: "fallback to JSON",
|
||||
record: map[string]interface{}{"id": 1, "price": 100},
|
||||
expectEmpty: false, // Should return JSON representation
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
content := h.extractContent(tt.record, mod)
|
||||
if tt.expectEmpty {
|
||||
assert.Empty(t, content)
|
||||
} else {
|
||||
assert.NotEmpty(t, content)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_ConvertToResultItems(t *testing.T) {
|
||||
h := NewHandler("builtin", nil)
|
||||
mod := &model.Model{
|
||||
PrimaryKey: "id",
|
||||
}
|
||||
|
||||
records := []map[string]interface{}{
|
||||
{
|
||||
"id": 1,
|
||||
"name": "Product 1",
|
||||
"description": "Description 1",
|
||||
"price": 99.99,
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"title": "Product 2",
|
||||
"content": "Content 2",
|
||||
},
|
||||
}
|
||||
|
||||
items := h.convertToResultItems(records, "product", mod, types.SourceAuto)
|
||||
|
||||
assert.Len(t, items, 2)
|
||||
|
||||
// First item
|
||||
assert.Equal(t, types.SearchTypeDB, items[0].Type)
|
||||
assert.Equal(t, types.SourceAuto, items[0].Source)
|
||||
assert.Equal(t, "product", items[0].Model)
|
||||
assert.Equal(t, 1, items[0].RecordID)
|
||||
assert.Equal(t, "Product 1", items[0].Title)
|
||||
assert.Equal(t, "Description 1", items[0].Content)
|
||||
assert.NotNil(t, items[0].Data)
|
||||
|
||||
// Second item
|
||||
assert.Equal(t, 2, items[1].RecordID)
|
||||
assert.Equal(t, "Product 2", items[1].Title)
|
||||
assert.Equal(t, "Content 2", items[1].Content)
|
||||
}
|
||||
|
||||
func TestHandler_ConvertToResultItems_NilModel(t *testing.T) {
|
||||
h := NewHandler("builtin", nil)
|
||||
|
||||
records := []map[string]interface{}{
|
||||
{"id": 1, "name": "Test"},
|
||||
}
|
||||
|
||||
// Should use default primary key "id" when model is nil
|
||||
items := h.convertToResultItems(records, "test", nil, types.SourceHook)
|
||||
|
||||
assert.Len(t, items, 1)
|
||||
assert.Equal(t, 1, items[0].RecordID)
|
||||
assert.Equal(t, "Test", items[0].Title)
|
||||
}
|
||||
|
||||
func TestHandler_Search_ScenarioTypes(t *testing.T) {
|
||||
// Test that all scenario types are valid
|
||||
scenarios := []types.ScenarioType{
|
||||
types.ScenarioFilter,
|
||||
types.ScenarioAggregation,
|
||||
types.ScenarioJoin,
|
||||
types.ScenarioComplex,
|
||||
}
|
||||
|
||||
for _, scenario := range scenarios {
|
||||
t.Run(string(scenario), func(t *testing.T) {
|
||||
h := NewHandler("builtin", &types.DBConfig{Models: []string{"product"}})
|
||||
req := &types.Request{
|
||||
Type: types.SearchTypeDB,
|
||||
Query: "test query",
|
||||
Source: types.SourceAuto,
|
||||
Models: []string{"product"},
|
||||
Scenario: scenario,
|
||||
}
|
||||
|
||||
// Without context, should return error (but scenario should be preserved in request)
|
||||
result, err := h.Search(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
// Verify request scenario is set correctly
|
||||
assert.Equal(t, scenario, req.Scenario)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestScenarioTypeConstants(t *testing.T) {
|
||||
// Verify scenario type constants match expected values
|
||||
assert.Equal(t, types.ScenarioType("filter"), types.ScenarioFilter)
|
||||
assert.Equal(t, types.ScenarioType("aggregation"), types.ScenarioAggregation)
|
||||
assert.Equal(t, types.ScenarioType("join"), types.ScenarioJoin)
|
||||
assert.Equal(t, types.ScenarioType("complex"), types.ScenarioComplex)
|
||||
}
|
||||
|
||||
func TestHandler_MergeDSLConditions(t *testing.T) {
|
||||
h := NewHandler("builtin", nil)
|
||||
|
||||
t.Run("merge wheres", func(t *testing.T) {
|
||||
dsl := &gou.QueryDSL{
|
||||
From: &gou.Table{Name: "users"},
|
||||
Wheres: []gou.Where{
|
||||
{Condition: gou.Condition{Field: &gou.Expression{Field: "status"}, Value: "active", OP: "="}},
|
||||
},
|
||||
}
|
||||
req := &types.Request{
|
||||
Wheres: []gou.Where{
|
||||
{Condition: gou.Condition{Field: &gou.Expression{Field: "tenant_id"}, Value: 1, OP: "="}},
|
||||
},
|
||||
}
|
||||
|
||||
h.mergeDSLConditions(dsl, req)
|
||||
|
||||
// Preset wheres should be prepended
|
||||
assert.Len(t, dsl.Wheres, 2)
|
||||
assert.Equal(t, "tenant_id", dsl.Wheres[0].Field.Field)
|
||||
assert.Equal(t, "status", dsl.Wheres[1].Field.Field)
|
||||
})
|
||||
|
||||
t.Run("merge orders", func(t *testing.T) {
|
||||
dsl := &gou.QueryDSL{
|
||||
From: &gou.Table{Name: "products"},
|
||||
Orders: gou.Orders{
|
||||
{Field: &gou.Expression{Field: "name"}, Sort: "asc"},
|
||||
},
|
||||
}
|
||||
req := &types.Request{
|
||||
Orders: gou.Orders{
|
||||
{Field: &gou.Expression{Field: "created_at"}, Sort: "desc"},
|
||||
},
|
||||
}
|
||||
|
||||
h.mergeDSLConditions(dsl, req)
|
||||
|
||||
// Preset orders should be prepended
|
||||
assert.Len(t, dsl.Orders, 2)
|
||||
assert.Equal(t, "created_at", dsl.Orders[0].Field.Field)
|
||||
assert.Equal(t, "name", dsl.Orders[1].Field.Field)
|
||||
})
|
||||
|
||||
t.Run("merge select fields", func(t *testing.T) {
|
||||
dsl := &gou.QueryDSL{
|
||||
From: &gou.Table{Name: "orders"},
|
||||
Select: []gou.Expression{
|
||||
{Field: "amount"},
|
||||
},
|
||||
}
|
||||
req := &types.Request{
|
||||
Select: []string{"id", "status"},
|
||||
}
|
||||
|
||||
h.mergeDSLConditions(dsl, req)
|
||||
|
||||
// Preset select should be prepended
|
||||
assert.Len(t, dsl.Select, 3)
|
||||
assert.Equal(t, "id", dsl.Select[0].Field)
|
||||
assert.Equal(t, "status", dsl.Select[1].Field)
|
||||
assert.Equal(t, "amount", dsl.Select[2].Field)
|
||||
})
|
||||
|
||||
t.Run("set limit from request", func(t *testing.T) {
|
||||
dsl := &gou.QueryDSL{
|
||||
From: &gou.Table{Name: "users"},
|
||||
Limit: 0,
|
||||
}
|
||||
req := &types.Request{
|
||||
Limit: 50,
|
||||
}
|
||||
|
||||
h.mergeDSLConditions(dsl, req)
|
||||
|
||||
assert.Equal(t, 50, dsl.Limit)
|
||||
})
|
||||
|
||||
t.Run("preserve dsl limit if set", func(t *testing.T) {
|
||||
dsl := &gou.QueryDSL{
|
||||
From: &gou.Table{Name: "users"},
|
||||
Limit: 10,
|
||||
}
|
||||
req := &types.Request{
|
||||
Limit: 50,
|
||||
}
|
||||
|
||||
h.mergeDSLConditions(dsl, req)
|
||||
|
||||
// DSL limit should be preserved
|
||||
assert.Equal(t, 10, dsl.Limit)
|
||||
})
|
||||
|
||||
t.Run("nil dsl", func(t *testing.T) {
|
||||
req := &types.Request{
|
||||
Wheres: []gou.Where{
|
||||
{Condition: gou.Condition{Field: &gou.Expression{Field: "id"}, Value: 1}},
|
||||
},
|
||||
}
|
||||
|
||||
// Should not panic
|
||||
h.mergeDSLConditions(nil, req)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -150,46 +150,20 @@ func (p *AgentProvider) Search(ctx *agentContext.Context, req *types.Request) (*
|
|||
}, nil
|
||||
}
|
||||
|
||||
// parseAgentResponse parses the agent response into search result items
|
||||
// The agent should return a JSON structure with search results
|
||||
func (p *AgentProvider) parseAgentResponse(response interface{}, source types.SourceType) ([]*types.ResultItem, int, string) {
|
||||
if response == nil {
|
||||
// parseAgentResponse parses the agent's *context.Response into search result items
|
||||
// Now that agent.Stream() returns *context.Response directly,
|
||||
// we can access fields without type assertions.
|
||||
//
|
||||
// The agent returns search results in response.Next field
|
||||
func (p *AgentProvider) parseAgentResponse(response *agentContext.Response, source types.SourceType) ([]*types.ResultItem, int, string) {
|
||||
if response == nil || response.Next == nil {
|
||||
return nil, 0, "Agent returned nil response"
|
||||
}
|
||||
|
||||
// Try to extract data from response
|
||||
var data map[string]interface{}
|
||||
|
||||
// Handle different response types
|
||||
switch v := response.(type) {
|
||||
case map[string]interface{}:
|
||||
data = v
|
||||
case string:
|
||||
// Try to parse as JSON
|
||||
if err := json.Unmarshal([]byte(v), &data); err != nil {
|
||||
return nil, 0, fmt.Sprintf("Failed to parse agent response as JSON: %v", err)
|
||||
}
|
||||
default:
|
||||
// Try to marshal and unmarshal
|
||||
jsonBytes, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Sprintf("Failed to serialize agent response: %v", err)
|
||||
}
|
||||
if err := json.Unmarshal(jsonBytes, &data); err != nil {
|
||||
return nil, 0, fmt.Sprintf("Failed to parse agent response: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Check for "next" field (custom hook data)
|
||||
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 {
|
||||
// Try to parse as JSON
|
||||
if err := json.Unmarshal([]byte(nextStr), &data); err != nil {
|
||||
return nil, 0, fmt.Sprintf("Failed to parse next hook data: %v", err)
|
||||
}
|
||||
}
|
||||
// Extract data from Next field
|
||||
data := extractNextData(response.Next)
|
||||
if data == nil {
|
||||
return nil, 0, "Failed to extract data from agent response"
|
||||
}
|
||||
|
||||
// Extract items from data
|
||||
|
|
@ -230,3 +204,34 @@ func (p *AgentProvider) parseAgentResponse(response interface{}, source types.So
|
|||
|
||||
return items, total, ""
|
||||
}
|
||||
|
||||
// extractNextData extracts the actual data from response.Next field
|
||||
// Handles nested structures like { "data": { ... } }
|
||||
func extractNextData(next interface{}) map[string]interface{} {
|
||||
if next == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch v := next.(type) {
|
||||
case map[string]interface{}:
|
||||
// Check for "data" wrapper
|
||||
if data, ok := v["data"].(map[string]interface{}); ok {
|
||||
return data
|
||||
}
|
||||
return v
|
||||
case string:
|
||||
// Try to parse as JSON
|
||||
var data map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(v), &data); err == nil {
|
||||
return extractNextData(data)
|
||||
}
|
||||
}
|
||||
// Try to handle other types by converting to JSON and back
|
||||
if bytes, err := json.Marshal(next); err == nil {
|
||||
var data map[string]interface{}
|
||||
if err := json.Unmarshal(bytes, &data); err == nil {
|
||||
return extractNextData(data)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
package interfaces
|
||||
|
||||
import (
|
||||
"github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
)
|
||||
|
||||
|
|
@ -12,3 +13,12 @@ type Handler interface {
|
|||
// Search executes the search and returns results
|
||||
Search(req *types.Request) (*types.Result, error)
|
||||
}
|
||||
|
||||
// ContextHandler extends Handler with context support
|
||||
// Handlers that need context (e.g., DB handler for QueryDSL generation) should implement this
|
||||
type ContextHandler interface {
|
||||
Handler
|
||||
|
||||
// SearchWithContext executes the search with context and returns results
|
||||
SearchWithContext(ctx *context.Context, req *types.Request) (*types.Result, error)
|
||||
}
|
||||
|
|
|
|||
191
agent/search/jsapi_db_test.go
Normal file
191
agent/search/jsapi_db_test.go
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
package search_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/yaoapp/gou/model"
|
||||
"github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/search"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ============================================================================
|
||||
// DB Search JSAPI Integration Tests
|
||||
// ============================================================================
|
||||
|
||||
func TestJSAPI_DB_Integration(t *testing.T) {
|
||||
// Skip if running short tests
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
// Initialize test environment (loads models, database, query engine, etc.)
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Create test context
|
||||
ctx := newJSAPITestContext(t)
|
||||
|
||||
// Verify __yao.role model is loaded
|
||||
mod := model.Select("__yao.role")
|
||||
require.NotNil(t, mod, "__yao.role model should be loaded")
|
||||
|
||||
// Ensure test data exists
|
||||
ensureJSAPITestRole(t, mod)
|
||||
|
||||
t.Run("db_search_with_context", func(t *testing.T) {
|
||||
api := search.NewJSAPI(ctx, &types.Config{
|
||||
DB: &types.DBConfig{
|
||||
Models: []string{"__yao.role"},
|
||||
MaxResults: 10,
|
||||
},
|
||||
}, &search.Uses{QueryDSL: "builtin"})
|
||||
|
||||
result := api.DB("查询所有角色", map[string]interface{}{
|
||||
"models": []interface{}{"__yao.role"},
|
||||
"limit": float64(10),
|
||||
})
|
||||
|
||||
require.NotNil(t, result)
|
||||
r, ok := result.(*types.Result)
|
||||
require.True(t, ok)
|
||||
|
||||
assert.Equal(t, types.SearchTypeDB, r.Type)
|
||||
assert.Equal(t, "查询所有角色", r.Query)
|
||||
assert.Equal(t, types.SourceHook, r.Source)
|
||||
|
||||
if r.Error != "" {
|
||||
t.Logf("Search error: %s", r.Error)
|
||||
}
|
||||
assert.Empty(t, r.Error, "Should not have error")
|
||||
assert.Greater(t, len(r.Items), 0, "Should have results")
|
||||
})
|
||||
|
||||
t.Run("db_search_with_scenario", func(t *testing.T) {
|
||||
api := search.NewJSAPI(ctx, &types.Config{
|
||||
DB: &types.DBConfig{
|
||||
Models: []string{"__yao.role"},
|
||||
MaxResults: 5,
|
||||
},
|
||||
}, &search.Uses{QueryDSL: "builtin"})
|
||||
|
||||
result := api.DB("查询系统角色", map[string]interface{}{
|
||||
"models": []interface{}{"__yao.role"},
|
||||
"scenario": "filter",
|
||||
"limit": float64(5),
|
||||
})
|
||||
|
||||
require.NotNil(t, result)
|
||||
r, ok := result.(*types.Result)
|
||||
require.True(t, ok)
|
||||
|
||||
assert.Equal(t, types.SearchTypeDB, r.Type)
|
||||
assert.LessOrEqual(t, len(r.Items), 5, "Should respect limit")
|
||||
})
|
||||
|
||||
t.Run("db_search_with_select_fields", func(t *testing.T) {
|
||||
api := search.NewJSAPI(ctx, &types.Config{
|
||||
DB: &types.DBConfig{
|
||||
Models: []string{"__yao.role"},
|
||||
MaxResults: 10,
|
||||
},
|
||||
}, &search.Uses{QueryDSL: "builtin"})
|
||||
|
||||
result := api.DB("查询角色名称", map[string]interface{}{
|
||||
"models": []interface{}{"__yao.role"},
|
||||
"select": []interface{}{"id", "name", "description"},
|
||||
"limit": float64(10),
|
||||
})
|
||||
|
||||
require.NotNil(t, result)
|
||||
r, ok := result.(*types.Result)
|
||||
require.True(t, ok)
|
||||
|
||||
assert.Equal(t, types.SearchTypeDB, r.Type)
|
||||
if r.Error == "" && len(r.Items) > 0 {
|
||||
// Verify items have data
|
||||
for _, item := range r.Items {
|
||||
assert.NotNil(t, item.Data)
|
||||
assert.Equal(t, "__yao.role", item.Model)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("db_search_all_with_multiple_types", func(t *testing.T) {
|
||||
api := search.NewJSAPI(ctx, &types.Config{
|
||||
KB: &types.KBConfig{Collections: []string{"docs"}},
|
||||
DB: &types.DBConfig{
|
||||
Models: []string{"__yao.role"},
|
||||
MaxResults: 10,
|
||||
},
|
||||
}, &search.Uses{QueryDSL: "builtin"})
|
||||
|
||||
requests := []interface{}{
|
||||
map[string]interface{}{
|
||||
"type": "db",
|
||||
"query": "查询角色",
|
||||
"models": []interface{}{"__yao.role"},
|
||||
"limit": float64(5),
|
||||
},
|
||||
map[string]interface{}{
|
||||
"type": "kb",
|
||||
"query": "知识库查询",
|
||||
"collections": []interface{}{"docs"},
|
||||
"limit": float64(5),
|
||||
},
|
||||
}
|
||||
|
||||
results := api.All(requests)
|
||||
require.Len(t, results, 2)
|
||||
|
||||
// DB result
|
||||
r0, ok := results[0].(*types.Result)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, types.SearchTypeDB, r0.Type)
|
||||
|
||||
// KB result
|
||||
r1, ok := results[1].(*types.Result)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, types.SearchTypeKB, r1.Type)
|
||||
})
|
||||
}
|
||||
|
||||
// newJSAPITestContext creates a test context for JSAPI tests
|
||||
func newJSAPITestContext(t *testing.T) *context.Context {
|
||||
t.Helper()
|
||||
authorized := &oauthTypes.AuthorizedInfo{
|
||||
UserID: "test-user-jsapi",
|
||||
}
|
||||
chatID := "test-chat-jsapi-db"
|
||||
ctx := context.New(t.Context(), authorized, chatID)
|
||||
return ctx
|
||||
}
|
||||
|
||||
// ensureJSAPITestRole ensures there's at least one role in the database
|
||||
func ensureJSAPITestRole(t *testing.T, mod *model.Model) {
|
||||
t.Helper()
|
||||
|
||||
// Try to find existing roles
|
||||
rows, err := mod.Get(model.QueryParam{Limit: 1})
|
||||
if err == nil && len(rows) > 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Create a test role
|
||||
_, err = mod.Create(map[string]interface{}{
|
||||
"role_id": "jsapi_test_role",
|
||||
"name": "JSAPI Test Role",
|
||||
"description": "A test role for JSAPI unit testing",
|
||||
"is_active": true,
|
||||
"is_system": false,
|
||||
"level": 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Logf("Note: Could not create test role: %v", err)
|
||||
}
|
||||
}
|
||||
|
|
@ -23,8 +23,8 @@ func NewAgentProvider(agentID string) *AgentProvider {
|
|||
}
|
||||
|
||||
// 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) {
|
||||
// The agent receives the content and returns extracted keywords with weights
|
||||
func (p *AgentProvider) Extract(ctx *agentContext.Context, content string, opts *types.KeywordOptions) ([]types.Keyword, error) {
|
||||
if ctx == nil {
|
||||
return nil, fmt.Errorf("context is required for agent keyword extraction")
|
||||
}
|
||||
|
|
@ -64,73 +64,64 @@ func (p *AgentProvider) Extract(ctx *agentContext.Context, content string, opts
|
|||
},
|
||||
}
|
||||
|
||||
result, err := agent.Stream(ctx, messages, options)
|
||||
response, 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)
|
||||
// Parse the result from response.Next
|
||||
return p.parseResponse(response)
|
||||
}
|
||||
|
||||
// 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
|
||||
// parseResponse extracts keywords from the agent's *context.Response
|
||||
// Now that agent.Stream() returns *context.Response directly,
|
||||
// we can access fields without type assertions.
|
||||
//
|
||||
// The agent returns keywords in response.Next field as {data: {keywords: [{k, w}, ...]}}
|
||||
func (p *AgentProvider) parseResponse(response *agentContext.Response) ([]types.Keyword, error) {
|
||||
if response == nil || response.Next == nil {
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
|
||||
return p.parseNextData(response.Next)
|
||||
}
|
||||
|
||||
// parseNextData extracts keywords from Next hook data
|
||||
// Expected format: {data: {keywords: [{k: "keyword", w: 0.9}, ...]}}
|
||||
func (p *AgentProvider) parseNextData(next interface{}) ([]types.Keyword, error) {
|
||||
if next == nil {
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
|
||||
// Try to convert to map first (most common case)
|
||||
var data map[string]interface{}
|
||||
|
||||
switch v := result.(type) {
|
||||
switch v := next.(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
|
||||
// Not a JSON object, try as array of keywords
|
||||
var keywords []types.Keyword
|
||||
if err := json.Unmarshal([]byte(v), &keywords); err == nil {
|
||||
return keywords, nil
|
||||
}
|
||||
// Return as single keyword
|
||||
return []string{v}, nil
|
||||
// Return as single keyword with default weight
|
||||
return []types.Keyword{{K: v, W: 0.5}}, nil
|
||||
}
|
||||
case []string:
|
||||
case []types.Keyword:
|
||||
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
|
||||
return p.extractKeywordsFromArray(v)
|
||||
default:
|
||||
// Try to marshal and unmarshal
|
||||
jsonBytes, err := json.Marshal(result)
|
||||
jsonBytes, err := json.Marshal(next)
|
||||
if err != nil {
|
||||
return []string{}, nil
|
||||
return []types.Keyword{}, 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
|
||||
}
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -148,28 +139,48 @@ func (p *AgentProvider) parseResult(result interface{}) ([]string, error) {
|
|||
return p.extractKeywordsFromValue(d)
|
||||
}
|
||||
|
||||
return []string{}, nil
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
|
||||
// extractKeywordsFromValue extracts string array from various types
|
||||
func (p *AgentProvider) extractKeywordsFromValue(v interface{}) ([]string, error) {
|
||||
// extractKeywordsFromValue extracts Keyword array from various types
|
||||
func (p *AgentProvider) extractKeywordsFromValue(v interface{}) ([]types.Keyword, error) {
|
||||
switch kw := v.(type) {
|
||||
case []string:
|
||||
case []types.Keyword:
|
||||
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
|
||||
return p.extractKeywordsFromArray(kw)
|
||||
case string:
|
||||
var keywords []string
|
||||
var keywords []types.Keyword
|
||||
if err := json.Unmarshal([]byte(kw), &keywords); err == nil {
|
||||
return keywords, nil
|
||||
}
|
||||
return []string{kw}, nil
|
||||
return []types.Keyword{{K: kw, W: 0.5}}, nil
|
||||
}
|
||||
return []string{}, nil
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
|
||||
// extractKeywordsFromArray extracts keywords from []interface{}
|
||||
// Handles both {k, w} objects and plain strings
|
||||
func (p *AgentProvider) extractKeywordsFromArray(items []interface{}) ([]types.Keyword, error) {
|
||||
keywords := make([]types.Keyword, 0, len(items))
|
||||
for _, item := range items {
|
||||
switch v := item.(type) {
|
||||
case map[string]interface{}:
|
||||
// Handle {k: "keyword", w: 0.9} format
|
||||
k, _ := v["k"].(string)
|
||||
w, _ := v["w"].(float64)
|
||||
if k != "" {
|
||||
if w == 0 {
|
||||
w = 0.5 // Default weight
|
||||
}
|
||||
keywords = append(keywords, types.Keyword{K: k, W: w})
|
||||
}
|
||||
case string:
|
||||
// Plain string, use default weight
|
||||
if v != "" {
|
||||
keywords = append(keywords, types.Keyword{K: v, W: 0.5})
|
||||
}
|
||||
}
|
||||
}
|
||||
return keywords, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
@ -32,19 +34,20 @@ func NewExtractor(usesKeyword string, cfg *types.KeywordConfig) *Extractor {
|
|||
}
|
||||
|
||||
// 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) {
|
||||
// Returns a list of keywords with weights optimized for search queries
|
||||
func (e *Extractor) Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]types.Keyword, error) {
|
||||
// Merge options with config defaults
|
||||
mergedOpts := e.mergeOptions(opts)
|
||||
|
||||
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,29 +81,24 @@ 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) ([]types.Keyword, error) {
|
||||
if ctx == nil {
|
||||
return nil, fmt.Errorf("context is required for keyword extraction")
|
||||
}
|
||||
provider := NewAgentProvider(agentID)
|
||||
return provider.Extract(ctx, content, opts)
|
||||
}
|
||||
|
||||
// mcpExtract calls an external MCP tool
|
||||
// Format: "mcp:<server>.<tool>"
|
||||
func (e *Extractor) mcpExtract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) {
|
||||
func (e *Extractor) mcpExtract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]types.Keyword, 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)
|
||||
// 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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ func NewMCPProvider(mcpRef string) (*MCPProvider, error) {
|
|||
}
|
||||
|
||||
// Extract extracts keywords by calling the MCP tool
|
||||
func (p *MCPProvider) Extract(ctx *agentContext.Context, content string, opts *types.KeywordOptions) ([]string, error) {
|
||||
func (p *MCPProvider) Extract(ctx *agentContext.Context, content string, opts *types.KeywordOptions) ([]types.Keyword, error) {
|
||||
// Get MCP client
|
||||
client, err := mcp.Select(p.serverID)
|
||||
if err != nil {
|
||||
|
|
@ -56,9 +56,9 @@ func (p *MCPProvider) Extract(ctx *agentContext.Context, content string, opts *t
|
|||
}
|
||||
|
||||
// parseResult extracts keywords from the MCP tool response
|
||||
func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse) ([]string, error) {
|
||||
func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse) ([]types.Keyword, error) {
|
||||
if result == nil {
|
||||
return []string{}, nil
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
|
||||
// Check for errors in result
|
||||
|
|
@ -72,7 +72,7 @@ func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse) ([]strin
|
|||
|
||||
// Parse content - expect JSON data with "keywords" field
|
||||
if len(result.Content) == 0 {
|
||||
return []string{}, nil
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
|
||||
// Try to extract keywords from content
|
||||
|
|
@ -88,36 +88,50 @@ func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse) ([]strin
|
|||
}
|
||||
}
|
||||
|
||||
// Try to parse as direct array
|
||||
var keywords []string
|
||||
// Try to parse as direct array of keywords
|
||||
var keywords []types.Keyword
|
||||
if err := json.Unmarshal([]byte(content.Text), &keywords); err == nil {
|
||||
return keywords, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return []string{}, nil
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
|
||||
// extractKeywordsFromValue extracts string array from various types
|
||||
func (p *MCPProvider) extractKeywordsFromValue(v interface{}) ([]string, error) {
|
||||
// extractKeywordsFromValue extracts Keyword array from various types
|
||||
func (p *MCPProvider) extractKeywordsFromValue(v interface{}) ([]types.Keyword, error) {
|
||||
switch kw := v.(type) {
|
||||
case []string:
|
||||
case []types.Keyword:
|
||||
return kw, nil
|
||||
case []interface{}:
|
||||
keywords := make([]string, 0, len(kw))
|
||||
keywords := make([]types.Keyword, 0, len(kw))
|
||||
for _, item := range kw {
|
||||
if s, ok := item.(string); ok {
|
||||
keywords = append(keywords, s)
|
||||
switch v := item.(type) {
|
||||
case map[string]interface{}:
|
||||
// Handle {k: "keyword", w: 0.9} format
|
||||
k, _ := v["k"].(string)
|
||||
w, _ := v["w"].(float64)
|
||||
if k != "" {
|
||||
if w == 0 {
|
||||
w = 0.5 // Default weight
|
||||
}
|
||||
keywords = append(keywords, types.Keyword{K: k, W: w})
|
||||
}
|
||||
case string:
|
||||
// Plain string, use default weight
|
||||
if v != "" {
|
||||
keywords = append(keywords, types.Keyword{K: v, W: 0.5})
|
||||
}
|
||||
}
|
||||
}
|
||||
return keywords, nil
|
||||
case string:
|
||||
var keywords []string
|
||||
var keywords []types.Keyword
|
||||
if err := json.Unmarshal([]byte(kw), &keywords); err == nil {
|
||||
return keywords, nil
|
||||
}
|
||||
return []string{kw}, nil
|
||||
return []types.Keyword{{K: kw, W: 0.5}}, nil
|
||||
}
|
||||
return []string{}, nil
|
||||
return []types.Keyword{}, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -45,15 +45,14 @@ func (p *AgentProvider) Generate(ctx *agentContext.Context, input *Input) (*Resu
|
|||
var lastLintErrors string
|
||||
|
||||
for attempt := 1; attempt <= MaxRetries; attempt++ {
|
||||
// Build the request message
|
||||
requestData := p.buildRequestData(input, attempt, lastLintErrors)
|
||||
requestJSON, _ := json.Marshal(requestData)
|
||||
// Build the request message in the format expected by querydsl agent
|
||||
requestMessage := p.buildRequestMessage(input, attempt, lastLintErrors)
|
||||
|
||||
// Create message for the agent
|
||||
messages := []agentContext.Message{
|
||||
{
|
||||
Role: "user",
|
||||
Content: string(requestJSON),
|
||||
Content: requestMessage,
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -71,8 +70,8 @@ func (p *AgentProvider) Generate(ctx *agentContext.Context, input *Input) (*Resu
|
|||
continue
|
||||
}
|
||||
|
||||
// Parse the result
|
||||
genResult, err := p.parseResult(result)
|
||||
// Parse the result from response
|
||||
genResult, err := p.parseResponse(result)
|
||||
if err != nil {
|
||||
lastError = err
|
||||
continue
|
||||
|
|
@ -103,27 +102,32 @@ func (p *AgentProvider) Generate(ctx *agentContext.Context, input *Input) (*Resu
|
|||
return nil, fmt.Errorf("QueryDSL generation failed after %d attempts: %w", MaxRetries, lastError)
|
||||
}
|
||||
|
||||
// buildRequestData constructs the request data for the agent
|
||||
func (p *AgentProvider) buildRequestData(input *Input, attempt int, lastLintErrors string) map[string]interface{} {
|
||||
// buildRequestMessage constructs the request message for the agent
|
||||
// Returns JSON format for structured communication with the agent
|
||||
func (p *AgentProvider) buildRequestMessage(input *Input, attempt int, lastLintErrors string) string {
|
||||
// Build request data as JSON
|
||||
requestData := map[string]interface{}{
|
||||
"query": input.Query,
|
||||
"models": input.ModelIDs,
|
||||
"limit": input.Limit,
|
||||
}
|
||||
|
||||
// Add optional fields
|
||||
if len(input.Wheres) > 0 {
|
||||
requestData["wheres"] = input.Wheres
|
||||
// Add schema from extra params if provided
|
||||
if input.ExtraParams != nil {
|
||||
if schema, ok := input.ExtraParams["schema"]; ok {
|
||||
requestData["schema"] = schema
|
||||
}
|
||||
}
|
||||
if len(input.Orders) > 0 {
|
||||
requestData["orders"] = input.Orders
|
||||
|
||||
// Add scenario hint if specified (filter, aggregation, join, complex)
|
||||
if input.Scenario != "" {
|
||||
requestData["scenario"] = string(input.Scenario)
|
||||
}
|
||||
|
||||
// Add allowed fields if specified
|
||||
if len(input.AllowedFields) > 0 {
|
||||
requestData["allowed_fields"] = input.AllowedFields
|
||||
}
|
||||
if len(input.ExtraParams) > 0 {
|
||||
requestData["extra"] = input.ExtraParams
|
||||
}
|
||||
|
||||
// Add retry context if this is a retry attempt
|
||||
if attempt > 1 && lastLintErrors != "" {
|
||||
|
|
@ -134,7 +138,8 @@ func (p *AgentProvider) buildRequestData(input *Input, attempt int, lastLintErro
|
|||
}
|
||||
}
|
||||
|
||||
return requestData
|
||||
jsonBytes, _ := json.Marshal(requestData)
|
||||
return string(jsonBytes)
|
||||
}
|
||||
|
||||
// validateDSL validates the generated QueryDSL using the linter
|
||||
|
|
@ -150,18 +155,36 @@ func (p *AgentProvider) validateDSL(dsl *gou.QueryDSL) *linter.LintResult {
|
|||
return lintResult
|
||||
}
|
||||
|
||||
// parseResult extracts QueryDSL from the agent's response
|
||||
// The agent should return data in NextHookResponse format: { data: { dsl: {...}, explain: "..." } }
|
||||
// The Stream() response wraps this in: { next: { data: { dsl: {...} } } }
|
||||
func (p *AgentProvider) parseResult(result interface{}) (*Result, error) {
|
||||
if result == nil {
|
||||
// parseResponse extracts QueryDSL from the agent's *context.Response
|
||||
// Now that agent.Stream() returns *context.Response directly,
|
||||
// we can access fields without type assertions.
|
||||
//
|
||||
// The querydsl agent returns QueryDSL in response.Next field
|
||||
// Or returns error JSON: {"error": "code", "message": "..."}
|
||||
func (p *AgentProvider) parseResponse(response *agentContext.Response) (*Result, error) {
|
||||
if response == nil {
|
||||
return &Result{}, nil
|
||||
}
|
||||
|
||||
// Try to convert to map first (most common case)
|
||||
// Check Next field first (custom hook data)
|
||||
if response.Next != nil {
|
||||
return p.parseNextData(response.Next)
|
||||
}
|
||||
|
||||
// No Next data, return empty result
|
||||
return &Result{}, nil
|
||||
}
|
||||
|
||||
// parseNextData extracts QueryDSL from Next hook data
|
||||
func (p *AgentProvider) parseNextData(next interface{}) (*Result, error) {
|
||||
if next == nil {
|
||||
return &Result{}, nil
|
||||
}
|
||||
|
||||
// Try to convert to map first
|
||||
var data map[string]interface{}
|
||||
|
||||
switch v := result.(type) {
|
||||
switch v := next.(type) {
|
||||
case map[string]interface{}:
|
||||
data = v
|
||||
case string:
|
||||
|
|
@ -171,7 +194,7 @@ func (p *AgentProvider) parseResult(result interface{}) (*Result, error) {
|
|||
}
|
||||
default:
|
||||
// Try to marshal and unmarshal
|
||||
jsonBytes, err := json.Marshal(result)
|
||||
jsonBytes, err := json.Marshal(next)
|
||||
if err != nil {
|
||||
return &Result{}, nil
|
||||
}
|
||||
|
|
@ -180,49 +203,72 @@ func (p *AgentProvider) parseResult(result interface{}) (*Result, error) {
|
|||
}
|
||||
}
|
||||
|
||||
// Check for "next" field (custom hook data from NextHookResponse)
|
||||
// Stream() returns: { next: { data: { dsl: {...} } } }
|
||||
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 &Result{}, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Extract QueryDSL from data
|
||||
// Try common field names: "dsl", "data", "data.dsl"
|
||||
genResult := &Result{}
|
||||
|
||||
// Get explain if present
|
||||
if explain, ok := data["explain"].(string); ok {
|
||||
genResult.Explain = explain
|
||||
// Check for error response: {"error": "code", "message": "..."}
|
||||
if errCode, hasError := data["error"]; hasError {
|
||||
errMsg := ""
|
||||
if msg, ok := data["message"].(string); ok {
|
||||
errMsg = msg
|
||||
}
|
||||
return nil, fmt.Errorf("QueryDSL generation error [%v]: %s", errCode, errMsg)
|
||||
}
|
||||
|
||||
// Get warnings if present
|
||||
if warnings, ok := data["warnings"]; ok {
|
||||
genResult.Warnings = p.extractWarnings(warnings)
|
||||
// Check if this is a direct QueryDSL (has "from" or "select" field)
|
||||
// The querydsl agent returns QueryDSL directly, e.g., {"select": [...], "from": "table", ...}
|
||||
if _, hasFrom := data["from"]; hasFrom {
|
||||
genResult.DSL = p.extractDSL(data)
|
||||
return genResult, nil
|
||||
}
|
||||
if _, hasSelect := data["select"]; hasSelect {
|
||||
genResult.DSL = p.extractDSL(data)
|
||||
return genResult, nil
|
||||
}
|
||||
|
||||
// Get DSL
|
||||
// Check for "dsl" field wrapper: { dsl: {...} }
|
||||
if dsl, ok := data["dsl"]; ok {
|
||||
genResult.DSL = p.extractDSL(dsl)
|
||||
} else if d, ok := data["data"]; ok {
|
||||
if explain, ok := data["explain"].(string); ok {
|
||||
genResult.Explain = explain
|
||||
}
|
||||
if warnings, ok := data["warnings"]; ok {
|
||||
genResult.Warnings = p.extractWarnings(warnings)
|
||||
}
|
||||
return genResult, nil
|
||||
}
|
||||
|
||||
// Check for "data" field wrapper: { data: { dsl: {...}, explain: "...", warnings: [] } }
|
||||
if d, ok := data["data"]; ok {
|
||||
if dm, ok := d.(map[string]interface{}); ok {
|
||||
// Check if data.data contains dsl field: { data: { dsl: {...} } }
|
||||
if dsl, ok := dm["dsl"]; ok {
|
||||
genResult.DSL = p.extractDSL(dsl)
|
||||
} else if _, hasFrom := dm["from"]; hasFrom {
|
||||
// data.data is directly a QueryDSL (from __yao.querydsl Next hook)
|
||||
genResult.DSL = p.extractDSL(dm)
|
||||
} else if _, hasSelect := dm["select"]; hasSelect {
|
||||
// data.data is directly a QueryDSL
|
||||
genResult.DSL = p.extractDSL(dm)
|
||||
}
|
||||
// Extract explain and warnings from data.data
|
||||
if explain, ok := dm["explain"].(string); ok {
|
||||
genResult.Explain = explain
|
||||
}
|
||||
if warnings, ok := dm["warnings"]; ok {
|
||||
genResult.Warnings = p.extractWarnings(warnings)
|
||||
}
|
||||
return genResult, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: Get explain and warnings from top level
|
||||
if explain, ok := data["explain"].(string); ok {
|
||||
genResult.Explain = explain
|
||||
}
|
||||
if warnings, ok := data["warnings"]; ok {
|
||||
genResult.Warnings = p.extractWarnings(warnings)
|
||||
}
|
||||
|
||||
return genResult, nil
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,12 +2,14 @@ package querydsl
|
|||
|
||||
import (
|
||||
"github.com/yaoapp/gou/query/gou"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
)
|
||||
|
||||
// Input contains all information needed to generate QueryDSL
|
||||
type Input struct {
|
||||
Query string // Natural language query
|
||||
ModelIDs []string // Target model IDs (e.g., ["user", "order", "product"])
|
||||
Scenario types.ScenarioType // QueryDSL scenario: "filter", "aggregation", "join", "complex"
|
||||
Wheres []gou.Where // Pre-defined filters (optional)
|
||||
Orders gou.Orders // Sort orders (optional)
|
||||
AllowedFields []string // Allowed fields whitelist (optional, for security validation)
|
||||
|
|
|
|||
|
|
@ -59,20 +59,24 @@ func (p *AgentProvider) Rerank(ctx *context.Context, query string, items []*type
|
|||
},
|
||||
}
|
||||
|
||||
result, err := agent.Stream(ctx, messages, options)
|
||||
response, err := agent.Stream(ctx, messages, options)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("agent stream failed: %w", err)
|
||||
}
|
||||
|
||||
// Parse response
|
||||
return p.parseResponse(result, items, opts)
|
||||
// Parse response from response.Next
|
||||
return p.parseAgentResponse(response, items, opts)
|
||||
}
|
||||
|
||||
// parseResponse extracts reranked items from agent response
|
||||
// The response format from agent.Stream is typically:
|
||||
// { "next": { "data": { "order": [...] } } }
|
||||
func (p *AgentProvider) parseResponse(result interface{}, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
|
||||
if result == nil {
|
||||
// parseAgentResponse extracts reranked items from agent's *context.Response
|
||||
// Now that agent.Stream() returns *context.Response directly,
|
||||
// we can access fields without type assertions.
|
||||
//
|
||||
// Expected response.Next format:
|
||||
// { "order": ["ref_001", "ref_003", "ref_002"] }
|
||||
// Or: { "items": [{ "citation_id": "ref_001", ... }, ...] }
|
||||
func (p *AgentProvider) parseAgentResponse(response *context.Response, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) {
|
||||
if response == nil || response.Next == nil {
|
||||
return originalItems, nil
|
||||
}
|
||||
|
||||
|
|
@ -84,20 +88,20 @@ func (p *AgentProvider) parseResponse(result interface{}, originalItems []*types
|
|||
}
|
||||
}
|
||||
|
||||
// Extract response data
|
||||
response := extractResponseData(result)
|
||||
if response == nil {
|
||||
// Extract response data from Next field
|
||||
data := extractNextData(response.Next)
|
||||
if data == nil {
|
||||
return originalItems, nil
|
||||
}
|
||||
|
||||
// Try to get reranked order from response
|
||||
// Try to get reranked order from data
|
||||
// Expected format: { "order": ["ref_001", "ref_003", "ref_002"] }
|
||||
// Or: { "items": [{ "citation_id": "ref_001", ... }, ...] }
|
||||
|
||||
var reranked []*types.ResultItem
|
||||
|
||||
// Try "order" field (list of citation IDs)
|
||||
if order, ok := response["order"]; ok {
|
||||
if order, ok := data["order"]; ok {
|
||||
if orderList := toStringSlice(order); len(orderList) > 0 {
|
||||
for _, id := range orderList {
|
||||
if item, exists := itemMap[id]; exists {
|
||||
|
|
@ -116,7 +120,7 @@ func (p *AgentProvider) parseResponse(result interface{}, originalItems []*types
|
|||
|
||||
// Try "items" field (full items or items with citation_id)
|
||||
if len(reranked) == 0 {
|
||||
if items, ok := response["items"]; ok {
|
||||
if items, ok := data["items"]; ok {
|
||||
if itemsList := toItemsList(items); len(itemsList) > 0 {
|
||||
for _, respItem := range itemsList {
|
||||
// Check if it's just a reference or full item
|
||||
|
|
@ -150,20 +154,16 @@ func (p *AgentProvider) parseResponse(result interface{}, originalItems []*types
|
|||
return reranked, nil
|
||||
}
|
||||
|
||||
// extractResponseData extracts the actual response data from agent.Stream result
|
||||
// Handles nested structures like { "next": { "data": { ... } } }
|
||||
func extractResponseData(result interface{}) map[string]interface{} {
|
||||
switch v := result.(type) {
|
||||
// extractNextData extracts the actual data from response.Next field
|
||||
// Handles nested structures like { "data": { ... } }
|
||||
func extractNextData(next interface{}) map[string]interface{} {
|
||||
if next == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch v := next.(type) {
|
||||
case map[string]interface{}:
|
||||
// Check for "next" wrapper (from NextHookResponse)
|
||||
if next, ok := v["next"].(map[string]interface{}); ok {
|
||||
// Check for "data" inside next
|
||||
if data, ok := next["data"].(map[string]interface{}); ok {
|
||||
return data
|
||||
}
|
||||
return next
|
||||
}
|
||||
// Check for direct "data" wrapper
|
||||
// Check for "data" wrapper
|
||||
if data, ok := v["data"].(map[string]interface{}); ok {
|
||||
return data
|
||||
}
|
||||
|
|
@ -172,16 +172,14 @@ func extractResponseData(result interface{}) map[string]interface{} {
|
|||
// Try to parse as JSON
|
||||
var data map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(v), &data); err == nil {
|
||||
return extractResponseData(data)
|
||||
return extractNextData(data)
|
||||
}
|
||||
}
|
||||
// Try to handle other types by converting to JSON and back
|
||||
if result != nil {
|
||||
if bytes, err := json.Marshal(result); err == nil {
|
||||
var data map[string]interface{}
|
||||
if err := json.Unmarshal(bytes, &data); err == nil {
|
||||
return extractResponseData(data)
|
||||
}
|
||||
if bytes, err := json.Marshal(next); err == nil {
|
||||
var data map[string]interface{}
|
||||
if err := json.Unmarshal(bytes, &data); err == nil {
|
||||
return extractNextData(data)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
|
|
|
|||
|
|
@ -60,8 +60,14 @@ func (s *Searcher) Search(ctx *context.Context, req *types.Request) (*types.Resu
|
|||
return &types.Result{Error: "unsupported search type"}, nil
|
||||
}
|
||||
|
||||
// Execute search
|
||||
result, err := handler.Search(req)
|
||||
// Execute search - use context if handler supports it
|
||||
var result *types.Result
|
||||
var err error
|
||||
if ctxHandler, ok := handler.(interfaces.ContextHandler); ok {
|
||||
result, err = ctxHandler.SearchWithContext(ctx, req)
|
||||
} else {
|
||||
result, err = handler.Search(req)
|
||||
}
|
||||
if err != nil {
|
||||
return &types.Result{Error: err.Error()}, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,6 +14,17 @@ const (
|
|||
SearchTypeDB SearchType = "db" // Database search (Yao Model/QueryDSL)
|
||||
)
|
||||
|
||||
// ScenarioType represents the QueryDSL generation scenario
|
||||
type ScenarioType string
|
||||
|
||||
// ScenarioType constants for QueryDSL generation
|
||||
const (
|
||||
ScenarioFilter ScenarioType = "filter" // Simple filtering queries
|
||||
ScenarioAggregation ScenarioType = "aggregation" // Aggregation/grouping queries
|
||||
ScenarioJoin ScenarioType = "join" // Multi-table join queries
|
||||
ScenarioComplex ScenarioType = "complex" // Complex queries combining multiple features
|
||||
)
|
||||
|
||||
// SourceType represents where the search result came from
|
||||
type SourceType string
|
||||
|
||||
|
|
@ -42,10 +53,11 @@ type Request struct {
|
|||
Graph bool `json:"graph,omitempty"` // Enable graph association
|
||||
|
||||
// Database search specific
|
||||
Models []string `json:"models,omitempty"` // Model IDs (e.g., "user", "agents.mybot.product")
|
||||
Wheres []gou.Where `json:"wheres,omitempty"` // Pre-defined filters (optional), uses GOU QueryDSL Where
|
||||
Orders gou.Orders `json:"orders,omitempty"` // Sort orders (optional), uses GOU QueryDSL Orders
|
||||
Select []string `json:"select,omitempty"` // Fields to return (optional)
|
||||
Models []string `json:"models,omitempty"` // Model IDs (e.g., "user", "agents.mybot.product")
|
||||
Scenario ScenarioType `json:"scenario,omitempty"` // QueryDSL scenario: "filter", "aggregation", "join", "complex"
|
||||
Wheres []gou.Where `json:"wheres,omitempty"` // Pre-defined filters (optional), uses GOU QueryDSL Where
|
||||
Orders gou.Orders `json:"orders,omitempty"` // Sort orders (optional), uses GOU QueryDSL Orders
|
||||
Select []string `json:"select,omitempty"` // Fields to return (optional)
|
||||
|
||||
// Reranking
|
||||
Rerank *RerankOptions `json:"rerank,omitempty"`
|
||||
|
|
@ -108,6 +120,12 @@ type ProcessedQuery struct {
|
|||
DSL *gou.QueryDSL `json:"dsl,omitempty"` // For DB search, uses GOU QueryDSL
|
||||
}
|
||||
|
||||
// Keyword represents an extracted keyword with weight
|
||||
type Keyword struct {
|
||||
K string `json:"k"` // Keyword text
|
||||
W float64 `json:"w"` // Weight (0.1-1.0), higher means more relevant
|
||||
}
|
||||
|
||||
// Note: For QueryDSL and Model types, use GOU types directly:
|
||||
// - github.com/yaoapp/gou/query/gou.QueryDSL
|
||||
// - github.com/yaoapp/gou/model.Model
|
||||
|
|
|
|||
|
|
@ -110,10 +110,58 @@ Each line is a JSON object:
|
|||
| `expected` | any | No | Expected output (exact match) |
|
||||
| `user` | string | No | Override user ID |
|
||||
| `team` | string | No | Override team ID |
|
||||
| `options` | Options | No | Context options (see below) |
|
||||
| `timeout` | string | No | Override timeout (e.g., "30s") |
|
||||
| `skip` | bool | No | Skip this test |
|
||||
| `metadata` | map | No | Additional metadata |
|
||||
|
||||
### Options
|
||||
|
||||
The `options` field allows per-test-case configuration that maps to `context.Options`:
|
||||
|
||||
| Field | Type | Description |
|
||||
| ------------------------ | ------- | -------------------------------------------------- |
|
||||
| `connector` | string | Override connector (e.g., `"deepseek.v3"`) |
|
||||
| `mode` | string | Agent mode (default: `"chat"`) |
|
||||
| `search` | bool | Enable/disable search mode (default: `true`) |
|
||||
| `disable_global_prompts` | bool | Temporarily disable global prompts |
|
||||
| `metadata` | map | Custom data passed to hooks (e.g., scenario) |
|
||||
| `skip` | object | Skip configuration (see below) |
|
||||
|
||||
#### Options.skip
|
||||
|
||||
| Field | Type | Description |
|
||||
| --------- | ---- | ------------------------ |
|
||||
| `history` | bool | Skip history loading |
|
||||
| `trace` | bool | Skip trace logging |
|
||||
| `output` | bool | Skip output to client |
|
||||
| `keyword` | bool | Skip keyword extraction |
|
||||
| `search` | bool | Skip auto search |
|
||||
|
||||
**Example with options:**
|
||||
|
||||
```jsonl
|
||||
{
|
||||
"id": "T001",
|
||||
"input": "Query users with status active",
|
||||
"options": {
|
||||
"connector": "deepseek.v3",
|
||||
"metadata": {"scenario": "filter"},
|
||||
"skip": {"trace": true}
|
||||
},
|
||||
"assert": {"type": "json_path", "path": "from", "value": "users"}
|
||||
}
|
||||
```
|
||||
|
||||
**Using metadata for hook scenarios:**
|
||||
|
||||
The `options.metadata` field is passed to agent hooks. For example, a Create Hook can read `options.metadata.scenario` to select different prompt presets:
|
||||
|
||||
```jsonl
|
||||
{"id": "T001", "input": "...", "options": {"metadata": {"scenario": "aggregation"}}}
|
||||
{"id": "T002", "input": "...", "options": {"metadata": {"scenario": "join"}}}
|
||||
```
|
||||
|
||||
### Input Types
|
||||
|
||||
| Type | Description | Example |
|
||||
|
|
@ -225,9 +273,28 @@ return { pass: true, message: "Validation passed" };
|
|||
### JSON Path Notes
|
||||
|
||||
- Supports dot-notation: `$.field.subfield` or `field.subfield`
|
||||
- Supports array indexing: `field[0]`, `field[0].subfield`, `field[0].nested[1]`
|
||||
- Supports multiple expected values (OR logic): `"value": ["a", "b"]` - passes if actual matches any
|
||||
- Auto-extracts JSON from markdown code blocks (` ```json ... ``` `)
|
||||
- Works with both string output and structured objects
|
||||
|
||||
**Array index examples:**
|
||||
|
||||
```jsonl
|
||||
{"id": "T001", "assert": {"type": "json_path", "path": "wheres[0].like", "value": "%test%"}}
|
||||
{"id": "T002", "assert": {"type": "json_path", "path": "wheres[0].in[0]", "value": "pending"}}
|
||||
{"id": "T003", "assert": {"type": "json_path", "path": "joins[0].from", "value": "users"}}
|
||||
{"id": "T004", "assert": {"type": "json_path", "path": "groups[0]", "value": "category"}}
|
||||
```
|
||||
|
||||
**Multiple expected values (OR logic):**
|
||||
|
||||
```jsonl
|
||||
{"id": "T005", "assert": {"type": "json_path", "path": "error", "value": ["missing_schema", "missing_query"]}}
|
||||
```
|
||||
|
||||
This passes if `error` equals either `"missing_schema"` or `"missing_query"`.
|
||||
|
||||
## Output Formats
|
||||
|
||||
Determined by `-o` file extension:
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
|
|
@ -244,7 +245,7 @@ func (a *Asserter) assertJSONPath(assertion *Assertion, output interface{}) *Ass
|
|||
jsonData = v
|
||||
default:
|
||||
result.Passed = false
|
||||
result.Message = "output is not a JSON object or array"
|
||||
result.Message = fmt.Sprintf("output is not a JSON object or array, got: %T = %v", output, truncateOutput(output, 200))
|
||||
return result
|
||||
}
|
||||
|
||||
|
|
@ -253,6 +254,8 @@ func (a *Asserter) assertJSONPath(assertion *Assertion, output interface{}) *Ass
|
|||
actual := a.extractPath(jsonData, path)
|
||||
result.Actual = actual
|
||||
|
||||
// Compare expected value with actual value
|
||||
// First, try direct comparison (handles both primitive values and arrays)
|
||||
if validateOutput(actual, assertion.Value) {
|
||||
result.Passed = true
|
||||
result.Message = fmt.Sprintf("path '%s' equals expected value", assertion.Path)
|
||||
|
|
@ -264,27 +267,113 @@ func (a *Asserter) assertJSONPath(assertion *Assertion, output interface{}) *Ass
|
|||
return result
|
||||
}
|
||||
|
||||
// extractPath extracts a value from JSON using a simple dot-notation path
|
||||
// truncateOutput truncates output for error messages
|
||||
func truncateOutput(output interface{}, maxLen int) string {
|
||||
var s string
|
||||
switch v := output.(type) {
|
||||
case string:
|
||||
s = v
|
||||
case nil:
|
||||
return "<nil>"
|
||||
default:
|
||||
bytes, err := jsoniter.Marshal(v)
|
||||
if err != nil {
|
||||
s = fmt.Sprintf("%v", v)
|
||||
} else {
|
||||
s = string(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
if len(s) > maxLen {
|
||||
return s[:maxLen] + "..."
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// extractPath extracts a value from JSON using dot-notation path with array index support
|
||||
// Supports: "field", "field.nested", "field[0]", "field[0].nested", "field.nested[0].value"
|
||||
func (a *Asserter) extractPath(data interface{}, path string) interface{} {
|
||||
parts := strings.Split(path, ".")
|
||||
current := data
|
||||
|
||||
for _, part := range parts {
|
||||
if part == "" {
|
||||
// Parse path into segments, handling both dots and array indices
|
||||
// e.g., "wheres[0].like" -> ["wheres", "[0]", "like"]
|
||||
segments := parsePathSegments(path)
|
||||
|
||||
for _, segment := range segments {
|
||||
if segment == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
switch v := current.(type) {
|
||||
case map[string]interface{}:
|
||||
current = v[part]
|
||||
default:
|
||||
return nil
|
||||
// Check if this is an array index like "[0]"
|
||||
if strings.HasPrefix(segment, "[") && strings.HasSuffix(segment, "]") {
|
||||
indexStr := segment[1 : len(segment)-1]
|
||||
index, err := strconv.Atoi(indexStr)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
arr, ok := current.([]interface{})
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
if index < 0 || index >= len(arr) {
|
||||
return nil
|
||||
}
|
||||
current = arr[index]
|
||||
} else {
|
||||
// Regular field access
|
||||
switch v := current.(type) {
|
||||
case map[string]interface{}:
|
||||
current = v[segment]
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return current
|
||||
}
|
||||
|
||||
// parsePathSegments splits a path like "wheres[0].like" into ["wheres", "[0]", "like"]
|
||||
func parsePathSegments(path string) []string {
|
||||
var segments []string
|
||||
var current strings.Builder
|
||||
|
||||
for i := 0; i < len(path); i++ {
|
||||
ch := path[i]
|
||||
switch ch {
|
||||
case '.':
|
||||
if current.Len() > 0 {
|
||||
segments = append(segments, current.String())
|
||||
current.Reset()
|
||||
}
|
||||
case '[':
|
||||
if current.Len() > 0 {
|
||||
segments = append(segments, current.String())
|
||||
current.Reset()
|
||||
}
|
||||
// Find the closing bracket
|
||||
j := i + 1
|
||||
for j < len(path) && path[j] != ']' {
|
||||
j++
|
||||
}
|
||||
if j < len(path) {
|
||||
segments = append(segments, path[i:j+1]) // Include "[" and "]"
|
||||
i = j
|
||||
}
|
||||
default:
|
||||
current.WriteByte(ch)
|
||||
}
|
||||
}
|
||||
|
||||
if current.Len() > 0 {
|
||||
segments = append(segments, current.String())
|
||||
}
|
||||
|
||||
return segments
|
||||
}
|
||||
|
||||
// assertRegex checks if output matches a regex pattern
|
||||
func (a *Asserter) assertRegex(assertion *Assertion, output interface{}) *AssertionResult {
|
||||
result := &AssertionResult{
|
||||
|
|
|
|||
|
|
@ -595,13 +595,13 @@ func (r *AgentReporter) Write(report *Report, w io.Writer) error {
|
|||
},
|
||||
}
|
||||
|
||||
result, err := agent.Stream(ctx, messages, options)
|
||||
response, err := agent.Stream(ctx, messages, options)
|
||||
if err != nil {
|
||||
return fmt.Errorf("reporter agent call failed: %w", err)
|
||||
}
|
||||
|
||||
// Extract content from result
|
||||
content, err := r.extractContent(result)
|
||||
// Extract content from response
|
||||
content, err := r.extractContent(response)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to extract report content: %w", err)
|
||||
}
|
||||
|
|
@ -615,55 +615,25 @@ func (r *AgentReporter) Write(report *Report, w io.Writer) error {
|
|||
return nil
|
||||
}
|
||||
|
||||
// extractContent extracts the report content from the agent's response
|
||||
func (r *AgentReporter) extractContent(result interface{}) (string, error) {
|
||||
if result == nil {
|
||||
return "", fmt.Errorf("agent returned nil result")
|
||||
// extractContent extracts the report content from the agent's *context.Response
|
||||
// Now that agent.Stream() returns *context.Response directly,
|
||||
// we can access fields without type assertions.
|
||||
func (r *AgentReporter) extractContent(response *context.Response) (string, error) {
|
||||
if response == nil {
|
||||
return "", fmt.Errorf("agent returned nil response")
|
||||
}
|
||||
|
||||
// Try to convert to map first (context.Response)
|
||||
switch v := result.(type) {
|
||||
case string:
|
||||
return v, nil
|
||||
|
||||
case *context.Response:
|
||||
// Extract from completion content
|
||||
if v.Completion != nil && v.Completion.Content != nil {
|
||||
return r.contentToString(v.Completion.Content)
|
||||
}
|
||||
// Try next field
|
||||
if v.Next != nil {
|
||||
return r.contentToString(v.Next)
|
||||
}
|
||||
return "", fmt.Errorf("no content in response")
|
||||
|
||||
case map[string]interface{}:
|
||||
// Check for completion.content
|
||||
if completion, ok := v["completion"].(map[string]interface{}); ok {
|
||||
if content, ok := completion["content"]; ok {
|
||||
return r.contentToString(content)
|
||||
}
|
||||
}
|
||||
// Check for next
|
||||
if next, ok := v["next"]; ok {
|
||||
return r.contentToString(next)
|
||||
}
|
||||
// Check for content directly
|
||||
if content, ok := v["content"]; ok {
|
||||
return r.contentToString(content)
|
||||
}
|
||||
// Marshal the whole thing
|
||||
jsonBytes, _ := jsoniter.Marshal(v)
|
||||
return string(jsonBytes), nil
|
||||
|
||||
default:
|
||||
// Try to marshal as JSON
|
||||
jsonBytes, err := jsoniter.Marshal(result)
|
||||
if err != nil {
|
||||
return fmt.Sprintf("%v", result), nil
|
||||
}
|
||||
return string(jsonBytes), nil
|
||||
// Priority 1: Check Next field (custom hook data)
|
||||
if response.Next != nil {
|
||||
return r.contentToString(response.Next)
|
||||
}
|
||||
|
||||
// Priority 2: Extract from completion content
|
||||
if response.Completion != nil && response.Completion.Content != nil {
|
||||
return r.contentToString(response.Completion.Content)
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("no content in response")
|
||||
}
|
||||
|
||||
// contentToString converts various content types to string
|
||||
|
|
|
|||
|
|
@ -64,15 +64,8 @@ func (r *Executor) RunDirect() (*Report, error) {
|
|||
ctx := NewTestContextFromOptions(chatID, agentInfo.ID, r.opts, tc)
|
||||
defer ctx.Release()
|
||||
|
||||
// Set options: skip history (input already contains conversation), connector override
|
||||
opts := &context.Options{
|
||||
Skip: &context.Skip{
|
||||
History: true, // Skip history loading - input already contains full conversation
|
||||
},
|
||||
}
|
||||
if r.opts.Connector != "" {
|
||||
opts.Connector = r.opts.Connector
|
||||
}
|
||||
// Build context options
|
||||
opts := buildContextOptions(tc, r.opts)
|
||||
|
||||
// Create timeout context
|
||||
timeout := tc.GetTimeout(r.opts.Timeout)
|
||||
|
|
@ -103,12 +96,19 @@ func (r *Executor) RunDirect() (*Report, error) {
|
|||
output := extractOutput(response)
|
||||
r.output.DirectOutput(output)
|
||||
|
||||
// Determine connector: user-specified > agent default
|
||||
connector := r.opts.Connector
|
||||
if connector == "" {
|
||||
connector = agentInfo.Connector
|
||||
}
|
||||
|
||||
// Return minimal report (for exit code handling)
|
||||
return &Report{
|
||||
Summary: &Summary{
|
||||
Total: 1,
|
||||
Passed: 1,
|
||||
AgentID: agentInfo.ID,
|
||||
Total: 1,
|
||||
Passed: 1,
|
||||
AgentID: agentInfo.ID,
|
||||
Connector: connector,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
|
@ -165,13 +165,19 @@ func (r *Executor) RunTests() (*Report, error) {
|
|||
return nil, fmt.Errorf("failed to get assistant: %w", err)
|
||||
}
|
||||
|
||||
// Determine connector: user-specified > agent default
|
||||
connector := r.opts.Connector
|
||||
if connector == "" {
|
||||
connector = agentInfo.Connector
|
||||
}
|
||||
|
||||
// Create report
|
||||
report := &Report{
|
||||
Summary: &Summary{
|
||||
Total: len(testCases),
|
||||
AgentID: agentInfo.ID,
|
||||
AgentPath: agentInfo.Path,
|
||||
Connector: r.opts.Connector,
|
||||
Connector: connector,
|
||||
RunsPerCase: r.opts.Runs,
|
||||
},
|
||||
Environment: NewEnvironment(r.opts.UserID, r.opts.TeamID),
|
||||
|
|
@ -279,6 +285,7 @@ func (r *Executor) runSingleTest(ast *assistant.Assistant, tc *Case, agentID str
|
|||
ID: tc.ID,
|
||||
Input: tc.Input,
|
||||
Expected: tc.Expected,
|
||||
Options: tc.Options,
|
||||
}
|
||||
|
||||
// Parse input to messages
|
||||
|
|
@ -297,15 +304,8 @@ func (r *Executor) runSingleTest(ast *assistant.Assistant, tc *Case, agentID str
|
|||
ctx := NewTestContextFromOptions(chatID, agentID, r.opts, tc)
|
||||
defer ctx.Release()
|
||||
|
||||
// Set options: skip history (input already contains conversation), connector override
|
||||
opts := &context.Options{
|
||||
Skip: &context.Skip{
|
||||
History: true, // Skip history loading - input already contains full conversation
|
||||
},
|
||||
}
|
||||
if r.opts.Connector != "" {
|
||||
opts.Connector = r.opts.Connector
|
||||
}
|
||||
// Build context options from test case and runner options
|
||||
opts := buildContextOptions(tc, r.opts)
|
||||
|
||||
// Create timeout context
|
||||
timeout := tc.GetTimeout(r.opts.Timeout)
|
||||
|
|
@ -477,23 +477,96 @@ func writeJSONLine(writer *bufio.Writer, data interface{}) error {
|
|||
return err
|
||||
}
|
||||
|
||||
// buildContextOptions builds context.Options from test case and runner options
|
||||
// Priority: test case options > runner options > defaults
|
||||
func buildContextOptions(tc *Case, runnerOpts *Options) *context.Options {
|
||||
opts := &context.Options{
|
||||
Skip: &context.Skip{
|
||||
History: true, // Default: skip history loading - input already contains full conversation
|
||||
},
|
||||
}
|
||||
|
||||
// Apply test case options if specified
|
||||
if tc.Options != nil {
|
||||
// Connector: test case > runner
|
||||
if tc.Options.Connector != "" {
|
||||
opts.Connector = tc.Options.Connector
|
||||
}
|
||||
|
||||
// Mode
|
||||
if tc.Options.Mode != "" {
|
||||
opts.Mode = tc.Options.Mode
|
||||
}
|
||||
|
||||
// DisableGlobalPrompts
|
||||
if tc.Options.DisableGlobalPrompts {
|
||||
opts.DisableGlobalPrompts = true
|
||||
}
|
||||
|
||||
// Search (pointer to distinguish unset from false)
|
||||
if tc.Options.Search != nil {
|
||||
opts.Search = tc.Options.Search
|
||||
}
|
||||
|
||||
// Metadata for hooks
|
||||
if tc.Options.Metadata != nil {
|
||||
opts.Metadata = tc.Options.Metadata
|
||||
}
|
||||
|
||||
// Skip options from test case
|
||||
if tc.Options.Skip != nil {
|
||||
opts.Skip.Trace = tc.Options.Skip.Trace
|
||||
opts.Skip.Output = tc.Options.Skip.Output
|
||||
opts.Skip.Keyword = tc.Options.Skip.Keyword
|
||||
opts.Skip.Search = tc.Options.Skip.Search
|
||||
// Note: History defaults to true for tests
|
||||
}
|
||||
}
|
||||
|
||||
// Runner connector override (highest priority)
|
||||
if runnerOpts != nil && runnerOpts.Connector != "" {
|
||||
opts.Connector = runnerOpts.Connector
|
||||
}
|
||||
|
||||
return opts
|
||||
}
|
||||
|
||||
// extractOutput extracts the output from the agent response
|
||||
func extractOutput(response interface{}) interface{} {
|
||||
// Priority: Next hook data (if non-empty) > Completion content > nil
|
||||
func extractOutput(response *context.Response) interface{} {
|
||||
if response == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Try to get completion content from context.Response
|
||||
if resp, ok := response.(*context.Response); ok {
|
||||
if resp.Completion != nil {
|
||||
return resp.Completion.Content
|
||||
}
|
||||
if resp.Next != nil {
|
||||
return resp.Next
|
||||
}
|
||||
// Prefer Next hook data if available and non-empty
|
||||
// response.Next is already the Data value (not NextHookResponse struct)
|
||||
if response.Next != nil && !isEmptyValue(response.Next) {
|
||||
return response.Next
|
||||
}
|
||||
// Fall back to raw completion content
|
||||
if response.Completion != nil {
|
||||
return response.Completion.Content
|
||||
}
|
||||
|
||||
return response
|
||||
return nil
|
||||
}
|
||||
|
||||
// isEmptyValue checks if a value is considered "empty" for output purposes
|
||||
func isEmptyValue(v interface{}) bool {
|
||||
if v == nil {
|
||||
return true
|
||||
}
|
||||
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
return val == ""
|
||||
case map[string]interface{}:
|
||||
return len(val) == 0
|
||||
case []interface{}:
|
||||
return len(val) == 0
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// validateOutput validates the actual output against expected
|
||||
|
|
|
|||
|
|
@ -205,8 +205,13 @@ type Case struct {
|
|||
TeamID string `json:"team,omitempty"`
|
||||
|
||||
// Metadata contains additional metadata for the test case
|
||||
// This is passed to ctx.Metadata and can be used by Create Hook
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
|
||||
// Options contains context options for this test case
|
||||
// Supports: connector, skip (history, trace, output, keyword, search), mode
|
||||
Options *CaseOptions `json:"options,omitempty"`
|
||||
|
||||
// Skip indicates whether to skip this test case
|
||||
Skip bool `json:"skip,omitempty"`
|
||||
|
||||
|
|
@ -215,6 +220,38 @@ type Case struct {
|
|||
Timeout string `json:"timeout,omitempty"`
|
||||
}
|
||||
|
||||
// CaseOptions represents per-test-case context options
|
||||
// Maps to context.Options fields
|
||||
type CaseOptions struct {
|
||||
// Connector overrides the agent's default connector
|
||||
Connector string `json:"connector,omitempty"`
|
||||
|
||||
// Skip configuration
|
||||
Skip *CaseSkipOptions `json:"skip,omitempty"`
|
||||
|
||||
// DisableGlobalPrompts temporarily disables global prompts for this request
|
||||
DisableGlobalPrompts bool `json:"disable_global_prompts,omitempty"`
|
||||
|
||||
// Search mode, default is true (use pointer to distinguish unset from false)
|
||||
Search *bool `json:"search,omitempty"`
|
||||
|
||||
// Mode is the agent mode (default: "chat")
|
||||
Mode string `json:"mode,omitempty"`
|
||||
|
||||
// Metadata for passing custom data to hooks (e.g., scenario selection)
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// CaseSkipOptions represents skip configuration for a test case
|
||||
// Maps to context.Skip fields
|
||||
type CaseSkipOptions struct {
|
||||
History bool `json:"history,omitempty"` // Skip history loading
|
||||
Trace bool `json:"trace,omitempty"` // Skip trace logging
|
||||
Output bool `json:"output,omitempty"` // Skip output to client
|
||||
Keyword bool `json:"keyword,omitempty"` // Skip keyword extraction
|
||||
Search bool `json:"search,omitempty"` // Skip auto search
|
||||
}
|
||||
|
||||
// Assertion represents a single assertion rule
|
||||
type Assertion struct {
|
||||
// Type is the assertion type:
|
||||
|
|
@ -334,6 +371,9 @@ type Result struct {
|
|||
// Error contains the error message if status is failed/error/timeout
|
||||
Error string `json:"error,omitempty"`
|
||||
|
||||
// Options contains the context options used for this test case
|
||||
Options *CaseOptions `json:"options,omitempty"`
|
||||
|
||||
// Metadata contains additional result metadata
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,6 +3,12 @@ package testutils
|
|||
import (
|
||||
"testing"
|
||||
|
||||
_ "github.com/yaoapp/gou/encoding"
|
||||
"github.com/yaoapp/gou/model"
|
||||
"github.com/yaoapp/gou/query"
|
||||
"github.com/yaoapp/gou/query/gou"
|
||||
_ "github.com/yaoapp/gou/text"
|
||||
"github.com/yaoapp/xun/capsule"
|
||||
"github.com/yaoapp/yao/agent"
|
||||
"github.com/yaoapp/yao/config"
|
||||
"github.com/yaoapp/yao/kb"
|
||||
|
|
@ -28,6 +34,21 @@ func Prepare(t *testing.T, opts ...interface{}) {
|
|||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Register default query engine (required for DB search)
|
||||
// capsule.Global is initialized by test.Prepare
|
||||
if _, has := query.Engines["default"]; !has && capsule.Global != nil {
|
||||
query.Register("default", &gou.Query{
|
||||
Query: capsule.Query(),
|
||||
GetTableName: func(s string) string {
|
||||
if mod, has := model.Models[s]; has {
|
||||
return mod.MetaData.Table.Name
|
||||
}
|
||||
return s
|
||||
},
|
||||
AESKey: config.Conf.DB.AESKey,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Clean clean the test environment
|
||||
|
|
|
|||
443
data/bindata.go
443
data/bindata.go
File diff suppressed because one or more lines are too long
2
go.mod
2
go.mod
|
|
@ -13,7 +13,7 @@ require (
|
|||
github.com/elazarl/go-bindata-assetfs v1.0.1
|
||||
github.com/emersion/go-imap v1.2.1
|
||||
github.com/evanw/esbuild v0.25.4
|
||||
github.com/expr-lang/expr v1.17.3
|
||||
github.com/expr-lang/expr v1.17.7
|
||||
github.com/fatih/color v1.18.0
|
||||
github.com/fsnotify/fsnotify v1.9.0
|
||||
github.com/gin-gonic/gin v1.10.1
|
||||
|
|
|
|||
2
go.sum
2
go.sum
|
|
@ -81,6 +81,8 @@ github.com/evanw/esbuild v0.25.4 h1:k1bTSim+usBG27w7BfOCorhgx3tO+6bAfMj5pR+6SKg=
|
|||
github.com/evanw/esbuild v0.25.4/go.mod h1:D2vIQZqV/vIf/VRHtViaUtViZmG7o+kKmlBfVQuRi48=
|
||||
github.com/expr-lang/expr v1.17.3 h1:myeTTuDFz7k6eFe/JPlep/UsiIjVhG61FMHFu63U7j0=
|
||||
github.com/expr-lang/expr v1.17.3/go.mod h1:8/vRC7+7HBzESEqt5kKpYXxrxkr31SaO8r40VO/1IT4=
|
||||
github.com/expr-lang/expr v1.17.7 h1:Q0xY/e/2aCIp8g9s/LGvMDCC5PxYlvHgDZRQ4y16JX8=
|
||||
github.com/expr-lang/expr v1.17.7/go.mod h1:8/vRC7+7HBzESEqt5kKpYXxrxkr31SaO8r40VO/1IT4=
|
||||
github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
|
||||
github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
|
||||
github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
|
||||
|
|
|
|||
|
|
@ -1,21 +1,25 @@
|
|||
- role: system
|
||||
content: |
|
||||
Extract keywords from text content.
|
||||
You are a keyword extraction tool, NOT a chatbot. Do NOT answer questions or provide explanations.
|
||||
Your ONLY job: analyze text and output a JSON array of keywords with weights.
|
||||
|
||||
Task:
|
||||
1. Analyze input text
|
||||
2. Extract important keywords
|
||||
3. Return JSON format
|
||||
4. Match input language
|
||||
Output format: ["keyword:weight", ...]
|
||||
|
||||
Response Format (JSON only):
|
||||
```json
|
||||
{"keywords": ["keyword1", "keyword2", ...]}
|
||||
```
|
||||
Weight:
|
||||
- 1.0: Core topic
|
||||
- 0.8-0.9: Key concepts
|
||||
- 0.6-0.7: Supporting themes
|
||||
- 0.4-0.5: Peripheral concepts
|
||||
|
||||
Guidelines:
|
||||
- Extract 5-15 keywords based on content length
|
||||
- Prioritize nouns, proper nouns, key concepts
|
||||
- Include single words and short phrases
|
||||
- Exclude common stop words
|
||||
- Keywords MUST be in the same language as input
|
||||
Examples:
|
||||
- Input(EN): "Developers frustrated with callback hell. ES2017 async/await improved readability."
|
||||
- Output: ["async/await:1", "asynchronous programming:0.9", "ES2017:0.8", "code readability:0.7"]
|
||||
|
||||
- Input(中文): "用户反馈APP启动慢、页面卡顿。需要优化首屏加载和内存占用。"
|
||||
- Output: ["性能优化:1", "启动速度:0.9", "内存管理:0.8", "用户体验:0.7"]
|
||||
|
||||
Rules:
|
||||
- ONLY output JSON array, nothing else
|
||||
- Max 5 keywords, sorted by weight
|
||||
- Summarize related concepts
|
||||
- Match input language (EN→EN, 中文→中文)
|
||||
|
|
|
|||
|
|
@ -1,13 +1,20 @@
|
|||
/**
|
||||
* Keyword Extraction Agent - Next Hook
|
||||
* Parses LLM response and extracts keywords with error tolerance
|
||||
* Parses LLM response and extracts keywords with weight
|
||||
* Format: ["keyword:weight", ...] -> [{k, w}, ...]
|
||||
*/
|
||||
|
||||
// @ts-nocheck
|
||||
|
||||
/** Keyword with weight */
|
||||
interface Keyword {
|
||||
k: string; // keyword
|
||||
w: number; // weight (0.1-1.0)
|
||||
}
|
||||
|
||||
/**
|
||||
* Next hook - processes keyword extraction response
|
||||
* Uses text.ExtractJSON for fault-tolerant JSON extraction from LLM output
|
||||
* Parses format: ["keyword1:0.9", "keyword2:0.8", ...]
|
||||
*/
|
||||
function Next(
|
||||
ctx: agent.Context,
|
||||
|
|
@ -21,22 +28,17 @@ function Next(
|
|||
}
|
||||
|
||||
const content = completion.content;
|
||||
let keywords: string[] = [];
|
||||
let keywords: Keyword[] = [];
|
||||
|
||||
try {
|
||||
// Use text.ExtractJSON for fault-tolerant extraction
|
||||
// Handles markdown code blocks, broken JSON, etc.
|
||||
const parsed = Process("text.ExtractJSON", content) as {
|
||||
keywords?: string[];
|
||||
} | null;
|
||||
// Extract JSON array from response
|
||||
const parsed = Process("text.ExtractJSON", content) as string[] | null;
|
||||
|
||||
if (parsed && Array.isArray(parsed.keywords)) {
|
||||
keywords = parsed.keywords.filter(
|
||||
(k) => typeof k === "string" && k.trim().length > 0
|
||||
);
|
||||
if (parsed && Array.isArray(parsed)) {
|
||||
keywords = parseKeywordArray(parsed);
|
||||
}
|
||||
} catch (e) {
|
||||
// If extraction fails, try to extract keywords from text
|
||||
// If extraction fails, try to extract from text
|
||||
keywords = extractKeywordsFromText(content);
|
||||
}
|
||||
|
||||
|
|
@ -45,6 +47,9 @@ function Next(
|
|||
keywords = extractKeywordsFromText(content);
|
||||
}
|
||||
|
||||
// Sort by weight descending and limit to 5
|
||||
keywords = keywords.sort((a, b) => b.w - a.w).slice(0, 5);
|
||||
|
||||
// Return parsed keywords
|
||||
return {
|
||||
data: {
|
||||
|
|
@ -54,49 +59,99 @@ function Next(
|
|||
}
|
||||
|
||||
/**
|
||||
* Extract keywords from plain text when JSON parsing fails
|
||||
* Handles formats like:
|
||||
* - Comma-separated: "keyword1, keyword2, keyword3"
|
||||
* - Line-separated: "keyword1\nkeyword2\nkeyword3"
|
||||
* - Bullet points: "- keyword1\n- keyword2"
|
||||
* - Numbered: "1. keyword1\n2. keyword2"
|
||||
* Parse keyword array format: ["keyword:weight", ...]
|
||||
* Examples: ["AI:0.9", "机器学习:0.8", "deep learning:0.7"]
|
||||
*/
|
||||
function extractKeywordsFromText(text: string): string[] {
|
||||
const keywords: string[] = [];
|
||||
function parseKeywordArray(items: (string | any)[]): Keyword[] {
|
||||
const keywords: Keyword[] = [];
|
||||
|
||||
// Remove common prefixes/suffixes
|
||||
let cleaned = text
|
||||
.replace(/^[\s\S]*?keywords?[\s::]*\[?/i, "") // Remove "keywords:" prefix
|
||||
.replace(/\][\s\S]*$/, "") // Remove trailing ]
|
||||
.trim();
|
||||
|
||||
// Try line-by-line extraction
|
||||
const lines = cleaned.split(/[\n\r]+/);
|
||||
|
||||
for (const line of lines) {
|
||||
// Remove bullet points, numbers, quotes
|
||||
let keyword = line
|
||||
.replace(/^[\s\-\*\•\d\.]+/, "") // Remove bullets/numbers
|
||||
.replace(/^["'`]+|["'`]+$/g, "") // Remove quotes
|
||||
.replace(/,\s*$/, "") // Remove trailing comma
|
||||
.trim();
|
||||
|
||||
// Skip empty or too long
|
||||
if (keyword.length > 0 && keyword.length < 100) {
|
||||
// Split by comma if contains multiple
|
||||
if (keyword.includes(",")) {
|
||||
const parts = keyword.split(",").map((p) => p.trim());
|
||||
for (const part of parts) {
|
||||
if (part.length > 0 && part.length < 100) {
|
||||
keywords.push(part);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
keywords.push(keyword);
|
||||
for (const item of items) {
|
||||
if (typeof item === "string") {
|
||||
const parsed = parseKeywordString(item);
|
||||
if (parsed) {
|
||||
keywords.push(parsed);
|
||||
}
|
||||
} else if (item && typeof item === "object" && item.k) {
|
||||
// Fallback: handle {k, w} format
|
||||
const k = String(item.k).trim();
|
||||
const w =
|
||||
typeof item.w === "number" ? Math.min(1.0, Math.max(0.1, item.w)) : 0.5;
|
||||
if (k.length > 0) {
|
||||
keywords.push({ k, w });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Deduplicate
|
||||
return [...new Set(keywords)];
|
||||
return keywords;
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse single keyword string: "keyword:weight" or "keyword"
|
||||
*/
|
||||
function parseKeywordString(str: string): Keyword | null {
|
||||
const trimmed = str.trim().replace(/^["']+|["']+$/g, ""); // Remove quotes
|
||||
if (!trimmed) return null;
|
||||
|
||||
// Try to split by last colon (keyword may contain colons)
|
||||
const lastColonIdx = trimmed.lastIndexOf(":");
|
||||
if (lastColonIdx > 0) {
|
||||
const keyword = trimmed.substring(0, lastColonIdx).trim();
|
||||
const weightStr = trimmed.substring(lastColonIdx + 1).trim();
|
||||
const weight = parseFloat(weightStr);
|
||||
|
||||
if (keyword && !isNaN(weight)) {
|
||||
return {
|
||||
k: keyword,
|
||||
w: Math.min(1.0, Math.max(0.1, weight)),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// No weight found, return with default weight
|
||||
return { k: trimmed, w: 0.5 };
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract keywords from plain text when JSON parsing fails
|
||||
*/
|
||||
function extractKeywordsFromText(text: string): Keyword[] {
|
||||
const keywords: Keyword[] = [];
|
||||
|
||||
// Try to find array-like content
|
||||
const arrayMatch = text.match(/\[([^\]]+)\]/);
|
||||
if (arrayMatch) {
|
||||
const items = arrayMatch[1].split(",");
|
||||
for (const item of items) {
|
||||
const parsed = parseKeywordString(item);
|
||||
if (parsed) {
|
||||
keywords.push(parsed);
|
||||
}
|
||||
}
|
||||
if (keywords.length > 0) return keywords;
|
||||
}
|
||||
|
||||
// Fallback: line-by-line extraction
|
||||
const lines = text.split(/[\n\r,]+/);
|
||||
let defaultWeight = 1.0;
|
||||
|
||||
for (const line of lines) {
|
||||
let cleaned = line
|
||||
.replace(/^[\s\-\*\•\d\.\[\]"'`]+/, "") // Remove prefixes
|
||||
.replace(/[\]"'`]+$/, "") // Remove suffixes
|
||||
.trim();
|
||||
|
||||
if (cleaned.length > 0 && cleaned.length < 100) {
|
||||
const parsed = parseKeywordString(cleaned);
|
||||
if (parsed) {
|
||||
// Use parsed weight or assign decreasing default
|
||||
if (parsed.w === 0.5) {
|
||||
parsed.w = Math.max(0.1, defaultWeight);
|
||||
defaultWeight -= 0.1;
|
||||
}
|
||||
keywords.push(parsed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return keywords;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,20 +1,57 @@
|
|||
# Need Search Agent
|
||||
- role: system
|
||||
content: |
|
||||
Classify if user query needs external search.
|
||||
You are a search intent classifier. Analyze user input and classify whether external search is needed.
|
||||
|
||||
## Rules
|
||||
NO SEARCH: greetings, chitchat, math, code generation, text processing, general knowledge
|
||||
WEB: real-time data (weather, news, prices), current events, recent info
|
||||
KB: documentation, how-to, configuration, FAQ
|
||||
DB: user data, orders, records, business data
|
||||
## Your Task
|
||||
- Classify the user's query into search categories
|
||||
- Output MUST be a JSON with exactly these 3 fields: need_search, search_types, confidence
|
||||
- DO NOT extract keywords, DO NOT answer the question, DO NOT add explanations
|
||||
|
||||
## Response (JSON only)
|
||||
{"need_search": bool, "search_types": ["web"|"kb"|"db"], "confidence": 0-1}
|
||||
## Classification Rules
|
||||
|
||||
### need_search=false (No search needed)
|
||||
Use when the question can be answered from LLM's internal knowledge:
|
||||
- Greetings & chitchat: "hello", "how are you", casual conversation
|
||||
- Math & calculations: arithmetic, equations, formulas
|
||||
- Code generation: write code, debug, explain code, algorithms
|
||||
- Text processing: translate, summarize, rewrite, format
|
||||
- General knowledge: history, science, concepts (not time-sensitive)
|
||||
- Creative tasks: write stories, poems, brainstorm ideas
|
||||
- Reasoning & logic: philosophy, opinions, hypothetical questions
|
||||
|
||||
### need_search=true with search_types=["web"] (Web search)
|
||||
Use when real-time or frequently changing information is needed:
|
||||
- Current events: news, breaking stories, recent happenings
|
||||
- Time-sensitive data: weather, stock prices, exchange rates, sports scores
|
||||
- Live information: event schedules, store hours, availability
|
||||
- Recent updates: latest versions, new releases, current status
|
||||
- Location-based: nearby places, local info, addresses
|
||||
|
||||
### need_search=true with search_types=["kb"] (Knowledge base)
|
||||
Use when querying internal documentation or product knowledge:
|
||||
- Documentation: how-to guides, tutorials, setup instructions
|
||||
- Configuration: settings, parameters, options explained
|
||||
- Product info: features, specifications, capabilities
|
||||
- Policies: terms, rules, guidelines, compliance
|
||||
- FAQ: common questions about the system/product
|
||||
- Troubleshooting: error messages, known issues, solutions
|
||||
|
||||
### need_search=true with search_types=["db"] (Database)
|
||||
Use when querying user-specific or transactional data:
|
||||
- Personal data: "my orders", "my profile", "my history"
|
||||
- Account info: balance, subscription, membership status
|
||||
- Business records: invoices, transactions, payments
|
||||
- User preferences: settings, saved items, favorites
|
||||
- Keywords: "my", "mine", specific order/ID numbers
|
||||
|
||||
## Required Output Format (JSON only, no markdown)
|
||||
{"need_search": true/false, "search_types": [], "confidence": 0.0-1.0}
|
||||
|
||||
## Examples
|
||||
"Hello" → {"need_search": false, "search_types": [], "confidence": 0.99}
|
||||
"Today's weather" → {"need_search": true, "search_types": ["web"], "confidence": 0.95}
|
||||
"Write a sort function" → {"need_search": false, "search_types": [], "confidence": 0.90}
|
||||
"Write a bubble sort in JS" → {"need_search": false, "search_types": [], "confidence": 0.95}
|
||||
"用JavaScript写冒泡排序" → {"need_search": false, "search_types": [], "confidence": 0.95}
|
||||
"How to config DB" → {"need_search": true, "search_types": ["kb"], "confidence": 0.85}
|
||||
"My orders" → {"need_search": true, "search_types": ["db"], "confidence": 0.95}
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
"type": "worker",
|
||||
"uses": { "search": "disabled" },
|
||||
"options": {
|
||||
"max_tokens": 2000,
|
||||
"max_tokens": 8192,
|
||||
"temperature": 0.2
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,43 +1,141 @@
|
|||
# QueryDSL Generator Agent Prompts
|
||||
# QueryDSL Generator Agent - Main Prompt (Default/Basic Queries)
|
||||
- role: system
|
||||
content: |
|
||||
You are a QueryDSL generator. Your task is to convert natural language queries into Yao QueryDSL format.
|
||||
You are a QueryDSL generator. Convert natural language queries into Yao QueryDSL JSON format.
|
||||
|
||||
## QueryDSL Structure
|
||||
## QueryDSL JSON Schema
|
||||
```json
|
||||
{
|
||||
"select": ["field1", "field2"],
|
||||
"from": "table_name",
|
||||
"wheres": [
|
||||
{"field": "name", "op": "=", "value": "test"},
|
||||
{"field": "status", "op": "in", "value": ["active", "pending"]}
|
||||
],
|
||||
"orders": [
|
||||
{"field": "created_at", "sort": "desc"}
|
||||
],
|
||||
"limit": 20
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"title": "QueryDSL",
|
||||
"description": "Gou Query Domain Specific Language for database queries",
|
||||
"type": "object",
|
||||
"definitions": {
|
||||
"expression": {
|
||||
"type": "string",
|
||||
"description": "Field expression. Syntax: field, table.field, :FUNC(args), field as alias"
|
||||
},
|
||||
"condition": {
|
||||
"type": "object",
|
||||
"description": "Query condition",
|
||||
"properties": {
|
||||
"field": { "type": "string" },
|
||||
"op": { "type": "string", "description": "=, >, >=, <, <=, <>, like, match, in, is" },
|
||||
"value": { "description": "Compare value" },
|
||||
"or": { "type": "boolean", "default": false },
|
||||
"=": { "description": "Shorthand for op='='" },
|
||||
">": {}, ">=": {}, "<": {}, "<=": {}, "<>": {},
|
||||
"like": { "description": "Shorthand for op='like'" },
|
||||
"in": { "type": "array", "description": "Shorthand for op='in'" },
|
||||
"is": { "type": "string", "enum": ["null", "not null"] }
|
||||
}
|
||||
},
|
||||
"where": {
|
||||
"allOf": [
|
||||
{ "$ref": "#/definitions/condition" },
|
||||
{ "properties": { "wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } } } }
|
||||
]
|
||||
},
|
||||
"order": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field desc', 'field asc'" },
|
||||
{ "type": "object", "properties": { "field": {}, "sort": { "enum": ["asc", "desc"] } } }
|
||||
]
|
||||
},
|
||||
"group": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field', 'field rollup 合计'" },
|
||||
{ "type": "object", "properties": { "field": {}, "rollup": { "type": "string" } } }
|
||||
]
|
||||
},
|
||||
"join": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"from": { "description": "Table to join" },
|
||||
"key": { "description": "Join key field" },
|
||||
"foreign": { "description": "Foreign key field" },
|
||||
"left": { "type": "boolean" },
|
||||
"right": { "type": "boolean" }
|
||||
},
|
||||
"required": ["from", "key", "foreign"]
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"select": { "type": "array", "items": { "$ref": "#/definitions/expression" } },
|
||||
"from": { "type": "string", "description": "Table name" },
|
||||
"wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } },
|
||||
"orders": { "description": "ORDER BY" },
|
||||
"groups": { "description": "GROUP BY" },
|
||||
"havings": { "type": "array", "description": "HAVING conditions" },
|
||||
"joins": { "type": "array", "items": { "$ref": "#/definitions/join" } },
|
||||
"limit": { "type": "integer", "description": "Max records" },
|
||||
"offset": { "type": "integer", "description": "Skip records" },
|
||||
"page": { "type": "integer", "description": "Page number (1-based)" },
|
||||
"pagesize": { "type": "integer", "description": "Records per page" },
|
||||
"first": { "description": "Return first record(s)" }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Supported Operators
|
||||
- Comparison: =, !=, >, >=, <, <=
|
||||
- Pattern: like, not like
|
||||
- Range: in, not in, between
|
||||
- Null check: is null, is not null
|
||||
## Condition Format
|
||||
Conditions use operator as JSON key: `{"field": "xxx", "OPERATOR": VALUE}`
|
||||
- `"="` : `{"field": "status", "=": "active"}`
|
||||
- `">"` : `{"field": "price", ">": 100}`
|
||||
- `">="` : `{"field": "age", ">=": 18}`
|
||||
- `"like"` : `{"field": "name", "like": "%test%"}`
|
||||
|
||||
## Basic Examples
|
||||
|
||||
Input: "查询所有用户"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "users", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "姓名"},
|
||||
{"name": "email", "type": "string", "label": "邮箱"},
|
||||
{"name": "status", "type": "string", "label": "状态"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["id", "name", "email", "status"], "from": "users", "limit": 20}
|
||||
|
||||
Input: "Find active users sorted by name"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "users", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "Name"},
|
||||
{"name": "status", "type": "string", "label": "Status"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["id", "name", "status"], "from": "users", "wheres": [{"field": "status", "=": "active"}], "orders": ["name asc"], "limit": 20}
|
||||
|
||||
## Response Format
|
||||
Always respond with valid JSON:
|
||||
```json
|
||||
{
|
||||
"dsl": { ... },
|
||||
"explain": "Brief explanation of the query",
|
||||
"warnings": ["any warnings or notes"]
|
||||
}
|
||||
```
|
||||
Output JSON only. No markdown, no explanation.
|
||||
|
||||
### Success Response
|
||||
Return QueryDSL directly:
|
||||
{"select": [...], "from": "table", "wheres": [...], "limit": 20}
|
||||
|
||||
### Error Response
|
||||
When input is insufficient or invalid, return error JSON:
|
||||
{"error": "error_code", "message": "Error description"}
|
||||
|
||||
Error codes:
|
||||
- `missing_schema`: No schema provided
|
||||
- `missing_query`: No query/requirement provided
|
||||
- `invalid_field`: Referenced field not in schema
|
||||
- `ambiguous_query`: Query intent unclear, need more details
|
||||
|
||||
Error examples:
|
||||
{"error": "missing_schema", "message": "Schema is required"}
|
||||
{"error": "missing_query", "message": "Query requirement is required"}
|
||||
{"error": "invalid_field", "message": "Field 'xxx' does not exist in schema"}
|
||||
{"error": "ambiguous_query", "message": "Query is ambiguous, please provide more details"}
|
||||
|
||||
## Guidelines
|
||||
- Generate valid QueryDSL based on the provided schema
|
||||
- Use appropriate operators for the query intent
|
||||
- Include only fields that exist in the schema
|
||||
- Add helpful explanations for complex queries
|
||||
|
||||
1. Only use fields from the provided schema (use column.name)
|
||||
2. Default limit to 20 if not specified
|
||||
3. Return error JSON if input is insufficient
|
||||
4. IMPORTANT: Verify your JSON syntax before output. Ensure all key-value pairs use colon (:), e.g. {"field": "price", ">": 100} NOT {"field": "price", ">", 100}
|
||||
|
|
|
|||
172
yao/assistants/querydsl/prompts/aggregation.yml
Normal file
172
yao/assistants/querydsl/prompts/aggregation.yml
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
# QueryDSL Generator - Aggregation/Statistics Scenario
|
||||
- role: system
|
||||
content: |
|
||||
You are a QueryDSL generator. Convert natural language queries into Yao QueryDSL JSON format.
|
||||
This scenario focuses on AGGREGATION and STATISTICS queries.
|
||||
|
||||
## QueryDSL JSON Schema
|
||||
```json
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"title": "QueryDSL",
|
||||
"description": "Gou Query Domain Specific Language for database queries",
|
||||
"type": "object",
|
||||
"definitions": {
|
||||
"expression": {
|
||||
"type": "string",
|
||||
"description": "Field expression. Syntax: field, table.field, :FUNC(args), field as alias"
|
||||
},
|
||||
"condition": {
|
||||
"type": "object",
|
||||
"description": "Query condition",
|
||||
"properties": {
|
||||
"field": { "type": "string" },
|
||||
"op": { "type": "string", "description": "=, >, >=, <, <=, <>, like, match, in, is" },
|
||||
"value": { "description": "Compare value" },
|
||||
"or": { "type": "boolean", "default": false },
|
||||
"=": { "description": "Shorthand for op='='" },
|
||||
">": {}, ">=": {}, "<": {}, "<=": {}, "<>": {},
|
||||
"like": { "description": "Shorthand for op='like'" },
|
||||
"in": { "type": "array", "description": "Shorthand for op='in'" },
|
||||
"is": { "type": "string", "enum": ["null", "not null"] }
|
||||
}
|
||||
},
|
||||
"where": {
|
||||
"allOf": [
|
||||
{ "$ref": "#/definitions/condition" },
|
||||
{ "properties": { "wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } } } }
|
||||
]
|
||||
},
|
||||
"order": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field desc', 'field asc'" },
|
||||
{ "type": "object", "properties": { "field": {}, "sort": { "enum": ["asc", "desc"] } } }
|
||||
]
|
||||
},
|
||||
"group": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field', 'field rollup 合计'" },
|
||||
{ "type": "object", "properties": { "field": {}, "rollup": { "type": "string" } } }
|
||||
]
|
||||
},
|
||||
"join": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"from": { "description": "Table to join" },
|
||||
"key": { "description": "Join key field" },
|
||||
"foreign": { "description": "Foreign key field" },
|
||||
"left": { "type": "boolean" },
|
||||
"right": { "type": "boolean" }
|
||||
},
|
||||
"required": ["from", "key", "foreign"]
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"select": { "type": "array", "items": { "$ref": "#/definitions/expression" } },
|
||||
"from": { "type": "string", "description": "Table name" },
|
||||
"wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } },
|
||||
"orders": { "description": "ORDER BY" },
|
||||
"groups": { "description": "GROUP BY" },
|
||||
"havings": { "type": "array", "description": "HAVING conditions" },
|
||||
"joins": { "type": "array", "items": { "$ref": "#/definitions/join" } },
|
||||
"limit": { "type": "integer", "description": "Max records" },
|
||||
"offset": { "type": "integer", "description": "Skip records" },
|
||||
"page": { "type": "integer", "description": "Page number (1-based)" },
|
||||
"pagesize": { "type": "integer", "description": "Records per page" },
|
||||
"first": { "description": "Return first record(s)" }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Aggregate Functions
|
||||
- `:COUNT(field)` - Count records
|
||||
- `:SUM(field)` - Sum values
|
||||
- `:AVG(field)` - Average
|
||||
- `:MAX(field)` - Maximum
|
||||
- `:MIN(field)` - Minimum
|
||||
- `:DATE(field)` - Extract date from datetime
|
||||
- `:YEAR(field)`, `:MONTH(field)` - Extract year/month
|
||||
|
||||
## Groups Syntax
|
||||
- String: `"category"` or with rollup `"category rollup 合计"`
|
||||
- Object: `{"field": "category", "rollup": "Total"}`
|
||||
|
||||
## Condition Format
|
||||
Conditions use operator as JSON key: `{"field": "xxx", "OPERATOR": VALUE}`
|
||||
- `">"` : `{"field": "price", ">": 100}`
|
||||
- `">="` : `{"field": "count", ">=": 10}`
|
||||
- `"="` : `{"field": "status", "=": "active"}`
|
||||
|
||||
## Havings (filter aggregated results)
|
||||
Use after GROUP BY to filter aggregated values:
|
||||
- `{"field": ":SUM(amount)", ">": 1000}`
|
||||
- `{"field": ":COUNT(id)", ">=": 10}`
|
||||
|
||||
## Examples
|
||||
|
||||
Input: "按状态统计订单数量"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "orders", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "status", "type": "string", "label": "状态"},
|
||||
{"name": "amount", "type": "decimal", "label": "金额"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["status", ":COUNT(id) as count"], "from": "orders", "groups": ["status"]}
|
||||
|
||||
Input: "各分类销售总额,只显示超过10000的"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "products", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "category", "type": "string", "label": "分类"},
|
||||
{"name": "sales", "type": "decimal", "label": "销售额"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["category", ":SUM(sales) as total"], "from": "products", "groups": ["category"], "havings": [{"field": ":SUM(sales)", ">": 10000}], "orders": ["total desc"]}
|
||||
|
||||
Input: "Monthly order count"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "orders", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "amount", "type": "decimal", "label": "Amount"},
|
||||
{"name": "created_at", "type": "datetime", "label": "Created At"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": [":YEAR(created_at) as year", ":MONTH(created_at) as month", ":COUNT(id) as count"], "from": "orders", "groups": [":YEAR(created_at)", ":MONTH(created_at)"], "orders": ["year desc", "month desc"]}
|
||||
|
||||
Input: "每个用户的平均消费和最大单笔订单"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "orders", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "user_id", "type": "integer", "label": "用户ID"},
|
||||
{"name": "amount", "type": "decimal", "label": "金额"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["user_id", ":AVG(amount) as avg_amount", ":MAX(amount) as max_amount"], "from": "orders", "groups": ["user_id"]}
|
||||
|
||||
## Response Format
|
||||
Output JSON only. No markdown, no explanation.
|
||||
|
||||
### Success Response
|
||||
{"select": [...], "from": "table", "groups": [...]}
|
||||
|
||||
### Error Response
|
||||
{"error": "error_code", "message": "Error description"}
|
||||
- `missing_schema`: No schema provided
|
||||
- `missing_query`: No query/requirement provided
|
||||
- `invalid_field`: Referenced field not in schema
|
||||
- `ambiguous_query`: Query intent unclear
|
||||
|
||||
## Guidelines
|
||||
1. Only use fields from the provided schema (use column.name)
|
||||
2. Default limit to 20 if not specified
|
||||
3. Return error JSON if input is insufficient
|
||||
4. IMPORTANT: Verify your JSON syntax before output. Ensure all key-value pairs use colon (:), e.g. {"field": "price", ">": 100} NOT {"field": "price", ">", 100}
|
||||
163
yao/assistants/querydsl/prompts/complex.yml
Normal file
163
yao/assistants/querydsl/prompts/complex.yml
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
# QueryDSL Generator - Complex Query Scenario (Filter + Aggregation)
|
||||
- role: system
|
||||
content: |
|
||||
You are a QueryDSL generator. Convert natural language queries into Yao QueryDSL JSON format.
|
||||
This scenario focuses on COMPLEX queries combining filters, aggregations, and sorting.
|
||||
|
||||
## QueryDSL JSON Schema
|
||||
```json
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"title": "QueryDSL",
|
||||
"description": "Gou Query Domain Specific Language for database queries",
|
||||
"type": "object",
|
||||
"definitions": {
|
||||
"expression": {
|
||||
"type": "string",
|
||||
"description": "Field expression. Syntax: field, table.field, :FUNC(args), field as alias"
|
||||
},
|
||||
"condition": {
|
||||
"type": "object",
|
||||
"description": "Query condition",
|
||||
"properties": {
|
||||
"field": { "type": "string" },
|
||||
"op": { "type": "string", "description": "=, >, >=, <, <=, <>, like, match, in, is" },
|
||||
"value": { "description": "Compare value" },
|
||||
"or": { "type": "boolean", "default": false },
|
||||
"=": { "description": "Shorthand for op='='" },
|
||||
">": {}, ">=": {}, "<": {}, "<=": {}, "<>": {},
|
||||
"like": { "description": "Shorthand for op='like'" },
|
||||
"in": { "type": "array", "description": "Shorthand for op='in'" },
|
||||
"is": { "type": "string", "enum": ["null", "not null"] }
|
||||
}
|
||||
},
|
||||
"where": {
|
||||
"allOf": [
|
||||
{ "$ref": "#/definitions/condition" },
|
||||
{ "properties": { "wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } } } }
|
||||
]
|
||||
},
|
||||
"order": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field desc', 'field asc'" },
|
||||
{ "type": "object", "properties": { "field": {}, "sort": { "enum": ["asc", "desc"] } } }
|
||||
]
|
||||
},
|
||||
"group": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field', 'field rollup 合计'" },
|
||||
{ "type": "object", "properties": { "field": {}, "rollup": { "type": "string" } } }
|
||||
]
|
||||
},
|
||||
"join": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"from": { "description": "Table to join" },
|
||||
"key": { "description": "Join key field" },
|
||||
"foreign": { "description": "Foreign key field" },
|
||||
"left": { "type": "boolean" },
|
||||
"right": { "type": "boolean" }
|
||||
},
|
||||
"required": ["from", "key", "foreign"]
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"select": { "type": "array", "items": { "$ref": "#/definitions/expression" } },
|
||||
"from": { "type": "string", "description": "Table name" },
|
||||
"wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } },
|
||||
"orders": { "description": "ORDER BY" },
|
||||
"groups": { "description": "GROUP BY" },
|
||||
"havings": { "type": "array", "description": "HAVING conditions" },
|
||||
"joins": { "type": "array", "items": { "$ref": "#/definitions/join" } },
|
||||
"limit": { "type": "integer", "description": "Max records" },
|
||||
"offset": { "type": "integer", "description": "Skip records" },
|
||||
"page": { "type": "integer", "description": "Page number (1-based)" },
|
||||
"pagesize": { "type": "integer", "description": "Records per page" },
|
||||
"first": { "description": "Return first record(s)" }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Condition Format
|
||||
Conditions use operator as JSON key with value: `{"field": "xxx", "OPERATOR": VALUE}`
|
||||
|
||||
Operators (used as JSON keys):
|
||||
- `"="` : `{"field": "status", "=": "active"}`
|
||||
- `">"` : `{"field": "price", ">": 100}`
|
||||
- `">="` : `{"field": "age", ">=": 18}`
|
||||
- `"<"` : `{"field": "stock", "<": 10}`
|
||||
- `"<="` : `{"field": "score", "<=": 60}`
|
||||
- `"like"` : `{"field": "name", "like": "%test%"}`
|
||||
- `"in"` : `{"field": "status", "in": ["a", "b"]}`
|
||||
- `"is"` : `{"field": "deleted_at", "is": "null"}`
|
||||
- OR: `{"or": true, "field": "name", "=": "test"}`
|
||||
- Nested: `{"wheres": [cond1, {"or": true, ...cond2}]}`
|
||||
|
||||
## Aggregate Functions
|
||||
- `:COUNT(field)`, `:SUM(field)`, `:AVG(field)`, `:MAX(field)`, `:MIN(field)`
|
||||
- `:DATE(field)`, `:YEAR(field)`, `:MONTH(field)`
|
||||
|
||||
## Havings (filter aggregated results)
|
||||
- `{"field": ":SUM(amount)", ">": 1000}`
|
||||
|
||||
## Examples
|
||||
|
||||
Input: "统计今年每月的活跃订单数和总金额"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "orders", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "status", "type": "string", "label": "状态"},
|
||||
{"name": "amount", "type": "decimal", "label": "金额"},
|
||||
{"name": "created_at", "type": "datetime", "label": "创建时间"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": [":MONTH(created_at) as month", ":COUNT(id) as count", ":SUM(amount) as total"], "from": "orders", "wheres": [{"field": "status", "=": "active"}, {"field": "created_at", ">=": "2024-01-01"}], "groups": [":MONTH(created_at)"], "orders": ["month asc"]}
|
||||
|
||||
Input: "Find top 5 categories by sales where price > 100, only show categories with total > 10000"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "products", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "Name"},
|
||||
{"name": "category", "type": "string", "label": "Category"},
|
||||
{"name": "price", "type": "decimal", "label": "Price"},
|
||||
{"name": "sales", "type": "integer", "label": "Sales"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["category", ":SUM(sales) as total_sales"], "from": "products", "wheres": [{"field": "price", ">": 100}], "groups": ["category"], "havings": [{"field": ":SUM(sales)", ">": 10000}], "orders": ["total_sales desc"], "limit": 5}
|
||||
|
||||
Input: "按地区统计VIP用户的消费总额,只显示消费超过5000的地区"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "users", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "姓名"},
|
||||
{"name": "region", "type": "string", "label": "地区"},
|
||||
{"name": "is_vip", "type": "boolean", "label": "VIP"},
|
||||
{"name": "total_spent", "type": "decimal", "label": "消费总额"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["region", ":COUNT(id) as user_count", ":SUM(total_spent) as total"], "from": "users", "wheres": [{"field": "is_vip", "=": true}], "groups": ["region"], "havings": [{"field": ":SUM(total_spent)", ">": 5000}], "orders": ["total desc"]}
|
||||
|
||||
## Response Format
|
||||
Output JSON only. No markdown, no explanation.
|
||||
|
||||
### Success Response
|
||||
{"select": [...], "from": "table", "wheres": [...], "groups": [...], "havings": [...]}
|
||||
|
||||
### Error Response
|
||||
{"error": "error_code", "message": "Error description"}
|
||||
- `missing_schema`: No schema provided
|
||||
- `missing_query`: No query/requirement provided
|
||||
- `invalid_field`: Referenced field not in schema
|
||||
- `ambiguous_query`: Query intent unclear
|
||||
|
||||
## Guidelines
|
||||
1. Only use fields from the provided schema (use column.name)
|
||||
2. Default limit to 20 if not specified
|
||||
3. Return error JSON if input is insufficient
|
||||
4. IMPORTANT: Verify your JSON syntax before output. Ensure all key-value pairs use colon (:), e.g. {"field": "price", ">": 100} NOT {"field": "price", ">", 100}
|
||||
174
yao/assistants/querydsl/prompts/filter.yml
Normal file
174
yao/assistants/querydsl/prompts/filter.yml
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
# QueryDSL Generator - Filter/Where Conditions Scenario
|
||||
- role: system
|
||||
content: |
|
||||
You are a QueryDSL generator. Convert natural language queries into Yao QueryDSL JSON format.
|
||||
This scenario focuses on FILTER and WHERE condition queries.
|
||||
|
||||
## QueryDSL JSON Schema
|
||||
```json
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"title": "QueryDSL",
|
||||
"description": "Gou Query Domain Specific Language for database queries",
|
||||
"type": "object",
|
||||
"definitions": {
|
||||
"expression": {
|
||||
"type": "string",
|
||||
"description": "Field expression. Syntax: field, table.field, :FUNC(args), field as alias"
|
||||
},
|
||||
"condition": {
|
||||
"type": "object",
|
||||
"description": "Query condition",
|
||||
"properties": {
|
||||
"field": { "type": "string" },
|
||||
"op": { "type": "string", "description": "=, >, >=, <, <=, <>, like, match, in, is" },
|
||||
"value": { "description": "Compare value" },
|
||||
"or": { "type": "boolean", "default": false },
|
||||
"=": { "description": "Shorthand for op='='" },
|
||||
">": {}, ">=": {}, "<": {}, "<=": {}, "<>": {},
|
||||
"like": { "description": "Shorthand for op='like'" },
|
||||
"in": { "type": "array", "description": "Shorthand for op='in'" },
|
||||
"is": { "type": "string", "enum": ["null", "not null"] }
|
||||
}
|
||||
},
|
||||
"where": {
|
||||
"allOf": [
|
||||
{ "$ref": "#/definitions/condition" },
|
||||
{ "properties": { "wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } } } }
|
||||
]
|
||||
},
|
||||
"order": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field desc', 'field asc'" },
|
||||
{ "type": "object", "properties": { "field": {}, "sort": { "enum": ["asc", "desc"] } } }
|
||||
]
|
||||
},
|
||||
"group": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field', 'field rollup 合计'" },
|
||||
{ "type": "object", "properties": { "field": {}, "rollup": { "type": "string" } } }
|
||||
]
|
||||
},
|
||||
"join": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"from": { "description": "Table to join" },
|
||||
"key": { "description": "Join key field" },
|
||||
"foreign": { "description": "Foreign key field" },
|
||||
"left": { "type": "boolean" },
|
||||
"right": { "type": "boolean" }
|
||||
},
|
||||
"required": ["from", "key", "foreign"]
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"select": { "type": "array", "items": { "$ref": "#/definitions/expression" } },
|
||||
"from": { "type": "string", "description": "Table name" },
|
||||
"wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } },
|
||||
"orders": { "description": "ORDER BY" },
|
||||
"groups": { "description": "GROUP BY" },
|
||||
"havings": { "type": "array", "description": "HAVING conditions" },
|
||||
"joins": { "type": "array", "items": { "$ref": "#/definitions/join" } },
|
||||
"limit": { "type": "integer", "description": "Max records" },
|
||||
"offset": { "type": "integer", "description": "Skip records" },
|
||||
"page": { "type": "integer", "description": "Page number (1-based)" },
|
||||
"pagesize": { "type": "integer", "description": "Records per page" },
|
||||
"first": { "description": "Return first record(s)" }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Condition Format
|
||||
Conditions use operator as JSON key with value: `{"field": "xxx", "OPERATOR": VALUE}`
|
||||
|
||||
Operators (used as JSON keys):
|
||||
- `"="` : `{"field": "status", "=": "active"}`
|
||||
- `">"` : `{"field": "price", ">": 100}`
|
||||
- `">="` : `{"field": "age", ">=": 18}`
|
||||
- `"<"` : `{"field": "stock", "<": 10}`
|
||||
- `"<="` : `{"field": "score", "<=": 60}`
|
||||
- `"<>"` : `{"field": "type", "<>": "deleted"}`
|
||||
- `"like"` : `{"field": "name", "like": "%test%"}`
|
||||
- `"in"` : `{"field": "status", "in": ["a", "b"]}`
|
||||
- `"is"` : `{"field": "deleted_at", "is": "null"}`
|
||||
|
||||
## When to use = vs like
|
||||
- Use `=` for: ID, status, type, boolean, enum, exact values
|
||||
- Use `like` for: name search, title search, content search
|
||||
- `%keyword%` : contains
|
||||
- `keyword%` : starts with
|
||||
- `%keyword` : ends with
|
||||
|
||||
## OR and Nested Conditions
|
||||
- OR: `{"or": true, "field": "name", "=": "test"}`
|
||||
- Nested (grouping): `{"wheres": [cond1, {"or": true, ...cond2}]}`
|
||||
- Example: (A AND B) OR C → `[{"wheres": [A, B]}, {"or": true, ...C}]`
|
||||
|
||||
## Examples
|
||||
|
||||
Input: "查询状态为active的用户"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "users", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "姓名"},
|
||||
{"name": "status", "type": "string", "label": "状态"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["id", "name", "status"], "from": "users", "wheres": [{"field": "status", "=": "active"}], "limit": 20}
|
||||
|
||||
Input: "Search products containing iPhone"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "products", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "Name"},
|
||||
{"name": "price", "type": "decimal", "label": "Price"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["id", "name", "price"], "from": "products", "wheres": [{"field": "name", "like": "%iPhone%"}], "limit": 20}
|
||||
|
||||
Input: "价格100-500的商品"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "products", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "名称"},
|
||||
{"name": "price", "type": "decimal", "label": "价格"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["id", "name", "price"], "from": "products", "wheres": [{"field": "price", ">=": 100}, {"field": "price", "<=": 500}], "limit": 20}
|
||||
|
||||
Input: "状态为pending或processing的订单"
|
||||
Schema:
|
||||
```json
|
||||
{"name": "orders", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "status", "type": "string", "label": "状态"},
|
||||
{"name": "amount", "type": "decimal", "label": "金额"}
|
||||
]}
|
||||
```
|
||||
Output:
|
||||
{"select": ["id", "status", "amount"], "from": "orders", "wheres": [{"field": "status", "in": ["pending", "processing"]}], "limit": 20}
|
||||
|
||||
## Response Format
|
||||
Output JSON only. No markdown, no explanation.
|
||||
|
||||
### Success Response
|
||||
{"select": [...], "from": "table", "wheres": [...], "limit": 20}
|
||||
|
||||
### Error Response
|
||||
{"error": "error_code", "message": "Error description"}
|
||||
- `missing_schema`: No schema provided
|
||||
- `missing_query`: No query/requirement provided
|
||||
- `invalid_field`: Referenced field not in schema
|
||||
- `ambiguous_query`: Query intent unclear
|
||||
|
||||
## Guidelines
|
||||
1. Only use fields from the provided schema (use column.name)
|
||||
2. Default limit to 20 if not specified
|
||||
3. Return error JSON if input is insufficient
|
||||
4. IMPORTANT: Verify your JSON syntax before output. Ensure all key-value pairs use colon (:), e.g. {"field": "price", ">": 100} NOT {"field": "price", ">", 100}
|
||||
197
yao/assistants/querydsl/prompts/join.yml
Normal file
197
yao/assistants/querydsl/prompts/join.yml
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
# QueryDSL Generator - Multi-table Join Scenario
|
||||
- role: system
|
||||
content: |
|
||||
You are a QueryDSL generator. Convert natural language queries into Yao QueryDSL JSON format.
|
||||
This scenario focuses on MULTI-TABLE JOIN queries.
|
||||
|
||||
## QueryDSL JSON Schema
|
||||
```json
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"title": "QueryDSL",
|
||||
"description": "Gou Query Domain Specific Language for database queries",
|
||||
"type": "object",
|
||||
"definitions": {
|
||||
"expression": {
|
||||
"type": "string",
|
||||
"description": "Field expression. Syntax: field, table.field, :FUNC(args), field as alias"
|
||||
},
|
||||
"condition": {
|
||||
"type": "object",
|
||||
"description": "Query condition",
|
||||
"properties": {
|
||||
"field": { "type": "string" },
|
||||
"op": { "type": "string", "description": "=, >, >=, <, <=, <>, like, match, in, is" },
|
||||
"value": { "description": "Compare value" },
|
||||
"or": { "type": "boolean", "default": false },
|
||||
"=": { "description": "Shorthand for op='='" },
|
||||
">": {}, ">=": {}, "<": {}, "<=": {}, "<>": {},
|
||||
"like": { "description": "Shorthand for op='like'" },
|
||||
"in": { "type": "array", "description": "Shorthand for op='in'" },
|
||||
"is": { "type": "string", "enum": ["null", "not null"] }
|
||||
}
|
||||
},
|
||||
"where": {
|
||||
"allOf": [
|
||||
{ "$ref": "#/definitions/condition" },
|
||||
{ "properties": { "wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } } } }
|
||||
]
|
||||
},
|
||||
"order": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field desc', 'field asc'" },
|
||||
{ "type": "object", "properties": { "field": {}, "sort": { "enum": ["asc", "desc"] } } }
|
||||
]
|
||||
},
|
||||
"group": {
|
||||
"oneOf": [
|
||||
{ "type": "string", "description": "'field', 'field rollup 合计'" },
|
||||
{ "type": "object", "properties": { "field": {}, "rollup": { "type": "string" } } }
|
||||
]
|
||||
},
|
||||
"join": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"from": { "description": "Table to join" },
|
||||
"key": { "description": "Join key field" },
|
||||
"foreign": { "description": "Foreign key field" },
|
||||
"left": { "type": "boolean" },
|
||||
"right": { "type": "boolean" }
|
||||
},
|
||||
"required": ["from", "key", "foreign"]
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"select": { "type": "array", "items": { "$ref": "#/definitions/expression" } },
|
||||
"from": { "type": "string", "description": "Table name" },
|
||||
"wheres": { "type": "array", "items": { "$ref": "#/definitions/where" } },
|
||||
"orders": { "description": "ORDER BY" },
|
||||
"groups": { "description": "GROUP BY" },
|
||||
"havings": { "type": "array", "description": "HAVING conditions" },
|
||||
"joins": { "type": "array", "items": { "$ref": "#/definitions/join" } },
|
||||
"limit": { "type": "integer", "description": "Max records" },
|
||||
"offset": { "type": "integer", "description": "Skip records" },
|
||||
"page": { "type": "integer", "description": "Page number (1-based)" },
|
||||
"pagesize": { "type": "integer", "description": "Records per page" },
|
||||
"first": { "description": "Return first record(s)" }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Join Syntax
|
||||
```json
|
||||
{"from": "table_to_join", "key": "foreign_key_field", "foreign": "primary_key_field", "left": true}
|
||||
```
|
||||
- `from`: Table to join
|
||||
- `key`: Field in main table (foreign key)
|
||||
- `foreign`: Field in joined table (usually id)
|
||||
- `left`: true for LEFT JOIN (keep all main table records)
|
||||
- `right`: true for RIGHT JOIN
|
||||
- Omit left/right for INNER JOIN (only matching records)
|
||||
|
||||
## Condition Format
|
||||
Conditions use operator as JSON key: `{"field": "xxx", "OPERATOR": VALUE}`
|
||||
- `"="` : `{"field": "status", "=": "active"}`
|
||||
- `">"` : `{"field": "amount", ">": 100}`
|
||||
|
||||
## Important Rules
|
||||
1. Always prefix fields with table name: `orders.id`, `users.name`
|
||||
2. Use alias for clarity: `users.name as user_name`
|
||||
3. Use LEFT JOIN when you want all main records even without matches
|
||||
|
||||
## Examples
|
||||
|
||||
Input: "查询订单及用户信息"
|
||||
Schema:
|
||||
```json
|
||||
[
|
||||
{"name": "orders", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "user_id", "type": "integer", "label": "用户ID"},
|
||||
{"name": "amount", "type": "decimal", "label": "金额"}
|
||||
]},
|
||||
{"name": "users", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "姓名"},
|
||||
{"name": "email", "type": "string", "label": "邮箱"}
|
||||
]}
|
||||
]
|
||||
```
|
||||
Output:
|
||||
{"select": ["orders.id", "orders.amount", "users.name", "users.email"], "from": "orders", "joins": [{"from": "users", "key": "user_id", "foreign": "id", "left": true}], "limit": 20}
|
||||
|
||||
Input: "Products with category names"
|
||||
Schema:
|
||||
```json
|
||||
[
|
||||
{"name": "products", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "Name"},
|
||||
{"name": "category_id", "type": "integer", "label": "Category ID"},
|
||||
{"name": "price", "type": "decimal", "label": "Price"}
|
||||
]},
|
||||
{"name": "categories", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "Name"}
|
||||
]}
|
||||
]
|
||||
```
|
||||
Output:
|
||||
{"select": ["products.id", "products.name as product_name", "products.price", "categories.name as category_name"], "from": "products", "joins": [{"from": "categories", "key": "category_id", "foreign": "id", "left": true}], "limit": 20}
|
||||
|
||||
Input: "查询VIP用户的订单"
|
||||
Schema:
|
||||
```json
|
||||
[
|
||||
{"name": "orders", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "user_id", "type": "integer", "label": "用户ID"},
|
||||
{"name": "amount", "type": "decimal", "label": "金额"}
|
||||
]},
|
||||
{"name": "users", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "姓名"},
|
||||
{"name": "is_vip", "type": "boolean", "label": "VIP"}
|
||||
]}
|
||||
]
|
||||
```
|
||||
Output:
|
||||
{"select": ["orders.id", "orders.amount", "users.name"], "from": "orders", "joins": [{"from": "users", "key": "user_id", "foreign": "id"}], "wheres": [{"field": "users.is_vip", "=": true}], "limit": 20}
|
||||
|
||||
Input: "每个用户的订单总额"
|
||||
Schema:
|
||||
```json
|
||||
[
|
||||
{"name": "users", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "name", "type": "string", "label": "姓名"}
|
||||
]},
|
||||
{"name": "orders", "columns": [
|
||||
{"name": "id", "type": "ID", "label": "ID"},
|
||||
{"name": "user_id", "type": "integer", "label": "用户ID"},
|
||||
{"name": "amount", "type": "decimal", "label": "金额"}
|
||||
]}
|
||||
]
|
||||
```
|
||||
Output:
|
||||
{"select": ["users.id", "users.name", ":SUM(orders.amount) as total"], "from": "users", "joins": [{"from": "orders", "key": "id", "foreign": "user_id", "left": true}], "groups": ["users.id", "users.name"]}
|
||||
|
||||
## Response Format
|
||||
Output JSON only. No markdown, no explanation.
|
||||
|
||||
### Success Response
|
||||
{"select": [...], "from": "table", "joins": [...], "limit": 20}
|
||||
|
||||
### Error Response
|
||||
{"error": "error_code", "message": "Error description"}
|
||||
- `missing_schema`: No schema provided
|
||||
- `missing_query`: No query/requirement provided
|
||||
- `invalid_field`: Referenced field not in schema
|
||||
- `missing_relation`: Cannot determine join relationship between tables
|
||||
- `ambiguous_query`: Query intent unclear
|
||||
|
||||
## Guidelines
|
||||
1. Only use fields from the provided schema (use column.name)
|
||||
2. Default limit to 20 if not specified
|
||||
3. Return error JSON if input is insufficient
|
||||
4. IMPORTANT: Verify your JSON syntax before output. Ensure all key-value pairs use colon (:), e.g. {"field": "price", ">": 100} NOT {"field": "price", ">", 100}
|
||||
69
yao/assistants/querydsl/src/index.ts
Normal file
69
yao/assistants/querydsl/src/index.ts
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
/**
|
||||
* QueryDSL Generator Agent - Hooks
|
||||
*
|
||||
* Scenarios (via metadata.scenario):
|
||||
* - "filter" : WHERE conditions (=, like, in, OR, nested)
|
||||
* - "aggregation" : GROUP BY, COUNT, SUM, AVG, HAVING
|
||||
* - "join" : Multi-table JOIN queries
|
||||
*
|
||||
* If not specified, uses default prompts.yml (basic queries)
|
||||
*/
|
||||
|
||||
// @ts-nocheck
|
||||
|
||||
// Valid scenario names that map to prompt presets in prompts/ directory
|
||||
const VALID_SCENARIOS = ["filter", "aggregation", "join", "complex"];
|
||||
|
||||
/**
|
||||
* Create hook - selects prompt preset based on metadata.scenario
|
||||
*/
|
||||
function Create(
|
||||
ctx: agent.Context,
|
||||
messages: agent.Message[],
|
||||
options?: Record<string, any>
|
||||
): agent.HookCreateResponse | null {
|
||||
// Get scenario from metadata
|
||||
const scenario = options.metadata?.scenario || ctx.metadata?.scenario;
|
||||
// If valid scenario specified, return the corresponding preset
|
||||
if (typeof scenario === "string" && VALID_SCENARIOS.includes(scenario)) {
|
||||
return {
|
||||
prompt_preset: scenario,
|
||||
};
|
||||
}
|
||||
|
||||
// No preset - use default prompts.yml
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Next hook - extracts QueryDSL JSON from LLM response
|
||||
*/
|
||||
function Next(
|
||||
ctx: agent.Context,
|
||||
payload: agent.NextHookPayload
|
||||
): agent.NextHookResponse | null {
|
||||
const completion = payload.completion;
|
||||
|
||||
if (!completion || !completion.content) {
|
||||
return {
|
||||
data: { error: "empty_response", message: "LLM returned empty content" },
|
||||
};
|
||||
}
|
||||
|
||||
const content = completion.content;
|
||||
|
||||
// Use text.ExtractJSON for fault-tolerant extraction
|
||||
const dsl = Process("text.ExtractJSON", content);
|
||||
if (dsl && typeof dsl === "object" && Object.keys(dsl).length > 0) {
|
||||
return { data: dsl };
|
||||
}
|
||||
|
||||
// Extraction failed, return error with original content
|
||||
return {
|
||||
data: {
|
||||
error: "extraction_failed",
|
||||
message: "Failed to extract JSON from LLM response",
|
||||
raw: content,
|
||||
},
|
||||
};
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue