From 01bf7dce58753f4bea2c6f856b93dd10fd78379a Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 18 Dec 2025 12:37:35 +0800 Subject: [PATCH] Refactor buildRequestMessage to use structured JSON format for agent communication - Updated the buildRequestMessage function to construct a JSON object for structured communication with the agent, replacing the previous string formatting approach. - Enhanced the handling of extra parameters, scenarios, and retry context by incorporating them into the JSON structure. - Improved the parseResult function to extract additional fields such as explain and warnings from the response, ensuring comprehensive result handling. - Streamlined the code for better readability and maintainability while preserving existing functionality. --- agent/search/nlp/querydsl/agent.go | 69 +++++++++++++++++++----------- 1 file changed, 43 insertions(+), 26 deletions(-) diff --git a/agent/search/nlp/querydsl/agent.go b/agent/search/nlp/querydsl/agent.go index 901c41e0..b4b5531b 100644 --- a/agent/search/nlp/querydsl/agent.go +++ b/agent/search/nlp/querydsl/agent.go @@ -103,35 +103,43 @@ func (p *AgentProvider) Generate(ctx *agentContext.Context, input *Input) (*Resu } // buildRequestMessage constructs the request message for the agent -// Format follows the querydsl agent prompts.yml: -// "用户查询\nSchema:\n{schema JSON}" +// Returns JSON format for structured communication with the agent func (p *AgentProvider) buildRequestMessage(input *Input, attempt int, lastLintErrors string) string { - // Build schema from extra params if provided - var schemaJSON string - if input.ExtraParams != nil { - if schema, ok := input.ExtraParams["schema"]; ok { - schemaBytes, _ := json.Marshal(schema) - schemaJSON = string(schemaBytes) - } + // Build request data as JSON + requestData := map[string]interface{}{ + "query": input.Query, + "models": input.ModelIDs, + "limit": input.Limit, } - // Build message in the expected format - message := input.Query - if schemaJSON != "" { - message = fmt.Sprintf("%s\nSchema:\n%s", input.Query, schemaJSON) + // Add schema from extra params if provided + if input.ExtraParams != nil { + if schema, ok := input.ExtraParams["schema"]; ok { + requestData["schema"] = schema + } } // Add scenario hint if specified (filter, aggregation, join, complex) if input.Scenario != "" { - message = fmt.Sprintf("%s\nScenario: %s", message, input.Scenario) + requestData["scenario"] = string(input.Scenario) + } + + // Add allowed fields if specified + if len(input.AllowedFields) > 0 { + requestData["allowed_fields"] = input.AllowedFields } // Add retry context if this is a retry attempt if attempt > 1 && lastLintErrors != "" { - message = fmt.Sprintf("%s\n\nPrevious attempt failed with errors:\n%s\n\nPlease fix the errors and regenerate.", message, lastLintErrors) + requestData["retry"] = map[string]interface{}{ + "attempt": attempt, + "lint_errors": lastLintErrors, + "instructions": "The previous QueryDSL was invalid. Please fix the errors and regenerate.", + } } - return message + jsonBytes, _ := json.Marshal(requestData) + return string(jsonBytes) } // validateDSL validates the generated QueryDSL using the linter @@ -158,12 +166,11 @@ func (p *AgentProvider) parseResult(result interface{}) (*Result, error) { // Handle *context.Response directly (most common case from Stream()) if resp, ok := result.(*agentContext.Response); ok { - genResult := &Result{} if resp.Next != nil { - // Next contains the QueryDSL from hook - genResult.DSL = p.extractDSL(resp.Next) + // Next contains the hook response, recursively parse it + return p.parseResult(resp.Next) } - return genResult, nil + return &Result{}, nil } // Try to convert to map first @@ -238,12 +245,22 @@ func (p *AgentProvider) parseResult(result interface{}) (*Result, error) { } } - // Check for "dsl" field wrapper + // 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 is a wrapped DSL: { dsl: {...} } + // 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 { @@ -253,21 +270,21 @@ func (p *AgentProvider) parseResult(result interface{}) (*Result, error) { // 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 } } - // Get explain if present + // Fallback: Get explain and warnings from top level 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) }