diff --git a/agent/search/nlp/querydsl/agent.go b/agent/search/nlp/querydsl/agent.go index 702db21e..7142d9a5 100644 --- a/agent/search/nlp/querydsl/agent.go +++ b/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 diff --git a/agent/search/nlp/querydsl/agent_test.go b/agent/search/nlp/querydsl/agent_test.go index 51f052a2..60300be6 100644 --- a/agent/search/nlp/querydsl/agent_test.go +++ b/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() diff --git a/agent/search/nlp/querydsl/mcp.go b/agent/search/nlp/querydsl/mcp.go index a4b75339..fbdf27a5 100644 --- a/agent/search/nlp/querydsl/mcp.go +++ b/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 diff --git a/agent/search/nlp/querydsl/mcp_test.go b/agent/search/nlp/querydsl/mcp_test.go index 650cc3d8..be413001 100644 --- a/agent/search/nlp/querydsl/mcp_test.go +++ b/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") + } + }) +}