Implement Retry Mechanism and Lint Validation for QueryDSL Generation

- Added retry logic to the Generate method in both AgentProvider and MCPProvider to handle failures in QueryDSL generation.
- Integrated lint validation to ensure generated QueryDSL meets required standards, with detailed error reporting for invalid DSL.
- Enhanced test coverage by introducing new tests for retry behavior in both agent and MCP contexts, ensuring robustness against lint failures.
- Updated documentation to reflect changes in QueryDSL generation processes and error handling mechanisms.

agent/search/nlp/querydsl/agent.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"github.com/yaoapp/gou/query/gou"
+	"github.com/yaoapp/gou/query/linter"
 	"github.com/yaoapp/yao/agent/caller"
 	agentContext "github.com/yaoapp/yao/agent/context"
 )
@@ -22,7 +23,7 @@ func NewAgentProvider(agentID string) *AgentProvider {
 	}
 }
 
-// Generate generates QueryDSL by calling the target agent
+// Generate generates QueryDSL by calling the target agent with retry and lint validation
 // The agent receives the query and schema, returns generated QueryDSL
 func (p *AgentProvider) Generate(ctx *agentContext.Context, input *Input) (*Result, error) {
 	if ctx == nil {
@@ -40,8 +41,70 @@ func (p *AgentProvider) Generate(ctx *agentContext.Context, input *Input) (*Resu
 		return nil, fmt.Errorf("failed to get agent %s: %w", p.agentID, err)
 	}
 
-	// Build the request message
-	// Note: Agent will load model metadata internally based on model IDs
+	var lastError error
+	var lastLintErrors string
+
+	for attempt := 1; attempt <= MaxRetries; attempt++ {
+		// Build the request message
+		requestData := p.buildRequestData(input, attempt, lastLintErrors)
+		requestJSON, _ := json.Marshal(requestData)
+
+		// Create message for the agent
+		messages := []agentContext.Message{
+			{
+				Role:    "user",
+				Content: string(requestJSON),
+			},
+		}
+
+		// Call the agent with skip options (no history, no output)
+		options := &agentContext.Options{
+			Skip: &agentContext.Skip{
+				History: true,
+				Output:  true,
+			},
+		}
+
+		result, err := agent.Stream(ctx, messages, options)
+		if err != nil {
+			lastError = fmt.Errorf("agent call failed: %w", err)
+			continue
+		}
+
+		// Parse the result
+		genResult, err := p.parseResult(result)
+		if err != nil {
+			lastError = err
+			continue
+		}
+
+		// Validate with linter if DSL is present
+		if genResult.DSL != nil {
+			lintResult := p.validateDSL(genResult.DSL)
+			if lintResult.Valid {
+				return genResult, nil
+			}
+
+			// Lint failed, prepare error message for retry
+			lastLintErrors = lintResult.FormatDiagnostics()
+			lastError = fmt.Errorf("QueryDSL validation failed: %s", lastLintErrors)
+
+			// Add lint warnings to result warnings
+			for _, diag := range lintResult.Diagnostics {
+				genResult.Warnings = append(genResult.Warnings, fmt.Sprintf("[%s] %s: %s", diag.Code, diag.Path, diag.Message))
+			}
+			continue
+		}
+
+		// No DSL returned
+		lastError = fmt.Errorf("no QueryDSL returned from agent")
+	}
+
+	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{} {
 	requestData := map[string]interface{}{
 		"query":  input.Query,
 		"models": input.ModelIDs,
@@ -62,31 +125,29 @@ func (p *AgentProvider) Generate(ctx *agentContext.Context, input *Input) (*Resu
 		requestData["extra"] = input.ExtraParams
 	}
 
-	requestJSON, _ := json.Marshal(requestData)
-
-	// Create message for the agent
-	messages := []agentContext.Message{
-		{
-			Role:    "user",
-			Content: string(requestJSON),
-		},
+	// Add retry context if this is a retry attempt
+	if attempt > 1 && lastLintErrors != "" {
+		requestData["retry"] = map[string]interface{}{
+			"attempt":      attempt,
+			"lint_errors":  lastLintErrors,
+			"instructions": "The previous QueryDSL was invalid. Please fix the errors and regenerate.",
+		}
 	}
 
-	// Call the agent with skip options (no history, no output)
-	options := &agentContext.Options{
-		Skip: &agentContext.Skip{
-			History: true,
-			Output:  true,
-		},
-	}
+	return requestData
+}
 
-	result, err := agent.Stream(ctx, messages, options)
+// validateDSL validates the generated QueryDSL using the linter
+func (p *AgentProvider) validateDSL(dsl *gou.QueryDSL) *linter.LintResult {
+	// Marshal DSL to JSON for linting
+	jsonBytes, err := json.Marshal(dsl)
 	if err != nil {
-		return nil, fmt.Errorf("agent call failed: %w", err)
+		result := &linter.LintResult{Valid: false}
+		return result
 	}
 
-	// Parse the result
-	return p.parseResult(result)
+	_, lintResult := linter.Parse(string(jsonBytes))
+	return lintResult
 }
 
 // parseResult extracts QueryDSL from the agent's response

agent/search/nlp/querydsl/agent_test.go

@@ -187,6 +187,52 @@ func TestGenerator_Agent_Integration(t *testing.T) {
 	})
 }
 
+func TestAgentProvider_Generate_WithRetry(t *testing.T) {
+	// Skip if running short tests
+	if testing.Short() {
+		t.Skip("Skipping integration test")
+	}
+
+	// Initialize test environment
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	// Load the querydsl-agent-retry assistant
+	ast, err := assistant.Get("tests.querydsl-agent-retry")
+	require.NoError(t, err)
+	require.NotNil(t, ast)
+
+	// Create test context
+	ctx := newTestContext(t)
+
+	// Create Agent provider for tests.querydsl-agent-retry
+	// This agent returns invalid DSL on first call, valid on second
+	provider := querydsl.NewAgentProvider("tests.querydsl-agent-retry")
+	assert.NotNil(t, provider)
+
+	t.Run("retry_on_lint_failure", func(t *testing.T) {
+		input := &querydsl.Input{
+			Query:    "test retry mechanism",
+			ModelIDs: []string{"user"},
+			Limit:    10,
+		}
+
+		// This should succeed after retry
+		// First call returns invalid DSL (missing 'from')
+		// Second call (with lint_errors) returns valid DSL
+		result, err := provider.Generate(ctx, input)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		if result.DSL != nil {
+			// Should have valid DSL after retry
+			assert.NotNil(t, result.DSL.From, "DSL should have 'from' field after retry")
+			// Explain should indicate this was fixed after receiving lint errors
+			assert.Contains(t, result.Explain, "fixed after receiving lint errors")
+		}
+	})
+}
+
 // newTestContext creates a test context with required fields
 func newTestContext(t *testing.T) *context.Context {
 	t.Helper()

agent/search/nlp/querydsl/mcp.go

@@ -8,9 +8,13 @@ import (
 	"github.com/yaoapp/gou/mcp"
 	gouMCPTypes "github.com/yaoapp/gou/mcp/types"
 	"github.com/yaoapp/gou/query/gou"
+	"github.com/yaoapp/gou/query/linter"
 	agentContext "github.com/yaoapp/yao/agent/context"
 )
 
+// MaxRetries is the maximum number of retry attempts for QueryDSL generation
+const MaxRetries = 3
+
 // MCPProvider delegates QueryDSL generation to an MCP tool
 type MCPProvider struct {
 	serverID string // MCP server ID
@@ -30,7 +34,7 @@ func NewMCPProvider(mcpRef string) (*MCPProvider, error) {
 	}, nil
 }
 
-// Generate generates QueryDSL by calling the MCP tool
+// Generate generates QueryDSL by calling the MCP tool with retry and lint validation
 func (p *MCPProvider) Generate(ctx *agentContext.Context, input *Input) (*Result, error) {
 	// Get MCP client
 	client, err := mcp.Select(p.serverID)
@@ -38,8 +42,54 @@ func (p *MCPProvider) Generate(ctx *agentContext.Context, input *Input) (*Result
 		return nil, fmt.Errorf("MCP server '%s' not found: %w", p.serverID, err)
 	}
 
-	// Build arguments for the MCP tool
-	// Note: model metadata is loaded internally by the MCP tool
+	var lastError error
+	var lastLintErrors string
+
+	for attempt := 1; attempt <= MaxRetries; attempt++ {
+		// Build arguments for the MCP tool
+		arguments := p.buildArguments(input, attempt, lastLintErrors)
+
+		// Call the MCP tool
+		callResult, err := client.CallTool(ctx, p.toolName, arguments)
+		if err != nil {
+			lastError = fmt.Errorf("MCP tool call failed: %w", err)
+			continue
+		}
+
+		// Parse the result
+		result, err := p.parseResult(callResult)
+		if err != nil {
+			lastError = err
+			continue
+		}
+
+		// Validate with linter if DSL is present
+		if result.DSL != nil {
+			lintResult := p.validateDSL(result.DSL)
+			if lintResult.Valid {
+				return result, nil
+			}
+
+			// Lint failed, prepare error message for retry
+			lastLintErrors = lintResult.FormatDiagnostics()
+			lastError = fmt.Errorf("QueryDSL validation failed: %s", lastLintErrors)
+
+			// Add lint warnings to result warnings
+			for _, diag := range lintResult.Diagnostics {
+				result.Warnings = append(result.Warnings, fmt.Sprintf("[%s] %s: %s", diag.Code, diag.Path, diag.Message))
+			}
+			continue
+		}
+
+		// No DSL returned
+		lastError = fmt.Errorf("no QueryDSL returned from MCP tool")
+	}
+
+	return nil, fmt.Errorf("QueryDSL generation failed after %d attempts: %w", MaxRetries, lastError)
+}
+
+// buildArguments constructs the MCP tool arguments
+func (p *MCPProvider) buildArguments(input *Input, attempt int, lastLintErrors string) map[string]interface{} {
 	arguments := map[string]interface{}{
 		"query":  input.Query,
 		"models": input.ModelIDs,
@@ -60,14 +110,29 @@ func (p *MCPProvider) Generate(ctx *agentContext.Context, input *Input) (*Result
 		arguments["extra"] = input.ExtraParams
 	}
 
-	// Call the MCP tool (ctx embeds context.Context)
-	callResult, err := client.CallTool(ctx, p.toolName, arguments)
-	if err != nil {
-		return nil, fmt.Errorf("MCP tool call failed: %w", err)
+	// Add retry context if this is a retry attempt
+	if attempt > 1 && lastLintErrors != "" {
+		arguments["retry"] = map[string]interface{}{
+			"attempt":      attempt,
+			"lint_errors":  lastLintErrors,
+			"instructions": "The previous QueryDSL was invalid. Please fix the errors and regenerate.",
+		}
 	}
 
-	// Parse the result
-	return p.parseResult(callResult)
+	return arguments
+}
+
+// validateDSL validates the generated QueryDSL using the linter
+func (p *MCPProvider) validateDSL(dsl *gou.QueryDSL) *linter.LintResult {
+	// Marshal DSL to JSON for linting
+	jsonBytes, err := json.Marshal(dsl)
+	if err != nil {
+		result := &linter.LintResult{Valid: false}
+		return result
+	}
+
+	_, lintResult := linter.Parse(string(jsonBytes))
+	return lintResult
 }
 
 // parseResult extracts QueryDSL from the MCP tool response

agent/search/nlp/querydsl/mcp_test.go

@@ -200,3 +200,38 @@ func TestGenerator_MCP_Integration(t *testing.T) {
 		assert.NotEmpty(t, result.Warnings)
 	})
 }
+
+func TestMCPProvider_Generate_WithRetry(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ctx := newTestContext()
+
+	// Create MCP provider for search.generate_querydsl_with_retry
+	// This tool returns invalid DSL on first call, valid on second
+	provider, err := NewMCPProvider("search.generate_querydsl_with_retry")
+	assert.NoError(t, err)
+	assert.NotNil(t, provider)
+
+	t.Run("retry_on_lint_failure", func(t *testing.T) {
+		input := &Input{
+			Query:    "test retry mechanism",
+			ModelIDs: []string{"user"},
+			Limit:    10,
+		}
+
+		// This should succeed after retry
+		// First call returns invalid DSL (missing 'from')
+		// Second call (with lint_errors) returns valid DSL
+		result, err := provider.Generate(ctx, input)
+		assert.NoError(t, err)
+		assert.NotNil(t, result)
+
+		if result != nil && result.DSL != nil {
+			// Should have valid DSL after retry
+			assert.NotNil(t, result.DSL.From, "DSL should have 'from' field after retry")
+			// Explain should indicate this was fixed after receiving lint errors
+			assert.Contains(t, result.Explain, "fixed after receiving lint errors")
+		}
+	})
+}