From 36939cd53ac3c2aeaf943c7a14aa25781056ceab Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Fri, 19 Dec 2025 09:38:55 +0800
Subject: [PATCH] Enhance Search Execution Results with DSL Handling

- Added a `Results` field to the `SearchExecutionResult` struct to store raw search results, facilitating the extraction of DSL from database searches.
- Implemented logic in the `saveSearch` method to extract and store the first DSL from DB search results, improving data retention for search operations.
- Introduced a `dslToMap` method in the DB handler to convert QueryDSL to a map format for storage, enhancing the flexibility of result handling.
- Updated the `Result` struct to include a `DSL` field for generated QueryDSL, ensuring comprehensive data representation for database search results.
---
 agent/assistant/search.go           | 12 ++++++++++++
 agent/search/handlers/db/handler.go | 24 ++++++++++++++++++++++++
 agent/search/types/types.go         |  3 +++
 3 files changed, 39 insertions(+)

diff --git a/agent/assistant/search.go b/agent/assistant/search.go
index c0bc05a2..4e344225 100644
--- a/agent/assistant/search.go
+++ b/agent/assistant/search.go
@@ -518,6 +518,7 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context
 		Keywords:   extractedKeywords,
 		Config:     ast.configToMap(searchConfig),
 		RefCtx:     refCtx,
+		Results:    results,
 		Duration:   duration,
 		SearchType: "auto",
 	})
@@ -925,6 +926,7 @@ type SearchExecutionResult struct {
 	Keywords   []searchTypes.Keyword         // Extracted keywords with weights
 	Config     map[string]any                // Search config used
 	RefCtx     *searchTypes.ReferenceContext // Reference context with results
+	Results    []*searchTypes.Result         // Raw search results (for extracting DSL, etc.)
 	Duration   int64                         // Search duration in ms
 	Error      error                         // Error if failed
 	SearchType string                        // "auto", "web", "kb", "db"
@@ -1012,6 +1014,16 @@ func (ast *Assistant) saveSearch(ctx *context.Context, execResult *SearchExecuti
 		searchRecord.Prompt = execResult.RefCtx.Prompt
 	}
 
+	// Extract DSL from DB search results
+	if execResult.Results != nil {
+		for _, result := range execResult.Results {
+			if result != nil && result.Type == searchTypes.SearchTypeDB && result.DSL != nil {
+				searchRecord.DSL = result.DSL
+				break // Only store the first DSL (usually there's only one DB search)
+			}
+		}
+	}
+
 	// Save to store
 	if err := store.SaveSearch(searchRecord); err != nil {
 		ctx.Logger.Warn("Failed to save search record: %v", err)
diff --git a/agent/search/handlers/db/handler.go b/agent/search/handlers/db/handler.go
index a728d39e..3bb8bf6c 100644
--- a/agent/search/handlers/db/handler.go
+++ b/agent/search/handlers/db/handler.go
@@ -207,6 +207,9 @@ func (h *Handler) SearchWithContext(ctx *agentContext.Context, req *types.Reques
 		items = items[:maxResults]
 	}
 
+	// 7. Convert DSL to map for storage
+	dslMap := h.dslToMap(result.DSL)
+
 	return &types.Result{
 		Type:     types.SearchTypeDB,
 		Query:    req.Query,
@@ -214,6 +217,7 @@ func (h *Handler) SearchWithContext(ctx *agentContext.Context, req *types.Reques
 		Items:    items,
 		Total:    len(items),
 		Duration: time.Since(start).Milliseconds(),
+		DSL:      dslMap,
 	}, nil
 }
 
@@ -377,3 +381,23 @@ func (h *Handler) extractContent(rec map[string]interface{}, mod *model.Model) s
 	}
 	return string(content)
 }
+
+// dslToMap converts QueryDSL to map for storage
+func (h *Handler) dslToMap(dsl *gou.QueryDSL) map[string]interface{} {
+	if dsl == nil {
+		return nil
+	}
+
+	// Marshal and unmarshal to get a clean map
+	data, err := json.Marshal(dsl)
+	if err != nil {
+		return nil
+	}
+
+	var result map[string]interface{}
+	if err := json.Unmarshal(data, &result); err != nil {
+		return nil
+	}
+
+	return result
+}
diff --git a/agent/search/types/types.go b/agent/search/types/types.go
index 6cb70207..2b8237de 100644
--- a/agent/search/types/types.go
+++ b/agent/search/types/types.go
@@ -81,6 +81,9 @@ type Result struct {
 
 	// Graph associations (KB only, if enabled)
 	GraphNodes []*GraphNode `json:"graph_nodes,omitempty"`
+
+	// DB specific
+	DSL map[string]interface{} `json:"dsl,omitempty"` // Generated QueryDSL (DB only)
 }
 
 // ResultItem represents a single search result item