diff --git a/agent/search/nlp/querydsl/agent.go b/agent/search/nlp/querydsl/agent.go new file mode 100644 index 00000000..702db21e --- /dev/null +++ b/agent/search/nlp/querydsl/agent.go @@ -0,0 +1,205 @@ +package querydsl + +import ( + "encoding/json" + "fmt" + + "github.com/yaoapp/gou/query/gou" + "github.com/yaoapp/yao/agent/caller" + agentContext "github.com/yaoapp/yao/agent/context" +) + +// AgentProvider delegates QueryDSL generation to an LLM-powered assistant +// The assistant can understand context and generate semantically correct QueryDSL +type AgentProvider struct { + agentID string // Assistant ID to delegate to +} + +// NewAgentProvider creates a new agent-based QueryDSL generator +func NewAgentProvider(agentID string) *AgentProvider { + return &AgentProvider{ + agentID: agentID, + } +} + +// Generate generates QueryDSL by calling the target agent +// The agent receives the query and schema, returns generated QueryDSL +func (p *AgentProvider) Generate(ctx *agentContext.Context, input *Input) (*Result, error) { + if ctx == nil { + return nil, fmt.Errorf("context is required for agent QueryDSL generation") + } + + // Check if AgentGetterFunc is initialized + if caller.AgentGetterFunc == nil { + return nil, fmt.Errorf("AgentGetterFunc not initialized") + } + + // Get the agent + agent, err := caller.AgentGetterFunc(p.agentID) + if err != nil { + 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 + 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 + } + if len(input.Orders) > 0 { + requestData["orders"] = input.Orders + } + if len(input.AllowedFields) > 0 { + requestData["allowed_fields"] = input.AllowedFields + } + if len(input.ExtraParams) > 0 { + requestData["extra"] = input.ExtraParams + } + + 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 { + return nil, fmt.Errorf("agent call failed: %w", err) + } + + // Parse the result + return p.parseResult(result) +} + +// 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 { + return &Result{}, nil + } + + // Try to convert to map first (most common case) + var data map[string]interface{} + + switch v := result.(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, fmt.Errorf("failed to parse agent response: %w", err) + } + default: + // Try to marshal and unmarshal + jsonBytes, err := json.Marshal(result) + if err != nil { + return &Result{}, nil + } + if err := json.Unmarshal(jsonBytes, &data); err != nil { + return &Result{}, nil + } + } + + // 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 + } + + // Get warnings if present + if warnings, ok := data["warnings"]; ok { + genResult.Warnings = p.extractWarnings(warnings) + } + + // Get DSL + if dsl, ok := data["dsl"]; ok { + genResult.DSL = p.extractDSL(dsl) + } else if d, ok := data["data"]; ok { + if dm, ok := d.(map[string]interface{}); ok { + if dsl, ok := dm["dsl"]; ok { + genResult.DSL = p.extractDSL(dsl) + } + if explain, ok := dm["explain"].(string); ok { + genResult.Explain = explain + } + if warnings, ok := dm["warnings"]; ok { + genResult.Warnings = p.extractWarnings(warnings) + } + } + } + + return genResult, nil +} + +// extractDSL converts interface{} to gou.QueryDSL +func (p *AgentProvider) extractDSL(v interface{}) *gou.QueryDSL { + if v == nil { + return nil + } + + // Marshal and unmarshal to gou.QueryDSL + jsonBytes, err := json.Marshal(v) + if err != nil { + return nil + } + + var dsl gou.QueryDSL + if err := json.Unmarshal(jsonBytes, &dsl); err != nil { + return nil + } + + return &dsl +} + +// extractWarnings extracts warnings array from various types +func (p *AgentProvider) extractWarnings(v interface{}) []string { + switch w := v.(type) { + case []string: + return w + case []interface{}: + warnings := make([]string, 0, len(w)) + for _, item := range w { + if s, ok := item.(string); ok { + warnings = append(warnings, s) + } + } + return warnings + case string: + return []string{w} + } + return nil +} diff --git a/agent/search/nlp/querydsl/agent_test.go b/agent/search/nlp/querydsl/agent_test.go new file mode 100644 index 00000000..51f052a2 --- /dev/null +++ b/agent/search/nlp/querydsl/agent_test.go @@ -0,0 +1,199 @@ +package querydsl_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/assistant" + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/nlp/querydsl" + "github.com/yaoapp/yao/agent/testutils" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" +) + +func TestNewAgentProvider(t *testing.T) { + t.Run("create_provider", func(t *testing.T) { + provider := querydsl.NewAgentProvider("tests.querydsl-agent") + assert.NotNil(t, provider) + }) +} + +func TestAgentProvider_Generate(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 assistant + ast, err := assistant.Get("tests.querydsl-agent") + require.NoError(t, err) + require.NotNil(t, ast) + + // Create test context + ctx := newTestContext(t) + + // Create Agent provider for tests.querydsl-agent + provider := querydsl.NewAgentProvider("tests.querydsl-agent") + assert.NotNil(t, provider) + + t.Run("verify_fixed_structure", func(t *testing.T) { + input := &querydsl.Input{ + Query: "find active users", + ModelIDs: []string{"user"}, + Limit: 15, + } + + result, err := provider.Generate(ctx, input) + if err != nil { + t.Logf("Generate error: %v", err) + } + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.DSL, "DSL should not be nil") + + // Verify fixed DSL structure from mock + // select: ["id", "name", "status", "created_at"] + assert.Len(t, result.DSL.Select, 4) + if len(result.DSL.Select) >= 4 { + assert.Equal(t, "id", result.DSL.Select[0].Field) + assert.Equal(t, "name", result.DSL.Select[1].Field) + assert.Equal(t, "status", result.DSL.Select[2].Field) + assert.Equal(t, "created_at", result.DSL.Select[3].Field) + } + + // wheres: [{ field: "status", op: "=", value: "active" }] + assert.Len(t, result.DSL.Wheres, 1) + if len(result.DSL.Wheres) > 0 { + assert.Equal(t, "status", result.DSL.Wheres[0].Field.Field) + assert.Equal(t, "=", result.DSL.Wheres[0].OP) + assert.Equal(t, "active", result.DSL.Wheres[0].Value) + } + + // orders: [{ field: "created_at", sort: "desc" }] + assert.Len(t, result.DSL.Orders, 1) + if len(result.DSL.Orders) > 0 { + assert.Equal(t, "created_at", result.DSL.Orders[0].Field.Field) + assert.Equal(t, "desc", result.DSL.Orders[0].Sort) + } + + // limit: 15 (from input) + assert.Equal(t, float64(15), result.DSL.Limit) + + // explain should contain query + assert.Contains(t, result.Explain, "find active users") + + // warnings should be empty + assert.Empty(t, result.Warnings) + }) +} + +func TestAgentProvider_Generate_Error(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) + + // Create test context + ctx := newTestContext(t) + + t.Run("non-existent_agent", func(t *testing.T) { + provider := querydsl.NewAgentProvider("tests.nonexistent-agent") + result, err := provider.Generate(ctx, &querydsl.Input{ + Query: "test", + ModelIDs: []string{"user"}, + }) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "failed to get agent") + }) + + t.Run("nil_context", func(t *testing.T) { + provider := querydsl.NewAgentProvider("tests.querydsl-agent") + result, err := provider.Generate(nil, &querydsl.Input{ + Query: "test", + ModelIDs: []string{"user"}, + }) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "context is required") + }) +} + +func TestGenerator_Agent_Integration(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) + + // Create test context + ctx := newTestContext(t) + + // Create generator with Agent mode (assistant ID without mcp: prefix) + gen := querydsl.NewGenerator("tests.querydsl-agent", nil) + + t.Run("generate_via_agent", func(t *testing.T) { + input := &querydsl.Input{ + Query: "find active users", + ModelIDs: []string{"user"}, + Limit: 10, + } + + result, err := gen.Generate(ctx, input) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.DSL) + + // Verify structure from agent mock + assert.Len(t, result.DSL.Select, 4) + assert.Len(t, result.DSL.Wheres, 1) + assert.Len(t, result.DSL.Orders, 1) + assert.Contains(t, result.Explain, "find active users") + }) + + t.Run("allowed_fields_validation", func(t *testing.T) { + input := &querydsl.Input{ + Query: "find users", + ModelIDs: []string{"user"}, + AllowedFields: []string{"id", "name"}, // Only allow id and name + Limit: 10, + } + + result, err := gen.Generate(ctx, input) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.DSL) + + // "status" and "created_at" fields should be filtered out from select + // since they are not in AllowedFields + for _, expr := range result.DSL.Select { + assert.Contains(t, []string{"id", "name"}, expr.Field) + } + + // Should have warning about removed fields + assert.NotEmpty(t, result.Warnings) + }) +} + +// 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-querydsl" + ctx := context.New(t.Context(), authorized, chatID) + return ctx +} diff --git a/agent/search/nlp/querydsl/builtin.go b/agent/search/nlp/querydsl/builtin.go new file mode 100644 index 00000000..76e5cdc6 --- /dev/null +++ b/agent/search/nlp/querydsl/builtin.go @@ -0,0 +1,124 @@ +package querydsl + +import ( + "github.com/yaoapp/gou/model" + "github.com/yaoapp/gou/query/gou" +) + +// BuiltinGenerator implements template-based QueryDSL generation +// This is a placeholder implementation that returns a basic QueryDSL. +// +// TODO: Implement actual template-based generation: +// - Parse natural language query +// - Match against model schema +// - Generate appropriate where clauses +// - Handle common query patterns (search, filter, sort) +// +// For production use cases requiring high accuracy, use Agent or MCP mode. +type BuiltinGenerator struct{} + +// NewBuiltinGenerator creates a new builtin QueryDSL generator +func NewBuiltinGenerator() *BuiltinGenerator { + return &BuiltinGenerator{} +} + +// Generate generates QueryDSL from natural language +// Currently returns a placeholder QueryDSL that searches all searchable fields +func (g *BuiltinGenerator) Generate(input *Input) (*Result, error) { + if input == nil || input.Query == "" { + return &Result{ + Warnings: []string{"empty query, returning empty DSL"}, + }, nil + } + + // Build a basic QueryDSL + dsl := &gou.QueryDSL{} + + // Set limit + limit := input.Limit + if limit <= 0 { + limit = 20 + } + dsl.Limit = limit + + // Apply pre-defined wheres if provided + if len(input.Wheres) > 0 { + dsl.Wheres = input.Wheres + } + + // Apply orders if provided + if len(input.Orders) > 0 { + dsl.Orders = input.Orders + } + + // Load models and try to generate basic search conditions + // Use the first model as the primary table, others can be joined + if len(input.ModelIDs) > 0 { + primaryModelID := input.ModelIDs[0] + + // Check if model exists before selecting + if !model.Exists(primaryModelID) { + return &Result{ + DSL: dsl, + Explain: "Generated basic QueryDSL (model not found)", + Warnings: []string{ + "model '" + primaryModelID + "' not found, returning basic DSL without search conditions", + }, + }, nil + } + + primaryModel := model.Select(primaryModelID) + if primaryModel != nil && len(primaryModel.MetaData.Columns) > 0 { + // Find searchable text columns (string/text types with index) + var searchableColumns []string + for _, col := range primaryModel.MetaData.Columns { + // Use Index as a proxy for searchable, and check for text types + if col.Index && (col.Type == "string" || col.Type == "text" || col.Type == "longText") { + searchableColumns = append(searchableColumns, col.Name) + } + } + + // If we have searchable columns and no pre-defined wheres, add a basic search + if len(searchableColumns) > 0 && len(input.Wheres) == 0 { + // Build OR conditions for searchable columns + orWheres := make([]gou.Where, 0, len(searchableColumns)) + for _, col := range searchableColumns { + orWheres = append(orWheres, gou.Where{ + Condition: gou.Condition{ + Field: &gou.Expression{Field: col}, + OP: "match", + Value: input.Query, + }, + }) + } + + // Wrap in OR group if multiple columns + if len(orWheres) > 1 { + // Mark all but the first as OR conditions + for i := 1; i < len(orWheres); i++ { + orWheres[i].OR = true + } + dsl.Wheres = []gou.Where{ + { + Wheres: orWheres, + }, + } + } else if len(orWheres) == 1 { + dsl.Wheres = orWheres + } + } + } + + // TODO: For multi-model queries, generate joins based on model relations + // This requires analyzing the relations between models and generating + // appropriate JOIN clauses in the QueryDSL + } + + return &Result{ + DSL: dsl, + Explain: "Generated basic search QueryDSL using builtin template (placeholder implementation)", + Warnings: []string{ + "builtin generator is a placeholder, consider using Agent or MCP mode for production", + }, + }, nil +} diff --git a/agent/search/nlp/querydsl/generator.go b/agent/search/nlp/querydsl/generator.go new file mode 100644 index 00000000..b36c6285 --- /dev/null +++ b/agent/search/nlp/querydsl/generator.go @@ -0,0 +1,170 @@ +// Package querydsl provides QueryDSL generation from natural language for DB search +// Supports three modes via uses.querydsl configuration: +// - "builtin": Template-based generation (no external dependencies) +// - "": Delegate to an LLM-powered assistant for high-quality generation +// - "mcp:.": Call external MCP tool +// +// For production use cases requiring high accuracy, use Agent or MCP mode. +package querydsl + +import ( + "strings" + + "github.com/yaoapp/gou/query/gou" + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/types" +) + +// Generator generates QueryDSL from natural language +// Mode is determined by uses.querydsl configuration +type Generator struct { + usesQueryDSL string // "builtin", "", "mcp:." + config *types.QueryDSLConfig // QueryDSL generation options +} + +// NewGenerator creates a new QueryDSL generator +// usesQueryDSL: value from uses.querydsl config +// cfg: QueryDSL generation options from search config +func NewGenerator(usesQueryDSL string, cfg *types.QueryDSLConfig) *Generator { + return &Generator{ + usesQueryDSL: usesQueryDSL, + config: cfg, + } +} + +// Generate generates QueryDSL from natural language based on configured mode +// Returns a QueryDSL ready for execution +func (g *Generator) Generate(ctx *context.Context, input *Input) (*Result, error) { + var result *Result + var err error + + switch { + case g.usesQueryDSL == "builtin" || g.usesQueryDSL == "": + result, err = g.builtinGenerate(input) + 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) + } + + if err != nil { + return nil, err + } + + // Validate generated DSL against allowed fields whitelist + if result != nil && result.DSL != nil && len(input.AllowedFields) > 0 { + result = g.validateFields(result, input.AllowedFields) + } + + 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) + return provider.Generate(ctx, input) +} + +// mcpGenerate calls an external MCP tool +// Format: "mcp:." +func (g *Generator) mcpGenerate(ctx *context.Context, input *Input) (*Result, error) { + mcpRef := strings.TrimPrefix(g.usesQueryDSL, "mcp:") + provider, err := NewMCPProvider(mcpRef) + if err != nil { + // Fallback to builtin on invalid MCP format + return g.builtinGenerate(input) + } + return provider.Generate(ctx, input) +} + +// validateFields validates that all fields in the generated DSL are in the allowed list +// If a field is not allowed, it's removed and a warning is added +func (g *Generator) validateFields(result *Result, allowedFields []string) *Result { + if result.DSL == nil { + return result + } + + // Build allowed fields set for fast lookup + allowed := make(map[string]bool) + for _, f := range allowedFields { + allowed[f] = true + } + + var removedFields []string + + // Validate Select fields + if len(result.DSL.Select) > 0 { + validSelect := make([]gou.Expression, 0, len(result.DSL.Select)) + for _, expr := range result.DSL.Select { + if allowed[expr.Field] { + validSelect = append(validSelect, expr) + } else if expr.Field != "" { + removedFields = append(removedFields, "select:"+expr.Field) + } + } + result.DSL.Select = validSelect + } + + // Validate Where fields (recursive) + result.DSL.Wheres = g.validateWheres(result.DSL.Wheres, allowed, &removedFields) + + // Validate Order fields + if len(result.DSL.Orders) > 0 { + validOrders := make(gou.Orders, 0, len(result.DSL.Orders)) + for _, order := range result.DSL.Orders { + if order.Field != nil && allowed[order.Field.Field] { + validOrders = append(validOrders, order) + } else if order.Field != nil && order.Field.Field != "" { + removedFields = append(removedFields, "order:"+order.Field.Field) + } + } + result.DSL.Orders = validOrders + } + + // Add warnings for removed fields + if len(removedFields) > 0 { + warning := "removed fields not in allowed list: " + strings.Join(removedFields, ", ") + result.Warnings = append(result.Warnings, warning) + } + + return result +} + +// validateWheres recursively validates where conditions +func (g *Generator) validateWheres(wheres []gou.Where, allowed map[string]bool, removedFields *[]string) []gou.Where { + if len(wheres) == 0 { + return wheres + } + + validWheres := make([]gou.Where, 0, len(wheres)) + for _, w := range wheres { + // Check if the field is allowed + fieldAllowed := true + if w.Field != nil && w.Field.Field != "" { + if !allowed[w.Field.Field] { + *removedFields = append(*removedFields, "where:"+w.Field.Field) + fieldAllowed = false + } + } + + if fieldAllowed { + // Recursively validate nested wheres + if len(w.Wheres) > 0 { + w.Wheres = g.validateWheres(w.Wheres, allowed, removedFields) + } + validWheres = append(validWheres, w) + } + } + + return validWheres +} diff --git a/agent/search/nlp/querydsl/generator_test.go b/agent/search/nlp/querydsl/generator_test.go new file mode 100644 index 00000000..0f2b5736 --- /dev/null +++ b/agent/search/nlp/querydsl/generator_test.go @@ -0,0 +1,203 @@ +package querydsl + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/gou/query/gou" + "github.com/yaoapp/yao/agent/search/types" +) + +func TestNewGenerator(t *testing.T) { + tests := []struct { + name string + usesQueryDSL string + config *types.QueryDSLConfig + }{ + { + name: "builtin mode", + usesQueryDSL: "builtin", + config: nil, + }, + { + name: "empty defaults to builtin", + usesQueryDSL: "", + config: nil, + }, + { + name: "agent mode", + usesQueryDSL: "my-querydsl-agent", + config: &types.QueryDSLConfig{Strict: true}, + }, + { + name: "mcp mode", + usesQueryDSL: "mcp:nlp.generate_querydsl", + config: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gen := NewGenerator(tt.usesQueryDSL, tt.config) + assert.NotNil(t, gen) + assert.Equal(t, tt.usesQueryDSL, gen.usesQueryDSL) + assert.Equal(t, tt.config, gen.config) + }) + } +} + +func TestGenerator_Generate_Builtin(t *testing.T) { + 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) +} + +func TestGenerator_Generate_EmptyMode(t *testing.T) { + // Empty mode should default to builtin + gen := NewGenerator("", nil) + + input := &Input{ + Query: "search products", + ModelIDs: []string{"product"}, + Limit: 5, + } + + result, err := gen.Generate(nil, input) + assert.NoError(t, err) + assert.NotNil(t, result) +} + +func TestBuiltinGenerator_Generate(t *testing.T) { + gen := NewBuiltinGenerator() + + 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") + }) + + 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) + }) + + 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) + }) + + 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) + }) + + 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) + }) + + 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 TestResult(t *testing.T) { + result := &Result{ + DSL: &gou.QueryDSL{ + Limit: 10, + }, + Explain: "Generated query for finding users", + Warnings: []string{"using placeholder implementation"}, + } + + assert.NotNil(t, result.DSL) + assert.Equal(t, 10, result.DSL.Limit) + assert.NotEmpty(t, result.Explain) + assert.Len(t, result.Warnings, 1) +} diff --git a/agent/search/nlp/querydsl/mcp.go b/agent/search/nlp/querydsl/mcp.go new file mode 100644 index 00000000..a4b75339 --- /dev/null +++ b/agent/search/nlp/querydsl/mcp.go @@ -0,0 +1,164 @@ +package querydsl + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/yaoapp/gou/mcp" + gouMCPTypes "github.com/yaoapp/gou/mcp/types" + "github.com/yaoapp/gou/query/gou" + agentContext "github.com/yaoapp/yao/agent/context" +) + +// MCPProvider delegates QueryDSL generation to an MCP tool +type MCPProvider struct { + serverID string // MCP server ID + toolName string // Tool name to call +} + +// NewMCPProvider creates a new MCP-based QueryDSL generator +// mcpRef format: "server.tool" (e.g., "nlp.generate_querydsl") +func NewMCPProvider(mcpRef string) (*MCPProvider, error) { + parts := strings.SplitN(mcpRef, ".", 2) + if len(parts) != 2 { + return nil, fmt.Errorf("invalid MCP format, expected 'server.tool', got '%s'", mcpRef) + } + return &MCPProvider{ + serverID: parts[0], + toolName: parts[1], + }, nil +} + +// Generate generates QueryDSL by calling the MCP tool +func (p *MCPProvider) Generate(ctx *agentContext.Context, input *Input) (*Result, error) { + // Get MCP client + client, err := mcp.Select(p.serverID) + if err != nil { + 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 + arguments := map[string]interface{}{ + "query": input.Query, + "models": input.ModelIDs, + "limit": input.Limit, + } + + // Add optional fields + if len(input.Wheres) > 0 { + arguments["wheres"] = input.Wheres + } + if len(input.Orders) > 0 { + arguments["orders"] = input.Orders + } + if len(input.AllowedFields) > 0 { + arguments["allowed_fields"] = input.AllowedFields + } + if len(input.ExtraParams) > 0 { + 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) + } + + // Parse the result + return p.parseResult(callResult) +} + +// parseResult extracts QueryDSL from the MCP tool response +func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse) (*Result, error) { + if result == nil { + return &Result{}, nil + } + + // Check for errors in result + if result.IsError { + errMsg := "MCP tool returned error" + if len(result.Content) > 0 && result.Content[0].Text != "" { + errMsg = result.Content[0].Text + } + return nil, fmt.Errorf("%s", errMsg) + } + + // Parse content - expect JSON data with "dsl" field + if len(result.Content) == 0 { + return &Result{}, nil + } + + genResult := &Result{} + + // Try to extract QueryDSL from content + for _, content := range result.Content { + // Check text content type + if content.Type == gouMCPTypes.ToolContentTypeText && content.Text != "" { + // Try to parse as JSON + var data map[string]interface{} + if err := json.Unmarshal([]byte(content.Text), &data); err == nil { + // Look for "dsl" field + if dsl, ok := data["dsl"]; ok { + genResult.DSL = p.extractDSL(dsl) + } + if explain, ok := data["explain"].(string); ok { + genResult.Explain = explain + } + if warnings, ok := data["warnings"]; ok { + genResult.Warnings = p.extractWarnings(warnings) + } + return genResult, nil + } + + // Try to parse as direct QueryDSL + var dsl gou.QueryDSL + if err := json.Unmarshal([]byte(content.Text), &dsl); err == nil { + genResult.DSL = &dsl + return genResult, nil + } + } + } + + return genResult, nil +} + +// extractDSL converts interface{} to gou.QueryDSL +func (p *MCPProvider) extractDSL(v interface{}) *gou.QueryDSL { + if v == nil { + return nil + } + + // Marshal and unmarshal to gou.QueryDSL + jsonBytes, err := json.Marshal(v) + if err != nil { + return nil + } + + var dsl gou.QueryDSL + if err := json.Unmarshal(jsonBytes, &dsl); err != nil { + return nil + } + + return &dsl +} + +// extractWarnings extracts warnings array from various types +func (p *MCPProvider) extractWarnings(v interface{}) []string { + switch w := v.(type) { + case []string: + return w + case []interface{}: + warnings := make([]string, 0, len(w)) + for _, item := range w { + if s, ok := item.(string); ok { + warnings = append(warnings, s) + } + } + return warnings + case string: + return []string{w} + } + return nil +} diff --git a/agent/search/nlp/querydsl/mcp_test.go b/agent/search/nlp/querydsl/mcp_test.go new file mode 100644 index 00000000..650cc3d8 --- /dev/null +++ b/agent/search/nlp/querydsl/mcp_test.go @@ -0,0 +1,202 @@ +package querydsl + +import ( + stdContext "context" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/gou/plan" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/test" +) + +// newTestContext creates a test context for MCP testing +func newTestContext() *agentContext.Context { + ctx := &agentContext.Context{ + Context: stdContext.Background(), + Space: plan.NewMemorySharedSpace(), + ID: "test-querydsl", + ChatID: "test-chat", + AssistantID: "test-assistant", + Locale: "en", + Referer: agentContext.RefererAPI, + } + stack, _, _ := agentContext.EnterStack(ctx, "test-assistant", &agentContext.Options{}) + ctx.Stack = stack + return ctx +} + +func TestMCPProvider_Generate(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + // Create context + ctx := newTestContext() + + // Create MCP provider for search.generate_querydsl + provider, err := NewMCPProvider("search.generate_querydsl") + assert.NoError(t, err) + assert.NotNil(t, provider) + assert.Equal(t, "search", provider.serverID) + assert.Equal(t, "generate_querydsl", provider.toolName) + + t.Run("verify_fixed_structure", func(t *testing.T) { + input := &Input{ + Query: "find active users", + ModelIDs: []string{"user"}, + Limit: 10, + } + + result, err := provider.Generate(ctx, input) + if err != nil { + t.Logf("Generate error: %v", err) + } + assert.NoError(t, err) + assert.NotNil(t, result) + + if result == nil { + t.Fatal("result is nil") + } + + if !assert.NotNil(t, result.DSL, "DSL should not be nil") { + t.Logf("Result: Explain=%s, Warnings=%v", result.Explain, result.Warnings) + return + } + + // Verify fixed DSL structure from mock + // select: ["id", "name", "status"] - parsed as Expression with Field property + assert.Len(t, result.DSL.Select, 3) + if len(result.DSL.Select) >= 3 { + assert.Equal(t, "id", result.DSL.Select[0].Field) + assert.Equal(t, "name", result.DSL.Select[1].Field) + assert.Equal(t, "status", result.DSL.Select[2].Field) + } + + // wheres: [{ field: "status", op: "=", value: "active" }] + assert.Len(t, result.DSL.Wheres, 1) + if len(result.DSL.Wheres) > 0 { + assert.Equal(t, "status", result.DSL.Wheres[0].Field.Field) + assert.Equal(t, "=", result.DSL.Wheres[0].OP) + assert.Equal(t, "active", result.DSL.Wheres[0].Value) + } + + // orders: [{ field: "created_at", sort: "desc" }] + assert.Len(t, result.DSL.Orders, 1) + if len(result.DSL.Orders) > 0 { + assert.Equal(t, "created_at", result.DSL.Orders[0].Field.Field) + assert.Equal(t, "desc", result.DSL.Orders[0].Sort) + } + + // limit: 10 (from input, returned as float64 from JSON) + assert.Equal(t, float64(10), result.DSL.Limit) + + // explain should contain query + assert.Contains(t, result.Explain, "find active users") + + // warnings should be empty + assert.Empty(t, result.Warnings) + }) +} + +func TestNewMCPProvider(t *testing.T) { + t.Run("valid format", func(t *testing.T) { + provider, err := NewMCPProvider("nlp.generate_querydsl") + assert.NoError(t, err) + assert.NotNil(t, provider) + assert.Equal(t, "nlp", provider.serverID) + assert.Equal(t, "generate_querydsl", provider.toolName) + }) + + t.Run("invalid format - no dot", func(t *testing.T) { + provider, err := NewMCPProvider("invalid") + assert.Error(t, err) + assert.Nil(t, provider) + assert.Contains(t, err.Error(), "invalid MCP format") + }) + + t.Run("complex tool name", func(t *testing.T) { + provider, err := NewMCPProvider("server.tool.with.dots") + assert.NoError(t, err) + assert.NotNil(t, provider) + assert.Equal(t, "server", provider.serverID) + assert.Equal(t, "tool.with.dots", provider.toolName) + }) +} + +func TestMCPProvider_Generate_Error(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + ctx := newTestContext() + + t.Run("non-existent server", func(t *testing.T) { + provider, _ := NewMCPProvider("nonexistent.tool") + result, err := provider.Generate(ctx, &Input{ + Query: "test", + ModelIDs: []string{"user"}, + }) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "not found") + }) +} + +func TestGenerator_MCP_Integration(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + // Skip if not in integration test mode + if os.Getenv("YAO_TEST_MCP") != "true" { + t.Skip("Skipping MCP integration test (set YAO_TEST_MCP=true to run)") + } + + ctx := newTestContext() + + // Create generator with MCP mode + gen := NewGenerator("mcp:search.generate_querydsl", nil) + + t.Run("generate_via_mcp", func(t *testing.T) { + input := &Input{ + Query: "find active users", + ModelIDs: []string{"user"}, + Limit: 15, + } + + result, err := gen.Generate(ctx, input) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotNil(t, result.DSL) + + // Verify fixed structure is correctly parsed + assert.Len(t, result.DSL.Select, 3) + assert.Len(t, result.DSL.Wheres, 1) + assert.Len(t, result.DSL.Orders, 1) + assert.Equal(t, float64(15), result.DSL.Limit) + assert.Contains(t, result.Explain, "find active users") + }) + + t.Run("allowed_fields_validation", func(t *testing.T) { + input := &Input{ + Query: "find users", + ModelIDs: []string{"user"}, + AllowedFields: []string{"id", "name"}, // Only allow id and name + Limit: 10, + } + + result, err := gen.Generate(ctx, input) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotNil(t, result.DSL) + + // "status" field should be filtered out from select and wheres + // since it's not in AllowedFields + for _, expr := range result.DSL.Select { + assert.Contains(t, []string{"id", "name"}, expr.Field) + } + + // Should have warning about removed fields + assert.NotEmpty(t, result.Warnings) + }) +} diff --git a/agent/search/nlp/querydsl/types.go b/agent/search/nlp/querydsl/types.go new file mode 100644 index 00000000..0c9b6d3b --- /dev/null +++ b/agent/search/nlp/querydsl/types.go @@ -0,0 +1,23 @@ +package querydsl + +import ( + "github.com/yaoapp/gou/query/gou" +) + +// 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"]) + Wheres []gou.Where // Pre-defined filters (optional) + Orders gou.Orders // Sort orders (optional) + AllowedFields []string // Allowed fields whitelist (optional, for security validation) + Limit int // Max results + ExtraParams map[string]interface{} // Additional parameters +} + +// Result represents the result of QueryDSL generation +type Result struct { + DSL *gou.QueryDSL `json:"dsl"` // Generated QueryDSL (supports joins) + Explain string `json:"explain,omitempty"` // Human-readable explanation + Warnings []string `json:"warnings,omitempty"` // Any warnings during generation +}