Merge pull request #1384 from trheyi/main
Enhance Search Result Storage and Intermediate Data Handling
This commit is contained in:
Max
GitHub
commit
380cfa54bb
21 changed files with 3349 additions and 246 deletions
|
|
@ -1,12 +1,18 @@
|
|||
package assistant
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/i18n"
|
||||
"github.com/yaoapp/yao/agent/output/message"
|
||||
"github.com/yaoapp/yao/agent/search"
|
||||
"github.com/yaoapp/yao/agent/search/nlp/keyword"
|
||||
searchTypes "github.com/yaoapp/yao/agent/search/types"
|
||||
storeTypes "github.com/yaoapp/yao/agent/store/types"
|
||||
traceTypes "github.com/yaoapp/yao/trace/types"
|
||||
)
|
||||
|
||||
// shouldAutoSearch determines if auto search should be executed
|
||||
|
|
@ -98,12 +104,13 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context
|
|||
// Create searcher
|
||||
searcher := search.New(searchConfig, searchUses)
|
||||
|
||||
// Extract query from messages
|
||||
query := extractQueryFromMessages(messages)
|
||||
if query == "" {
|
||||
// Extract query from messages (save original for storage)
|
||||
originalQuery := extractQueryFromMessages(messages)
|
||||
if originalQuery == "" {
|
||||
ctx.Logger.Info("No query found in messages, skipping auto search")
|
||||
return nil
|
||||
}
|
||||
query := originalQuery
|
||||
|
||||
// Check if keyword extraction should be skipped
|
||||
skipKeyword := false
|
||||
|
|
@ -115,6 +122,7 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context
|
|||
// 1. uses.keyword is configured (not empty)
|
||||
// 2. Skip.Keyword is not true
|
||||
// 3. Web search is enabled
|
||||
var extractedKeywords []string
|
||||
webSearchEnabled := searchConfig != nil && searchConfig.Web != nil
|
||||
if webSearchEnabled && !skipKeyword && searchUses.Keyword != "" {
|
||||
extractor := keyword.NewExtractor(searchUses.Keyword, searchConfig.Keyword)
|
||||
|
|
@ -122,6 +130,7 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context
|
|||
if err != nil {
|
||||
ctx.Logger.Warn("Keyword extraction failed, using original query: %v", err)
|
||||
} else if len(keywords) > 0 {
|
||||
extractedKeywords = keywords
|
||||
// Use extracted keywords as the search query for web search
|
||||
optimizedQuery := strings.Join(keywords, " ")
|
||||
ctx.Logger.Info("Extracted keywords for web search: %s -> %s", truncateString(query, 30), optimizedQuery)
|
||||
|
|
@ -136,13 +145,39 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context
|
|||
return nil
|
||||
}
|
||||
|
||||
// === Output: Send loading message ===
|
||||
loadingID := ast.sendSearchLoading(ctx)
|
||||
|
||||
// === Trace: Create search trace node ===
|
||||
searchNode := ast.createSearchTrace(ctx, query, requests)
|
||||
|
||||
// Execute searches in parallel
|
||||
ctx.Logger.Info("Executing %d search requests for query: %s", len(requests), truncateString(query, 50))
|
||||
|
||||
startTime := time.Now()
|
||||
results, err := searcher.All(ctx, requests)
|
||||
duration := time.Since(startTime).Milliseconds()
|
||||
|
||||
if err != nil {
|
||||
// Log error but don't fail - search errors shouldn't block the main flow
|
||||
ctx.Logger.Error("Auto search failed: %v", err)
|
||||
|
||||
// === Output: Send failed message ===
|
||||
ast.sendSearchDone(ctx, loadingID, 0, true)
|
||||
|
||||
// === Trace: Mark as failed ===
|
||||
ast.completeSearchTrace(searchNode, 0, err)
|
||||
|
||||
// === Storage: Save failed search ===
|
||||
ast.saveSearch(ctx, &SearchExecutionResult{
|
||||
Query: originalQuery,
|
||||
Keywords: extractedKeywords,
|
||||
Config: ast.configToMap(searchConfig),
|
||||
Duration: duration,
|
||||
Error: err,
|
||||
SearchType: "auto",
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -153,15 +188,189 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context
|
|||
}
|
||||
refCtx := search.BuildReferenceContext(results, citationConfig)
|
||||
|
||||
if len(refCtx.References) == 0 {
|
||||
resultCount := len(refCtx.References)
|
||||
|
||||
// === Output: Send result message, then done ===
|
||||
ast.sendSearchResult(ctx, loadingID, resultCount)
|
||||
ast.sendSearchDone(ctx, loadingID, resultCount, false)
|
||||
|
||||
// === Trace: Mark as completed ===
|
||||
ast.completeSearchTrace(searchNode, resultCount, nil)
|
||||
|
||||
// === Storage: Save successful search ===
|
||||
ast.saveSearch(ctx, &SearchExecutionResult{
|
||||
Query: originalQuery,
|
||||
Keywords: extractedKeywords,
|
||||
Config: ast.configToMap(searchConfig),
|
||||
RefCtx: refCtx,
|
||||
Duration: duration,
|
||||
SearchType: "auto",
|
||||
})
|
||||
|
||||
if resultCount == 0 {
|
||||
ctx.Logger.Info("No search results found")
|
||||
return nil
|
||||
}
|
||||
|
||||
ctx.Logger.Info("Auto search completed: %d references", len(refCtx.References))
|
||||
ctx.Logger.Info("Auto search completed: %d references", resultCount)
|
||||
return refCtx
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Output: Loading Replace Pattern
|
||||
// ============================================================================
|
||||
|
||||
// sendSearchLoading sends the initial loading message
|
||||
// Returns the message ID for later replacement
|
||||
func (ast *Assistant) sendSearchLoading(ctx *context.Context) string {
|
||||
loadingMsg := i18n.T(ctx.Locale, "search.loading")
|
||||
|
||||
msg := &message.Message{
|
||||
Type: "loading",
|
||||
Props: map[string]any{
|
||||
"message": loadingMsg,
|
||||
},
|
||||
}
|
||||
|
||||
// Send and get message ID
|
||||
msgID, err := ctx.SendStream(msg)
|
||||
if err != nil {
|
||||
ctx.Logger.Warn("Failed to send search loading message: %v", err)
|
||||
return ""
|
||||
}
|
||||
|
||||
return msgID
|
||||
}
|
||||
|
||||
// sendSearchResult replaces loading with result message (without done flag)
|
||||
func (ast *Assistant) sendSearchResult(ctx *context.Context, loadingID string, count int) {
|
||||
if loadingID == "" {
|
||||
return
|
||||
}
|
||||
|
||||
var resultMsg string
|
||||
if count == 0 {
|
||||
resultMsg = i18n.T(ctx.Locale, "search.no_results")
|
||||
} else if count == 1 {
|
||||
resultMsg = i18n.T(ctx.Locale, "search.success.one")
|
||||
} else {
|
||||
resultMsg = fmt.Sprintf(i18n.T(ctx.Locale, "search.success"), count)
|
||||
}
|
||||
|
||||
msg := &message.Message{
|
||||
MessageID: loadingID,
|
||||
Delta: true,
|
||||
DeltaAction: message.DeltaReplace,
|
||||
Type: "loading",
|
||||
Props: map[string]any{
|
||||
"message": resultMsg,
|
||||
},
|
||||
}
|
||||
|
||||
if err := ctx.Send(msg); err != nil {
|
||||
ctx.Logger.Warn("Failed to send search result message: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// sendSearchDone sends the final done message (removes loading indicator)
|
||||
func (ast *Assistant) sendSearchDone(ctx *context.Context, loadingID string, count int, failed bool) {
|
||||
if loadingID == "" {
|
||||
return
|
||||
}
|
||||
|
||||
var resultMsg string
|
||||
if failed {
|
||||
resultMsg = i18n.T(ctx.Locale, "search.failed")
|
||||
} else if count == 0 {
|
||||
resultMsg = i18n.T(ctx.Locale, "search.no_results")
|
||||
} else if count == 1 {
|
||||
resultMsg = i18n.T(ctx.Locale, "search.success.one")
|
||||
} else {
|
||||
resultMsg = fmt.Sprintf(i18n.T(ctx.Locale, "search.success"), count)
|
||||
}
|
||||
|
||||
msg := &message.Message{
|
||||
MessageID: loadingID,
|
||||
Delta: true,
|
||||
DeltaAction: message.DeltaReplace,
|
||||
Type: "loading",
|
||||
Props: map[string]any{
|
||||
"message": resultMsg,
|
||||
"done": true, // Frontend will remove loading indicator
|
||||
},
|
||||
}
|
||||
|
||||
if err := ctx.Send(msg); err != nil {
|
||||
ctx.Logger.Warn("Failed to send search done message: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Trace: Search Node
|
||||
// ============================================================================
|
||||
|
||||
// createSearchTrace creates a trace node for search operation
|
||||
func (ast *Assistant) createSearchTrace(ctx *context.Context, query string, requests []*searchTypes.Request) traceTypes.Node {
|
||||
trace, _ := ctx.Trace()
|
||||
if trace == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Build search types list
|
||||
var searchTypes []string
|
||||
for _, req := range requests {
|
||||
searchTypes = append(searchTypes, string(req.Type))
|
||||
}
|
||||
|
||||
input := map[string]any{
|
||||
"query": query,
|
||||
"types": searchTypes,
|
||||
}
|
||||
|
||||
node, err := trace.Add(input, traceTypes.TraceNodeOption{
|
||||
Label: i18n.T(ctx.Locale, "search.trace.label"),
|
||||
Type: "search",
|
||||
Icon: "search",
|
||||
Description: i18n.T(ctx.Locale, "search.trace.description"),
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
ctx.Logger.Warn("Failed to create search trace node: %v", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Log search start
|
||||
node.Info("Starting search", map[string]any{
|
||||
"query": query,
|
||||
"types": searchTypes,
|
||||
})
|
||||
|
||||
return node
|
||||
}
|
||||
|
||||
// completeSearchTrace marks the search trace node as completed or failed
|
||||
func (ast *Assistant) completeSearchTrace(node traceTypes.Node, resultCount int, err error) {
|
||||
if node == nil {
|
||||
return
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
node.Warn("Search failed", map[string]any{"error": err.Error()})
|
||||
node.Fail(err)
|
||||
return
|
||||
}
|
||||
|
||||
// Log completion
|
||||
node.Info("Search completed", map[string]any{
|
||||
"result_count": resultCount,
|
||||
})
|
||||
|
||||
// Complete with output
|
||||
node.Complete(map[string]any{
|
||||
"result_count": resultCount,
|
||||
})
|
||||
}
|
||||
|
||||
// buildSearchRequests builds search requests based on assistant configuration
|
||||
func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Config) []*searchTypes.Request {
|
||||
var requests []*searchTypes.Request
|
||||
|
|
@ -299,3 +508,139 @@ func truncateString(s string, maxLen int) string {
|
|||
}
|
||||
return s[:maxLen] + "..."
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Storage: Save Search Results
|
||||
// ============================================================================
|
||||
|
||||
// SearchExecutionResult holds all data from search execution for storage
|
||||
type SearchExecutionResult struct {
|
||||
Query string // Original query (before keyword optimization)
|
||||
Keywords []string // Extracted keywords
|
||||
Config map[string]any // Search config used
|
||||
RefCtx *searchTypes.ReferenceContext // Reference context with results
|
||||
Duration int64 // Search duration in ms
|
||||
Error error // Error if failed
|
||||
SearchType string // "auto", "web", "kb", "db"
|
||||
}
|
||||
|
||||
// saveSearch saves search results to storage
|
||||
// Called after search execution completes (success or failure)
|
||||
func (ast *Assistant) saveSearch(ctx *context.Context, execResult *SearchExecutionResult) {
|
||||
// Get store
|
||||
store := GetStore()
|
||||
if store == nil {
|
||||
ctx.Logger.Debug("Storage not configured, skipping search save")
|
||||
return
|
||||
}
|
||||
|
||||
// Build search record
|
||||
searchRecord := &storeTypes.Search{
|
||||
RequestID: ctx.RequestID(),
|
||||
ChatID: ctx.ChatID,
|
||||
Query: execResult.Query,
|
||||
Keywords: execResult.Keywords,
|
||||
Config: execResult.Config,
|
||||
Source: execResult.SearchType,
|
||||
Duration: execResult.Duration,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
|
||||
// Set error if present
|
||||
if execResult.Error != nil {
|
||||
searchRecord.Error = execResult.Error.Error()
|
||||
}
|
||||
|
||||
// Convert references if available
|
||||
if execResult.RefCtx != nil {
|
||||
searchRecord.References = convertToStoreReferences(execResult.RefCtx.References)
|
||||
searchRecord.XML = execResult.RefCtx.XML
|
||||
searchRecord.Prompt = execResult.RefCtx.Prompt
|
||||
}
|
||||
|
||||
// Save to store
|
||||
if err := store.SaveSearch(searchRecord); err != nil {
|
||||
ctx.Logger.Warn("Failed to save search record: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
ctx.Logger.Debug("Search record saved: request_id=%s, refs=%d",
|
||||
searchRecord.RequestID, len(searchRecord.References))
|
||||
}
|
||||
|
||||
// convertToStoreReferences converts search References to store References
|
||||
func convertToStoreReferences(refs []*searchTypes.Reference) []storeTypes.Reference {
|
||||
if len(refs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
storeRefs := make([]storeTypes.Reference, len(refs))
|
||||
for i, ref := range refs {
|
||||
if ref == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Parse citation ID as integer (e.g., "1", "2", "3")
|
||||
index := i + 1 // Default to position-based index
|
||||
if ref.ID != "" {
|
||||
if n, err := fmt.Sscanf(ref.ID, "%d", &index); n != 1 || err != nil {
|
||||
index = i + 1
|
||||
}
|
||||
}
|
||||
|
||||
storeRefs[i] = storeTypes.Reference{
|
||||
Index: index,
|
||||
Type: string(ref.Type),
|
||||
Title: ref.Title,
|
||||
URL: ref.URL,
|
||||
Snippet: truncateString(ref.Content, 200), // Short snippet
|
||||
Content: ref.Content,
|
||||
Metadata: map[string]any{
|
||||
"weight": ref.Weight,
|
||||
"score": ref.Score,
|
||||
"source": string(ref.Source),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
return storeRefs
|
||||
}
|
||||
|
||||
// configToMap converts search config to map for storage
|
||||
func (ast *Assistant) configToMap(config *searchTypes.Config) map[string]any {
|
||||
if config == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
result := make(map[string]any)
|
||||
|
||||
if config.Web != nil {
|
||||
result["web"] = map[string]any{
|
||||
"provider": config.Web.Provider,
|
||||
"max_results": config.Web.MaxResults,
|
||||
}
|
||||
}
|
||||
|
||||
if config.KB != nil {
|
||||
result["kb"] = map[string]any{
|
||||
"threshold": config.KB.Threshold,
|
||||
"graph": config.KB.Graph,
|
||||
}
|
||||
}
|
||||
|
||||
if config.DB != nil {
|
||||
result["db"] = map[string]any{
|
||||
"max_results": config.DB.MaxResults,
|
||||
}
|
||||
}
|
||||
|
||||
if config.Weights != nil {
|
||||
result["weights"] = map[string]any{
|
||||
"user": config.Weights.User,
|
||||
"hook": config.Weights.Hook,
|
||||
"auto": config.Weights.Auto,
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
|
|
|||
|
|
@ -98,6 +98,24 @@ func init() {
|
|||
// KB: Chat collection
|
||||
"kb.chat.name": "Chat Knowledge Base",
|
||||
"kb.chat.description": "Auto-created knowledge base collection for chat sessions",
|
||||
|
||||
// Search: assistant/search.go - Output messages
|
||||
"search.loading": "Searching...",
|
||||
"search.success": "Found %d references",
|
||||
"search.success.one": "Found 1 reference",
|
||||
"search.partial": "Found %d references (some sources failed)",
|
||||
"search.failed": "Search failed",
|
||||
"search.no_results": "No references found",
|
||||
|
||||
// Search: assistant/search.go - Trace labels
|
||||
"search.trace.label": "Search",
|
||||
"search.trace.description": "Search the web and knowledge base for relevant information",
|
||||
"search.trace.web.label": "Web Search",
|
||||
"search.trace.web.description": "Searching the web",
|
||||
"search.trace.kb.label": "KB Search",
|
||||
"search.trace.kb.description": "Searching knowledge base",
|
||||
"search.trace.db.label": "DB Search",
|
||||
"search.trace.db.description": "Searching database",
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -164,6 +182,24 @@ func init() {
|
|||
// KB: Chat collection
|
||||
"kb.chat.name": "聊天知识库",
|
||||
"kb.chat.description": "自动为聊天会话创建的知识库集合",
|
||||
|
||||
// Search: assistant/search.go - Output messages
|
||||
"search.loading": "正在搜索...",
|
||||
"search.success": "找到 %d 条参考资料",
|
||||
"search.success.one": "找到 1 条参考资料",
|
||||
"search.partial": "找到 %d 条参考资料（部分来源失败）",
|
||||
"search.failed": "搜索失败",
|
||||
"search.no_results": "未找到相关资料",
|
||||
|
||||
// Search: assistant/search.go - Trace labels
|
||||
"search.trace.label": "搜索",
|
||||
"search.trace.description": "搜索网络和知识库获取相关信息",
|
||||
"search.trace.web.label": "网页搜索",
|
||||
"search.trace.web.description": "搜索网页获取相关信息",
|
||||
"search.trace.kb.label": "知识库搜索",
|
||||
"search.trace.kb.description": "搜索知识库获取相关信息",
|
||||
"search.trace.db.label": "数据库搜索",
|
||||
"search.trace.db.description": "搜索数据库获取相关信息",
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -258,6 +294,24 @@ func init() {
|
|||
// KB: Chat collection
|
||||
"kb.chat.name": "聊天知识库",
|
||||
"kb.chat.description": "自动为聊天会话创建的知识库集合",
|
||||
|
||||
// Search: assistant/search.go - Output messages
|
||||
"search.loading": "正在搜索...",
|
||||
"search.success": "找到 %d 条参考资料",
|
||||
"search.success.one": "找到 1 条参考资料",
|
||||
"search.partial": "找到 %d 条参考资料（部分来源失败）",
|
||||
"search.failed": "搜索失败",
|
||||
"search.no_results": "未找到相关资料",
|
||||
|
||||
// Search: assistant/search.go - Trace labels
|
||||
"search.trace.label": "搜索",
|
||||
"search.trace.description": "搜索网络和知识库获取相关信息",
|
||||
"search.trace.web.label": "网页搜索",
|
||||
"search.trace.web.description": "搜索网页获取相关信息",
|
||||
"search.trace.kb.label": "知识库搜索",
|
||||
"search.trace.kb.description": "搜索知识库获取相关信息",
|
||||
"search.trace.db.label": "数据库搜索",
|
||||
"search.trace.db.description": "搜索数据库获取相关信息",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -155,7 +155,7 @@ agent/search/
|
|||
│ │ ├── builtin.go # Builtin frequency-based extraction
|
||||
│ │ ├── agent.go # Agent mode (LLM-powered)
|
||||
│ │ └── mcp.go # MCP mode (external service)
|
||||
│ └── querydsl/ # QueryDSL generation for DB search (待实现)
|
||||
│ └── querydsl/ # QueryDSL generation for DB search (TODO)
|
||||
│ ├── generator.go # Main generator (mode dispatch)
|
||||
│ ├── builtin.go # Builtin template-based generation
|
||||
│ ├── agent.go # Agent mode (LLM-powered)
|
||||
|
|
@ -171,22 +171,21 @@ agent/search/
|
|||
│ │ ├── agent.go # Agent mode (AI Search)
|
||||
│ │ └── mcp.go # MCP mode (external service)
|
||||
│ │
|
||||
│ ├── kb/ # Knowledge base search (骨架)
|
||||
│ ├── kb/ # Knowledge base search (skeleton)
|
||||
│ │ ├── handler.go # KB search handler
|
||||
│ │ ├── vector.go # Vector similarity search (待实现)
|
||||
│ │ └── graph.go # Graph-based association (待实现)
|
||||
│ │ ├── vector.go # Vector similarity search (TODO)
|
||||
│ │ └── graph.go # Graph-based association (TODO)
|
||||
│ │
|
||||
│ └── db/ # Database search (骨架)
|
||||
│ └── db/ # Database search (skeleton)
|
||||
│ ├── handler.go # DB search handler
|
||||
│ ├── query.go # QueryDSL builder (待实现)
|
||||
│ └── schema.go # Model schema introspection (待实现)
|
||||
│ ├── query.go # QueryDSL builder (TODO)
|
||||
│ └── schema.go # Model schema introspection (TODO)
|
||||
│
|
||||
└── defaults/ # Default configuration values
|
||||
└── defaults.go # System built-in defaults (used by agent/load.go)
|
||||
|
||||
# 待实现文件:
|
||||
# - trace.go # Trace node creation and management
|
||||
# - output.go # Real-time output/streaming to client
|
||||
# Note: Output and Trace are integrated into assistant/search.go
|
||||
# No separate trace.go or output.go files needed
|
||||
```
|
||||
|
||||
### Dependency Graph
|
||||
|
|
@ -566,7 +565,7 @@ type RerankOptions struct {
|
|||
TopN int `json:"top_n,omitempty"` // Return top N after reranking
|
||||
}
|
||||
|
||||
// Result represents the search result
|
||||
// Result represents the search result with all intermediate processing data
|
||||
type Result struct {
|
||||
Type SearchType `json:"type"` // Search type
|
||||
Query string `json:"query"` // Original query
|
||||
|
|
@ -576,10 +575,31 @@ type Result struct {
|
|||
Duration int64 `json:"duration_ms"` // Search duration in ms
|
||||
Error string `json:"error,omitempty"` // Error message if failed
|
||||
|
||||
// Intermediate processing results (for storage and debugging)
|
||||
Keywords []string `json:"keywords,omitempty"` // Extracted keywords (Web/NLP)
|
||||
DSL map[string]any `json:"dsl,omitempty"` // Generated QueryDSL (DB)
|
||||
Entities []Entity `json:"entities,omitempty"` // Extracted entities (Graph RAG)
|
||||
Relations []Relation `json:"relations,omitempty"` // Extracted relations (Graph RAG)
|
||||
|
||||
// Graph associations (KB only, if enabled)
|
||||
GraphNodes []*GraphNode `json:"graph_nodes,omitempty"`
|
||||
}
|
||||
|
||||
// Entity represents an extracted entity (for Graph RAG)
|
||||
type Entity struct {
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
// Relation represents an extracted relation (for Graph RAG)
|
||||
type Relation struct {
|
||||
Subject string `json:"subject"`
|
||||
Predicate string `json:"predicate"`
|
||||
Object string `json:"object"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
// ResultItem represents a single search result item
|
||||
type ResultItem struct {
|
||||
// Citation
|
||||
|
|
@ -616,6 +636,43 @@ type ProcessedQuery struct {
|
|||
Vector []float32 `json:"vector,omitempty"` // For KB search
|
||||
DSL *gou.QueryDSL `json:"dsl,omitempty"` // For DB search, uses GOU QueryDSL
|
||||
}
|
||||
```
|
||||
|
||||
> **Design Note: Result with Intermediate Data**
|
||||
>
|
||||
> The `Result` type now includes intermediate processing results (`Keywords`, `DSL`, `Entities`, `Relations`)
|
||||
> that were previously only available during query processing. This design enables:
|
||||
>
|
||||
> 1. **Storage for Debugging**: All processing steps are captured for later analysis
|
||||
> 2. **System Tuning**: Analyze extracted keywords, generated DSL, and entity extraction quality
|
||||
> 3. **Unified Data Flow**: Handlers populate these fields during execution, eliminating the need
|
||||
> for separate data collection in `executeAutoSearch`
|
||||
>
|
||||
> **Handler Responsibilities**:
|
||||
>
|
||||
> - **Web Handler**: Populates `Keywords` from NLP extraction
|
||||
> - **DB Handler**: Populates `DSL` from QueryDSL generation
|
||||
> - **KB Handler**: Populates `Entities`, `Relations`, and `GraphNodes` from Graph RAG
|
||||
>
|
||||
> **Data Flow**:
|
||||
>
|
||||
> ```
|
||||
> Request → Handler → Result (with Keywords/DSL/Entities/Relations)
|
||||
> ↓
|
||||
> BuildReferenceContext
|
||||
> ↓
|
||||
> saveSearch (stores all intermediate data)
|
||||
> ```
|
||||
|
||||
```go
|
||||
// ProcessedQuery is DEPRECATED for external use
|
||||
// Handlers should populate Result.Keywords/DSL/Entities/Relations directly
|
||||
type ProcessedQuery struct {
|
||||
Type SearchType `json:"type"`
|
||||
Keywords []string `json:"keywords,omitempty"` // For web search
|
||||
Vector []float32 `json:"vector,omitempty"` // For KB search
|
||||
DSL *gou.QueryDSL `json:"dsl,omitempty"` // For DB search
|
||||
}
|
||||
|
||||
// Note: For QueryDSL and Model types, use GOU types directly:
|
||||
// - github.com/yaoapp/gou/query/gou.QueryDSL
|
||||
|
|
@ -836,51 +893,666 @@ search:
|
|||
|
||||
## Trace Integration
|
||||
|
||||
Search operations create trace nodes to report execution details to users, providing transparency about what the agent is doing.
|
||||
Search operations create minimal trace nodes to report execution status to users, providing transparency about what the agent is doing. Detailed information is recorded via LOG for debugging.
|
||||
|
||||
### Trace Node Structure
|
||||
|
||||
Uses `trace/types.NodeStatus` constants:
|
||||
|
||||
- `pending` - Node created but not started
|
||||
- `running` - Node is currently executing
|
||||
- `completed` - Node finished successfully
|
||||
- `failed` - Node failed with error
|
||||
|
||||
**Single Search:**
|
||||
|
||||
```
|
||||
search (type: "search")
|
||||
├── query // Original query
|
||||
├── search_type // "web", "kb", or "db"
|
||||
├── duration_ms
|
||||
├── status // "success", "failed"
|
||||
├── result_count
|
||||
└── children // Sub-operations
|
||||
├── embedding (kb only)
|
||||
├── vector_search (kb only)
|
||||
├── graph_search (kb, if enabled)
|
||||
├── querydsl_build (db only)
|
||||
├── db_query (db only)
|
||||
└── rerank (if enabled)
|
||||
├── label // i18n: "Search" / "搜索"
|
||||
├── status // "pending" | "running" | "completed" | "failed"
|
||||
├── input
|
||||
│ ├── query // Original query
|
||||
│ └── types // ["web"], ["kb"], ["web", "kb", "db"]
|
||||
└── output // (set on complete)
|
||||
└── result_count // Total results found
|
||||
```
|
||||
|
||||
**Parallel Search:**
|
||||
|
||||
```
|
||||
search (type: "search")
|
||||
├── label // i18n: "Search" / "搜索"
|
||||
├── status // "pending" | "running" | "completed" | "failed"
|
||||
├── input
|
||||
│ ├── query // Original query
|
||||
│ └── types // ["web", "kb", "db"]
|
||||
└── children
|
||||
├── web (type: "search_item")
|
||||
│ ├── label // i18n: "Web Search" / "网页搜索"
|
||||
│ ├── status // "pending" | "running" | "completed" | "failed"
|
||||
│ └── output
|
||||
│ └── result_count
|
||||
├── kb (type: "search_item")
|
||||
│ └── ...
|
||||
└── db (type: "search_item")
|
||||
└── ...
|
||||
```
|
||||
|
||||
### Trace Logging
|
||||
|
||||
Detailed search information is recorded via Trace node logging methods (broadcasts to client):
|
||||
|
||||
```go
|
||||
// Node logging methods (from trace/node.go):
|
||||
// - node.Info(message, args...) - Info level log
|
||||
// - node.Debug(message, args...) - Debug level log
|
||||
// - node.Warn(message, args...) - Warning level log
|
||||
// - node.Error(message, args...) - Error level log
|
||||
|
||||
// Search start
|
||||
searchNode.Info("Starting search", map[string]any{"query": query, "types": types})
|
||||
|
||||
// Per-type results (on parallel search children)
|
||||
webNode.Debug("Web search completed", map[string]any{"count": count, "duration_ms": duration})
|
||||
kbNode.Debug("KB search completed", map[string]any{"count": count, "duration_ms": duration})
|
||||
dbNode.Debug("DB search completed", map[string]any{"count": count, "duration_ms": duration})
|
||||
|
||||
// Errors (non-blocking, search continues)
|
||||
webNode.Warn("Web search failed", map[string]any{"error": err.Error()})
|
||||
|
||||
// Final summary (on parent node)
|
||||
searchNode.Info("Search completed", map[string]any{"total": total, "duration_ms": duration})
|
||||
```
|
||||
|
||||
**Log Event Structure** (broadcasted via SSE):
|
||||
|
||||
```go
|
||||
// types.TraceLog
|
||||
type TraceLog struct {
|
||||
Timestamp int64 `json:"timestamp"` // milliseconds since epoch
|
||||
Level string `json:"level"` // "info", "debug", "warn", "error"
|
||||
Message string `json:"message"` // Log message
|
||||
Data any `json:"data"` // Additional data
|
||||
NodeID string `json:"node_id"` // Parent node ID
|
||||
}
|
||||
```
|
||||
|
||||
## Real-time Output
|
||||
|
||||
Search progress is streamed to the client via the output system.
|
||||
Search progress is displayed to the client using **Loading component with Replace** pattern. Uses `ctx.Send()` and `ctx.Replace()` methods.
|
||||
|
||||
### Output Message Types
|
||||
### Output Flow
|
||||
|
||||
```
|
||||
1. Send Loading Message
|
||||
loading_id = ctx.Send({ type: "loading", props: { message: "Searching..." } })
|
||||
→ Client displays loading indicator
|
||||
|
||||
2. Execute Search (parallel web/kb/db)
|
||||
|
||||
3. Replace with Result Message (shows result to user)
|
||||
ctx.Replace(loading_id, { type: "loading", props: { message: "Found 5 references" } })
|
||||
→ Client displays result message
|
||||
|
||||
4. Mark as Done (removes the loading after brief display)
|
||||
ctx.Replace(loading_id, { type: "loading", props: { message: "Found 5 references", done: true } })
|
||||
→ Client removes loading indicator
|
||||
```
|
||||
|
||||
### Implementation
|
||||
|
||||
```go
|
||||
const (
|
||||
TypeSearchStart = "search_start" // Search initiated
|
||||
TypeSearchResult = "search_result" // Result item (streamed)
|
||||
TypeSearchComplete = "search_complete" // Search completed
|
||||
)
|
||||
// Send loading message
|
||||
loadingID := ctx.Send(map[string]any{
|
||||
"type": "loading",
|
||||
"props": map[string]any{
|
||||
"message": i18n.Tr("search.loading", locale), // "Searching..." / "正在搜索..."
|
||||
},
|
||||
})
|
||||
|
||||
// Execute search...
|
||||
|
||||
// Replace with result message (displayed to user)
|
||||
resultMessage := i18n.Tr("search.success", locale, count) // "Found 5 references"
|
||||
ctx.Replace(loadingID, map[string]any{
|
||||
"type": "loading",
|
||||
"props": map[string]any{
|
||||
"message": resultMessage,
|
||||
},
|
||||
})
|
||||
|
||||
// Mark as done (removes loading indicator after user sees the result)
|
||||
ctx.Replace(loadingID, map[string]any{
|
||||
"type": "loading",
|
||||
"props": map[string]any{
|
||||
"message": resultMessage,
|
||||
"done": true, // Frontend will remove loading indicator
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
### Loading Props
|
||||
|
||||
| Prop | Type | Description |
|
||||
| --------- | ------ | --------------------------------------------------- |
|
||||
| `message` | string | Localized message to display |
|
||||
| `done` | bool | When `true`, frontend removes the loading indicator |
|
||||
|
||||
### Localized Messages
|
||||
|
||||
| Scenario | English | Chinese |
|
||||
| ------------- | ---------------------------------------- | --------------------------------- |
|
||||
| Loading | Searching... | 正在搜索... |
|
||||
| Success (1) | Found 1 reference | 找到 1 条参考资料 |
|
||||
| Success (N) | Found N references | 找到 N 条参考资料 |
|
||||
| Partial Error | Found N references (some sources failed) | 找到 N 条参考资料（部分来源失败） |
|
||||
| All Failed | Search failed | 搜索失败 |
|
||||
| No Results | No references found | 未找到相关资料 |
|
||||
|
||||
### Client Display Example
|
||||
|
||||
```
|
||||
🔍 Searching "latest AI developments"...
|
||||
Frame 1 - During search:
|
||||
┌─────────────────────────────────┐
|
||||
│ Searching... │ ← Loading (done: false)
|
||||
└─────────────────────────────────┘
|
||||
|
||||
📄 Found 5 results:
|
||||
1. #ref:a1b2 - OpenAI Announces GPT-5
|
||||
2. #ref:c3d4 - Google's New AI Model
|
||||
...
|
||||
Frame 2 - Result displayed:
|
||||
┌─────────────────────────────────┐
|
||||
│ Found 5 references │ ← Result (done: false)
|
||||
└─────────────────────────────────┘
|
||||
|
||||
✅ Search complete (1.2s)
|
||||
Frame 3 - Removed:
|
||||
(loading indicator removed when done: true)
|
||||
```
|
||||
|
||||
## Search Result Storage
|
||||
|
||||
Search results are stored per request to support citation click-through and history replay.
|
||||
|
||||
### Data Model
|
||||
|
||||
```
|
||||
Relationships:
|
||||
Chat
|
||||
└── Request (request_id)
|
||||
├── Message[] (user, assistant, tool...)
|
||||
└── SearchResult[] (one request may have multiple searches)
|
||||
└── Reference[] (indexed references from each search)
|
||||
```
|
||||
|
||||
### Citation Locating
|
||||
|
||||
LLM output uses `<a>` tags with index:
|
||||
|
||||
```xml
|
||||
AI is artificial intelligence<a index="1" />, it has developed rapidly<a index="2" />...
|
||||
```
|
||||
|
||||
Location path: `request_id` + `index` → precisely locate reference
|
||||
|
||||
### Database Schema
|
||||
|
||||
**Table: `agent_search`**
|
||||
|
||||
| Column | Type | Description |
|
||||
| ---------- | ----------- | -------------------------------------- |
|
||||
| id | BIGINT | Auto-increment primary key |
|
||||
| request_id | VARCHAR(64) | Associated request ID (indexed) |
|
||||
| chat_id | VARCHAR(64) | Associated chat ID (indexed) |
|
||||
| query | TEXT | Original search query |
|
||||
| config | JSON | Search config used (for tuning) |
|
||||
| keywords | JSON | Extracted keywords (from NLP) |
|
||||
| entities | JSON | Extracted entities (for Graph search) |
|
||||
| relations | JSON | Extracted relations (for Graph search) |
|
||||
| dsl | JSON | Generated QueryDSL (for DB search) |
|
||||
| source | VARCHAR(32) | Search source: web/kb/db/auto |
|
||||
| references | JSON | Reference[] with global index |
|
||||
| graph | JSON | GraphNode[] from knowledge graph |
|
||||
| xml | TEXT | Formatted XML for LLM context |
|
||||
| prompt | TEXT | Citation instruction prompt |
|
||||
| duration | INT | Search duration in milliseconds |
|
||||
| error | TEXT | Error message if failed (nullable) |
|
||||
| created_at | TIMESTAMP | Creation time |
|
||||
| deleted_at | TIMESTAMP | Soft delete time (nullable) |
|
||||
|
||||
**Config Field Structure:**
|
||||
|
||||
```json
|
||||
{
|
||||
"uses": {
|
||||
"search": "builtin",
|
||||
"web": "builtin",
|
||||
"keyword": "builtin",
|
||||
"querydsl": "builtin",
|
||||
"rerank": "builtin"
|
||||
},
|
||||
"web": {
|
||||
"provider": "tavily",
|
||||
"max_results": 5
|
||||
},
|
||||
"kb": {
|
||||
"collections": ["docs", "faq"],
|
||||
"threshold": 0.7,
|
||||
"graph": true
|
||||
},
|
||||
"db": {
|
||||
"models": ["product", "order"],
|
||||
"max_results": 20
|
||||
},
|
||||
"rerank": {
|
||||
"provider": "builtin",
|
||||
"top_n": 10
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Type Definitions
|
||||
|
||||
```go
|
||||
// store/types/types.go
|
||||
|
||||
// Search represents stored search results for a request
|
||||
// Stores all intermediate processing results for debugging and replay
|
||||
type Search struct {
|
||||
ID int64 `json:"id"`
|
||||
RequestID string `json:"request_id"`
|
||||
ChatID string `json:"chat_id"`
|
||||
Query string `json:"query"` // Original query
|
||||
Config map[string]any `json:"config,omitempty"` // Search config used (for tuning)
|
||||
Keywords []string `json:"keywords,omitempty"` // Extracted keywords (Web/NLP)
|
||||
Entities []Entity `json:"entities,omitempty"` // Extracted entities (Graph)
|
||||
Relations []Relation `json:"relations,omitempty"` // Extracted relations (Graph)
|
||||
DSL map[string]any `json:"dsl,omitempty"` // Generated QueryDSL (DB)
|
||||
Source string `json:"source"` // web/kb/db/auto
|
||||
References []Reference `json:"references"`
|
||||
Graph []GraphNode `json:"graph,omitempty"` // Graph nodes from KB
|
||||
XML string `json:"xml,omitempty"` // Formatted XML for LLM
|
||||
Prompt string `json:"prompt,omitempty"` // Citation prompt
|
||||
Duration int64 `json:"duration_ms"` // Search duration
|
||||
Error string `json:"error,omitempty"` // Error if failed
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// Reference represents a single reference with global index (for storage)
|
||||
type Reference struct {
|
||||
Index int `json:"index"` // Global index: 1, 2, 3...
|
||||
Type string `json:"type"` // web/kb/db
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url,omitempty"`
|
||||
Snippet string `json:"snippet"`
|
||||
Content string `json:"content,omitempty"` // Full content (optional)
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// Entity represents an extracted entity from query (for Graph search)
|
||||
type Entity struct {
|
||||
Name string `json:"name"` // Entity name
|
||||
Type string `json:"type"` // Entity type: person, org, location, etc.
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// Relation represents an extracted relation from query (for Graph search)
|
||||
type Relation struct {
|
||||
Subject string `json:"subject"` // Source entity
|
||||
Predicate string `json:"predicate"` // Relation type
|
||||
Object string `json:"object"` // Target entity
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// GraphNode represents a node from knowledge graph (search result)
|
||||
type GraphNode struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // Entity type
|
||||
Name string `json:"name"` // Entity name
|
||||
Description string `json:"description,omitempty"`
|
||||
Relation string `json:"relation,omitempty"` // Relationship to query
|
||||
Score float64 `json:"score,omitempty"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// SearchFilter for querying search records
|
||||
type SearchFilter struct {
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
ChatID string `json:"chat_id,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
```
|
||||
|
||||
### Store Interface Extension
|
||||
|
||||
```go
|
||||
// store/types/store.go
|
||||
|
||||
// ChatStore interface extension
|
||||
type ChatStore interface {
|
||||
// ... existing methods ...
|
||||
|
||||
// ==========================================================================
|
||||
// Search Management
|
||||
// ==========================================================================
|
||||
|
||||
// SaveSearch saves search record for a request
|
||||
// search: Search record to save
|
||||
// Returns: Potential error
|
||||
SaveSearch(search *Search) error
|
||||
|
||||
// GetSearches retrieves search records for a request
|
||||
// requestID: Request ID
|
||||
// Returns: Search records and potential error
|
||||
GetSearches(requestID string) ([]*Search, error)
|
||||
|
||||
// GetReference retrieves a single reference by request ID and index
|
||||
// requestID: Request ID
|
||||
// index: Reference index (1-based)
|
||||
// Returns: Reference and potential error
|
||||
GetReference(requestID string, index int) (*Reference, error)
|
||||
|
||||
// DeleteSearches deletes all search records for a chat
|
||||
// chatID: Chat ID
|
||||
// Returns: Potential error
|
||||
DeleteSearches(chatID string) error
|
||||
}
|
||||
```
|
||||
|
||||
### Xun Implementation
|
||||
|
||||
```go
|
||||
// store/xun/search.go
|
||||
|
||||
// SaveSearch saves a search record
|
||||
func (store *Xun) SaveSearch(search *Search) error {
|
||||
if search.RequestID == "" {
|
||||
return fmt.Errorf("request_id is required")
|
||||
}
|
||||
|
||||
refsJSON, err := jsoniter.MarshalToString(search.References)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal references: %w", err)
|
||||
}
|
||||
|
||||
row := map[string]interface{}{
|
||||
"request_id": search.RequestID,
|
||||
"chat_id": search.ChatID,
|
||||
"query": search.Query,
|
||||
"config": search.Config, // Search config for tuning
|
||||
"keywords": search.Keywords,
|
||||
"entities": search.Entities, // Graph entities
|
||||
"relations": search.Relations, // Graph relations
|
||||
"dsl": search.DSL,
|
||||
"source": search.Source,
|
||||
"references": refsJSON,
|
||||
"graph": search.Graph, // Graph nodes
|
||||
"xml": search.XML,
|
||||
"prompt": search.Prompt,
|
||||
"duration": search.Duration,
|
||||
"error": search.Error,
|
||||
"created_at": time.Now(),
|
||||
}
|
||||
|
||||
return store.newQuerySearch().Insert(row)
|
||||
}
|
||||
|
||||
// GetSearches retrieves all search records for a request
|
||||
func (store *Xun) GetSearches(requestID string) ([]*Search, error) {
|
||||
rows, err := store.newQuerySearch().
|
||||
Where("request_id", requestID).
|
||||
WhereNull("deleted_at").
|
||||
OrderBy("created_at", "asc").
|
||||
Get()
|
||||
// ... convert rows to Search
|
||||
}
|
||||
|
||||
// GetReference retrieves a single reference
|
||||
func (store *Xun) GetReference(requestID string, index int) (*Reference, error) {
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Find reference by index across all search records
|
||||
for _, search := range searches {
|
||||
for _, ref := range search.References {
|
||||
if ref.Index == index {
|
||||
return &ref, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("reference %d not found in request %s", index, requestID)
|
||||
}
|
||||
```
|
||||
|
||||
### Model Definition
|
||||
|
||||
```json
|
||||
// yao/models/agent/search.mod.yao
|
||||
{
|
||||
"name": "Search",
|
||||
"label": "Search",
|
||||
"description": "Search records for citation support and debugging",
|
||||
"tags": ["agent", "system"],
|
||||
"builtin": true,
|
||||
"readonly": true,
|
||||
"table": {
|
||||
"name": "agent_search",
|
||||
"comment": "Agent search table"
|
||||
},
|
||||
"columns": [
|
||||
{ "name": "id", "type": "ID", "label": "ID" },
|
||||
{
|
||||
"name": "request_id",
|
||||
"type": "string",
|
||||
"length": 64,
|
||||
"nullable": false,
|
||||
"index": true
|
||||
},
|
||||
{
|
||||
"name": "chat_id",
|
||||
"type": "string",
|
||||
"length": 64,
|
||||
"nullable": false,
|
||||
"index": true
|
||||
},
|
||||
{ "name": "query", "type": "text", "nullable": true },
|
||||
{
|
||||
"name": "config",
|
||||
"type": "json",
|
||||
"nullable": true,
|
||||
"comment": "Search config used (for tuning)"
|
||||
},
|
||||
{ "name": "keywords", "type": "json", "nullable": true },
|
||||
{ "name": "entities", "type": "json", "nullable": true },
|
||||
{ "name": "relations", "type": "json", "nullable": true },
|
||||
{ "name": "dsl", "type": "json", "nullable": true },
|
||||
{ "name": "source", "type": "string", "length": 32, "nullable": false },
|
||||
{ "name": "references", "type": "json", "nullable": true },
|
||||
{ "name": "graph", "type": "json", "nullable": true },
|
||||
{ "name": "xml", "type": "text", "nullable": true },
|
||||
{ "name": "prompt", "type": "text", "nullable": true },
|
||||
{ "name": "duration", "type": "integer", "nullable": true },
|
||||
{ "name": "error", "type": "text", "nullable": true }
|
||||
],
|
||||
"option": { "timestamps": true, "soft_deletes": true }
|
||||
}
|
||||
```
|
||||
|
||||
### Stream Integration
|
||||
|
||||
Storage logic is encapsulated in `assistant/search.go` with a dedicated method:
|
||||
|
||||
```go
|
||||
// assistant/search.go
|
||||
|
||||
// SearchExecutionResult contains all intermediate results from search execution
|
||||
type SearchExecutionResult struct {
|
||||
Query string // Original query
|
||||
Config map[string]any // Search config used
|
||||
Keywords []string // Extracted keywords (Web/NLP)
|
||||
Entities []storeTypes.Entity // Extracted entities (Graph)
|
||||
Relations []storeTypes.Relation // Extracted relations (Graph)
|
||||
DSL map[string]any // Generated QueryDSL (DB)
|
||||
Source string // web/kb/db/auto
|
||||
RefCtx *searchTypes.ReferenceContext // Reference context for LLM
|
||||
Graph []storeTypes.GraphNode // Graph nodes from KB
|
||||
Duration int64 // Duration in ms
|
||||
Error string // Error message if failed
|
||||
}
|
||||
|
||||
// saveSearch saves search record to store for citation support and debugging
|
||||
func (ast *Assistant) saveSearch(ctx *context.Context, result *SearchExecutionResult) {
|
||||
if ctx.Store == nil || result == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Skip if no references and no error
|
||||
if result.RefCtx == nil && result.Error == "" {
|
||||
return
|
||||
}
|
||||
|
||||
var refs []storeTypes.Reference
|
||||
var xml, prompt string
|
||||
|
||||
if result.RefCtx != nil {
|
||||
refs = convertReferences(result.RefCtx.References)
|
||||
xml = result.RefCtx.XML
|
||||
prompt = result.RefCtx.Prompt
|
||||
}
|
||||
|
||||
search := &storeTypes.Search{
|
||||
RequestID: ctx.RequestID,
|
||||
ChatID: ctx.ID,
|
||||
Query: result.Query,
|
||||
Config: result.Config, // Search config for tuning analysis
|
||||
Keywords: result.Keywords,
|
||||
Entities: result.Entities, // Graph entities
|
||||
Relations: result.Relations, // Graph relations
|
||||
DSL: result.DSL,
|
||||
Source: result.Source,
|
||||
References: refs,
|
||||
Graph: result.Graph, // Graph nodes
|
||||
XML: xml,
|
||||
Prompt: prompt,
|
||||
Duration: result.Duration,
|
||||
Error: result.Error,
|
||||
}
|
||||
|
||||
if err := ctx.Store.SaveSearch(search); err != nil {
|
||||
ctx.Logger.Warn("Failed to save search: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// convertReferences converts search references to store format
|
||||
func convertReferences(refs []*searchTypes.Reference) []storeTypes.Reference {
|
||||
result := make([]storeTypes.Reference, len(refs))
|
||||
for i, ref := range refs {
|
||||
result[i] = storeTypes.Reference{
|
||||
Index: i + 1, // 1-based index
|
||||
Type: string(ref.Type),
|
||||
Title: ref.Title,
|
||||
URL: ref.URL,
|
||||
Snippet: ref.Content,
|
||||
Content: ref.Content,
|
||||
Metadata: ref.Meta,
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// In executeAutoSearch:
|
||||
func (ast *Assistant) executeAutoSearch(ctx *context.Context, ...) *searchTypes.ReferenceContext {
|
||||
start := time.Now()
|
||||
|
||||
// 1. Execute search (Result now contains all intermediate data)
|
||||
results, err := searcher.All(ctx, requests)
|
||||
duration := time.Since(start).Milliseconds()
|
||||
|
||||
// 2. Prepare execution result for storage
|
||||
execResult := &SearchExecutionResult{
|
||||
Query: query,
|
||||
Config: buildSearchConfig(searchConfig, searchUses),
|
||||
Source: "auto",
|
||||
Duration: duration,
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
execResult.Error = err.Error()
|
||||
ast.saveSearch(ctx, execResult)
|
||||
return nil
|
||||
}
|
||||
|
||||
// 3. Extract intermediate data from results
|
||||
// Result.Keywords, Result.DSL, Result.Entities, Result.Relations are populated by handlers
|
||||
for _, result := range results {
|
||||
if len(result.Keywords) > 0 {
|
||||
execResult.Keywords = result.Keywords
|
||||
}
|
||||
if result.DSL != nil {
|
||||
execResult.DSL = result.DSL
|
||||
}
|
||||
if len(result.Entities) > 0 {
|
||||
execResult.Entities = convertEntities(result.Entities)
|
||||
}
|
||||
if len(result.Relations) > 0 {
|
||||
execResult.Relations = convertRelations(result.Relations)
|
||||
}
|
||||
if len(result.GraphNodes) > 0 {
|
||||
execResult.Graph = convertGraphNodes(result.GraphNodes)
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Build reference context
|
||||
refCtx := search.BuildReferenceContext(results, citationConfig)
|
||||
execResult.RefCtx = refCtx
|
||||
|
||||
// 5. Save search record
|
||||
ast.saveSearch(ctx, execResult)
|
||||
|
||||
return refCtx
|
||||
}
|
||||
```
|
||||
|
||||
### Usage Scenarios
|
||||
|
||||
**Scenario 1: Single Search**
|
||||
|
||||
```
|
||||
Request: req_001
|
||||
└── Search: { source: "auto", references: [{index:1,...}, {index:2,...}, {index:3,...}] }
|
||||
```
|
||||
|
||||
**Scenario 2: Multiple Searches (e.g., Tool Call triggers another search)**
|
||||
|
||||
```
|
||||
Request: req_001
|
||||
├── Search[0]: { source: "web", references: [{index:1,...}, {index:2,...}] }
|
||||
└── Search[1]: { source: "kb", references: [{index:3,...}, {index:4,...}] }
|
||||
```
|
||||
|
||||
Index is globally incremented, so `request_id + index` is always unique.
|
||||
|
||||
### API Endpoints
|
||||
|
||||
```
|
||||
GET /api/chat/{chat_id}/request/{request_id}/references # Get all references for request
|
||||
GET /api/chat/{chat_id}/request/{request_id}/reference/{index} # Get single reference by index
|
||||
```
|
||||
|
||||
### Frontend Integration
|
||||
|
||||
```typescript
|
||||
// When user clicks citation [1]
|
||||
async function onCitationClick(requestId: string, index: number) {
|
||||
const ref = await api.get(
|
||||
`/chat/${chatId}/request/${requestId}/reference/${index}`
|
||||
);
|
||||
showReferenceCard({
|
||||
title: ref.title,
|
||||
url: ref.url,
|
||||
snippet: ref.snippet,
|
||||
content: ref.content,
|
||||
});
|
||||
}
|
||||
```
|
||||
|
||||
## JSAPI Integration
|
||||
|
|
@ -1978,7 +2650,7 @@ SerpAPI supports multiple search engines via the `engine` config:
|
|||
| ------------ | ---------------------------- |
|
||||
| `google` | Google Search (default) |
|
||||
| `bing` | Bing Search |
|
||||
| `baidu` | Baidu (百度) |
|
||||
| `baidu` | Baidu Search (Chinese) |
|
||||
| `yandex` | Yandex Search |
|
||||
| `yahoo` | Yahoo Search |
|
||||
| `duckduckgo` | DuckDuckGo Search |
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
package search
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// CitationGenerator generates unique citation IDs
|
||||
// CitationGenerator generates unique citation IDs (1-based integers)
|
||||
// Thread-safe for concurrent use within a single request
|
||||
type CitationGenerator struct {
|
||||
counter uint64
|
||||
}
|
||||
|
|
@ -15,13 +15,38 @@ func NewCitationGenerator() *CitationGenerator {
|
|||
return &CitationGenerator{}
|
||||
}
|
||||
|
||||
// Next generates the next citation ID
|
||||
// Next generates the next citation ID (1, 2, 3, ...)
|
||||
func (g *CitationGenerator) Next() string {
|
||||
n := atomic.AddUint64(&g.counter, 1)
|
||||
return fmt.Sprintf("ref_%03d", n)
|
||||
return uint64ToString(n)
|
||||
}
|
||||
|
||||
// NextInt generates the next citation ID as integer
|
||||
func (g *CitationGenerator) NextInt() int {
|
||||
return int(atomic.AddUint64(&g.counter, 1))
|
||||
}
|
||||
|
||||
// Current returns the current counter value without incrementing
|
||||
func (g *CitationGenerator) Current() int {
|
||||
return int(atomic.LoadUint64(&g.counter))
|
||||
}
|
||||
|
||||
// Reset resets the counter (for testing)
|
||||
func (g *CitationGenerator) Reset() {
|
||||
atomic.StoreUint64(&g.counter, 0)
|
||||
}
|
||||
|
||||
// uint64ToString converts uint64 to string without fmt package
|
||||
func uint64ToString(n uint64) string {
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
var buf [20]byte // max uint64 is 20 digits
|
||||
i := len(buf)
|
||||
for n > 0 {
|
||||
i--
|
||||
buf[i] = byte('0' + n%10)
|
||||
n /= 10
|
||||
}
|
||||
return string(buf[i:])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,17 +10,43 @@ import (
|
|||
func TestCitationGenerator_Next(t *testing.T) {
|
||||
gen := NewCitationGenerator()
|
||||
|
||||
// First ID should be ref_001
|
||||
// First ID should be "1"
|
||||
id1 := gen.Next()
|
||||
assert.Equal(t, "ref_001", id1)
|
||||
assert.Equal(t, "1", id1)
|
||||
|
||||
// Second ID should be ref_002
|
||||
// Second ID should be "2"
|
||||
id2 := gen.Next()
|
||||
assert.Equal(t, "ref_002", id2)
|
||||
assert.Equal(t, "2", id2)
|
||||
|
||||
// Third ID should be ref_003
|
||||
// Third ID should be "3"
|
||||
id3 := gen.Next()
|
||||
assert.Equal(t, "ref_003", id3)
|
||||
assert.Equal(t, "3", id3)
|
||||
}
|
||||
|
||||
func TestCitationGenerator_NextInt(t *testing.T) {
|
||||
gen := NewCitationGenerator()
|
||||
|
||||
// First ID should be 1
|
||||
id1 := gen.NextInt()
|
||||
assert.Equal(t, 1, id1)
|
||||
|
||||
// Second ID should be 2
|
||||
id2 := gen.NextInt()
|
||||
assert.Equal(t, 2, id2)
|
||||
}
|
||||
|
||||
func TestCitationGenerator_Current(t *testing.T) {
|
||||
gen := NewCitationGenerator()
|
||||
|
||||
// Initial should be 0
|
||||
assert.Equal(t, 0, gen.Current())
|
||||
|
||||
// After one Next, should be 1
|
||||
gen.Next()
|
||||
assert.Equal(t, 1, gen.Current())
|
||||
|
||||
// Current doesn't increment
|
||||
assert.Equal(t, 1, gen.Current())
|
||||
}
|
||||
|
||||
func TestCitationGenerator_Reset(t *testing.T) {
|
||||
|
|
@ -34,22 +60,22 @@ func TestCitationGenerator_Reset(t *testing.T) {
|
|||
// Reset
|
||||
gen.Reset()
|
||||
|
||||
// Next ID should be ref_001 again
|
||||
// Next ID should be "1" again
|
||||
id := gen.Next()
|
||||
assert.Equal(t, "ref_001", id)
|
||||
assert.Equal(t, "1", id)
|
||||
}
|
||||
|
||||
func TestCitationGenerator_Format(t *testing.T) {
|
||||
func TestCitationGenerator_LargeNumbers(t *testing.T) {
|
||||
gen := NewCitationGenerator()
|
||||
|
||||
// Generate 999 IDs to test padding
|
||||
// Generate 999 IDs
|
||||
for i := 0; i < 999; i++ {
|
||||
gen.Next()
|
||||
}
|
||||
|
||||
// 1000th ID should be ref_1000 (no padding limit)
|
||||
// 1000th ID should be "1000"
|
||||
id := gen.Next()
|
||||
assert.Equal(t, "ref_1000", id)
|
||||
assert.Equal(t, "1000", id)
|
||||
}
|
||||
|
||||
func TestCitationGenerator_Concurrent(t *testing.T) {
|
||||
|
|
@ -86,3 +112,23 @@ func TestNewCitationGenerator(t *testing.T) {
|
|||
gen := NewCitationGenerator()
|
||||
assert.NotNil(t, gen)
|
||||
}
|
||||
|
||||
func TestUint64ToString(t *testing.T) {
|
||||
tests := []struct {
|
||||
input uint64
|
||||
expected string
|
||||
}{
|
||||
{0, "0"},
|
||||
{1, "1"},
|
||||
{10, "10"},
|
||||
{100, "100"},
|
||||
{999, "999"},
|
||||
{1000, "1000"},
|
||||
{18446744073709551615, "18446744073709551615"}, // max uint64
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
result := uint64ToString(tt.input)
|
||||
assert.Equal(t, tt.expected, result, "uint64ToString(%d)", tt.input)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import (
|
|||
|
||||
// DefaultCitationPrompt is the default prompt for citation instructions
|
||||
const DefaultCitationPrompt = `You have access to reference data in <references> tags. Each <ref> has:
|
||||
- id: Citation identifier
|
||||
- id: Citation identifier (integer)
|
||||
- type: Data type (web/kb/db)
|
||||
- weight: Relevance weight (1.0=highest priority, 0.6=lowest)
|
||||
- source: Origin (user=user-provided, hook=assistant-searched, auto=auto-searched)
|
||||
|
|
@ -19,7 +19,7 @@ Prioritize higher-weight references when answering.
|
|||
When citing a reference, use this exact HTML format:
|
||||
<a class="ref" data-ref-id="{id}" data-ref-type="{type}" href="#ref:{id}">[{id}]</a>
|
||||
|
||||
Example: According to the product data<a class="ref" data-ref-id="ref_001" data-ref-type="db" href="#ref:ref_001">[ref_001]</a>, the price is $999.`
|
||||
Example: According to the product data<a class="ref" data-ref-id="1" data-ref-type="db" href="#ref:1">[1]</a>, the price is $999.`
|
||||
|
||||
// BuildReferences converts search results to unified Reference format
|
||||
func BuildReferences(results []*types.Result) []*types.Reference {
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ func TestBuildReferences(t *testing.T) {
|
|||
Query: "test query",
|
||||
Items: []*types.ResultItem{
|
||||
{
|
||||
CitationID: "ref_001",
|
||||
CitationID: "1",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Weight: 0.6,
|
||||
|
|
@ -42,7 +42,7 @@ func TestBuildReferences(t *testing.T) {
|
|||
URL: "https://example.com",
|
||||
},
|
||||
{
|
||||
CitationID: "ref_002",
|
||||
CitationID: "2",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Weight: 0.6,
|
||||
|
|
@ -62,19 +62,19 @@ func TestBuildReferences(t *testing.T) {
|
|||
{
|
||||
Type: types.SearchTypeWeb,
|
||||
Items: []*types.ResultItem{
|
||||
{CitationID: "ref_001", Type: types.SearchTypeWeb, Content: "Web content"},
|
||||
{CitationID: "1", Type: types.SearchTypeWeb, Content: "Web content"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: types.SearchTypeKB,
|
||||
Items: []*types.ResultItem{
|
||||
{CitationID: "ref_002", Type: types.SearchTypeKB, Content: "KB content"},
|
||||
{CitationID: "2", Type: types.SearchTypeKB, Content: "KB content"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: types.SearchTypeDB,
|
||||
Items: []*types.ResultItem{
|
||||
{CitationID: "ref_003", Type: types.SearchTypeDB, Content: "DB content"},
|
||||
{CitationID: "3", Type: types.SearchTypeDB, Content: "DB content"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
|
@ -86,9 +86,9 @@ func TestBuildReferences(t *testing.T) {
|
|||
{
|
||||
Type: types.SearchTypeWeb,
|
||||
Items: []*types.ResultItem{
|
||||
{CitationID: "ref_001", Content: "Content 1"},
|
||||
{CitationID: "1", Content: "Content 1"},
|
||||
nil,
|
||||
{CitationID: "ref_002", Content: "Content 2"},
|
||||
{CitationID: "2", Content: "Content 2"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
|
@ -100,14 +100,14 @@ func TestBuildReferences(t *testing.T) {
|
|||
{
|
||||
Type: types.SearchTypeWeb,
|
||||
Items: []*types.ResultItem{
|
||||
{CitationID: "ref_001", Content: "Content"},
|
||||
{CitationID: "1", Content: "Content"},
|
||||
},
|
||||
},
|
||||
nil,
|
||||
{
|
||||
Type: types.SearchTypeKB,
|
||||
Items: []*types.ResultItem{
|
||||
{CitationID: "ref_002", Content: "Content 2"},
|
||||
{CitationID: "2", Content: "Content 2"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
|
@ -125,7 +125,7 @@ func TestBuildReferences(t *testing.T) {
|
|||
|
||||
func TestBuildReferences_FieldMapping(t *testing.T) {
|
||||
item := &types.ResultItem{
|
||||
CitationID: "ref_001",
|
||||
CitationID: "1",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceHook,
|
||||
Weight: 0.8,
|
||||
|
|
@ -143,7 +143,7 @@ func TestBuildReferences_FieldMapping(t *testing.T) {
|
|||
assert.Equal(t, 1, len(refs))
|
||||
|
||||
ref := refs[0]
|
||||
assert.Equal(t, "ref_001", ref.ID)
|
||||
assert.Equal(t, "1", ref.ID)
|
||||
assert.Equal(t, types.SearchTypeWeb, ref.Type)
|
||||
assert.Equal(t, types.SourceHook, ref.Source)
|
||||
assert.Equal(t, 0.8, ref.Weight)
|
||||
|
|
@ -176,7 +176,7 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
name: "single ref with all fields",
|
||||
refs: []*types.Reference{
|
||||
{
|
||||
ID: "ref_001",
|
||||
ID: "1",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceUser,
|
||||
Weight: 1.0,
|
||||
|
|
@ -189,7 +189,7 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
contains: []string{
|
||||
"<references>",
|
||||
"</references>",
|
||||
`<ref id="ref_001" type="web" weight="1.0" source="user">`,
|
||||
`<ref id="1" type="web" weight="1.0" source="user">`,
|
||||
"</ref>",
|
||||
"Test Title",
|
||||
"Test Content",
|
||||
|
|
@ -200,7 +200,7 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
name: "ref without title",
|
||||
refs: []*types.Reference{
|
||||
{
|
||||
ID: "ref_001",
|
||||
ID: "1",
|
||||
Type: types.SearchTypeKB,
|
||||
Source: types.SourceHook,
|
||||
Weight: 0.8,
|
||||
|
|
@ -208,7 +208,7 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
},
|
||||
},
|
||||
contains: []string{
|
||||
`<ref id="ref_001" type="kb" weight="0.8" source="hook">`,
|
||||
`<ref id="1" type="kb" weight="0.8" source="hook">`,
|
||||
"Content without title",
|
||||
},
|
||||
excludes: []string{
|
||||
|
|
@ -219,7 +219,7 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
name: "ref without URL",
|
||||
refs: []*types.Reference{
|
||||
{
|
||||
ID: "ref_001",
|
||||
ID: "1",
|
||||
Type: types.SearchTypeDB,
|
||||
Source: types.SourceAuto,
|
||||
Weight: 0.6,
|
||||
|
|
@ -228,7 +228,7 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
},
|
||||
},
|
||||
contains: []string{
|
||||
`<ref id="ref_001" type="db" weight="0.6" source="auto">`,
|
||||
`<ref id="1" type="db" weight="0.6" source="auto">`,
|
||||
"DB Record",
|
||||
"Database content",
|
||||
},
|
||||
|
|
@ -239,16 +239,16 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
{
|
||||
name: "multiple refs",
|
||||
refs: []*types.Reference{
|
||||
{ID: "ref_001", Type: types.SearchTypeWeb, Source: types.SourceUser, Weight: 1.0, Content: "Content 1"},
|
||||
{ID: "ref_002", Type: types.SearchTypeKB, Source: types.SourceHook, Weight: 0.8, Content: "Content 2"},
|
||||
{ID: "ref_003", Type: types.SearchTypeDB, Source: types.SourceAuto, Weight: 0.6, Content: "Content 3"},
|
||||
{ID: "1", Type: types.SearchTypeWeb, Source: types.SourceUser, Weight: 1.0, Content: "Content 1"},
|
||||
{ID: "2", Type: types.SearchTypeKB, Source: types.SourceHook, Weight: 0.8, Content: "Content 2"},
|
||||
{ID: "3", Type: types.SearchTypeDB, Source: types.SourceAuto, Weight: 0.6, Content: "Content 3"},
|
||||
},
|
||||
contains: []string{
|
||||
"<references>",
|
||||
"</references>",
|
||||
`id="ref_001"`,
|
||||
`id="ref_002"`,
|
||||
`id="ref_003"`,
|
||||
`id="1"`,
|
||||
`id="2"`,
|
||||
`id="3"`,
|
||||
"Content 1",
|
||||
"Content 2",
|
||||
"Content 3",
|
||||
|
|
@ -257,13 +257,13 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
{
|
||||
name: "nil ref in slice",
|
||||
refs: []*types.Reference{
|
||||
{ID: "ref_001", Type: types.SearchTypeWeb, Weight: 1.0, Content: "Content 1"},
|
||||
{ID: "1", Type: types.SearchTypeWeb, Weight: 1.0, Content: "Content 1"},
|
||||
nil,
|
||||
{ID: "ref_002", Type: types.SearchTypeKB, Weight: 0.8, Content: "Content 2"},
|
||||
{ID: "2", Type: types.SearchTypeKB, Weight: 0.8, Content: "Content 2"},
|
||||
},
|
||||
contains: []string{
|
||||
`id="ref_001"`,
|
||||
`id="ref_002"`,
|
||||
`id="1"`,
|
||||
`id="2"`,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
|
@ -286,7 +286,7 @@ func TestFormatReferencesXML(t *testing.T) {
|
|||
func TestFormatReferencesXML_Structure(t *testing.T) {
|
||||
refs := []*types.Reference{
|
||||
{
|
||||
ID: "ref_001",
|
||||
ID: "1",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceUser,
|
||||
Weight: 1.0,
|
||||
|
|
@ -361,6 +361,8 @@ func TestDefaultCitationPrompt(t *testing.T) {
|
|||
assert.Contains(t, DefaultCitationPrompt, `<a class="ref"`)
|
||||
assert.Contains(t, DefaultCitationPrompt, "data-ref-id")
|
||||
assert.Contains(t, DefaultCitationPrompt, "data-ref-type")
|
||||
// Verify example uses simple integer ID
|
||||
assert.Contains(t, DefaultCitationPrompt, `data-ref-id="1"`)
|
||||
}
|
||||
|
||||
func TestBuildReferenceContext(t *testing.T) {
|
||||
|
|
@ -369,7 +371,7 @@ func TestBuildReferenceContext(t *testing.T) {
|
|||
Type: types.SearchTypeWeb,
|
||||
Items: []*types.ResultItem{
|
||||
{
|
||||
CitationID: "ref_001",
|
||||
CitationID: "1",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Weight: 0.6,
|
||||
|
|
@ -387,7 +389,7 @@ func TestBuildReferenceContext(t *testing.T) {
|
|||
assert.NotNil(t, ctx)
|
||||
assert.Equal(t, 1, len(ctx.References))
|
||||
assert.Contains(t, ctx.XML, "<references>")
|
||||
assert.Contains(t, ctx.XML, "ref_001")
|
||||
assert.Contains(t, ctx.XML, `id="1"`)
|
||||
assert.Equal(t, DefaultCitationPrompt, ctx.Prompt)
|
||||
})
|
||||
|
||||
|
|
@ -419,7 +421,7 @@ func TestBuildReferenceContext_Integration(t *testing.T) {
|
|||
Query: "AI developments",
|
||||
Items: []*types.ResultItem{
|
||||
{
|
||||
CitationID: "ref_001",
|
||||
CitationID: "1",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Weight: 0.6,
|
||||
|
|
@ -435,7 +437,7 @@ func TestBuildReferenceContext_Integration(t *testing.T) {
|
|||
Query: "AI developments",
|
||||
Items: []*types.ResultItem{
|
||||
{
|
||||
CitationID: "ref_002",
|
||||
CitationID: "2",
|
||||
Type: types.SearchTypeKB,
|
||||
Source: types.SourceHook,
|
||||
Weight: 0.8,
|
||||
|
|
@ -450,7 +452,7 @@ func TestBuildReferenceContext_Integration(t *testing.T) {
|
|||
Query: "AI developments",
|
||||
Items: []*types.ResultItem{
|
||||
{
|
||||
CitationID: "ref_003",
|
||||
CitationID: "3",
|
||||
Type: types.SearchTypeDB,
|
||||
Source: types.SourceUser,
|
||||
Weight: 1.0,
|
||||
|
|
@ -468,9 +470,9 @@ func TestBuildReferenceContext_Integration(t *testing.T) {
|
|||
assert.Equal(t, 3, len(ctx.References))
|
||||
|
||||
// Verify XML contains all references
|
||||
assert.Contains(t, ctx.XML, "ref_001")
|
||||
assert.Contains(t, ctx.XML, "ref_002")
|
||||
assert.Contains(t, ctx.XML, "ref_003")
|
||||
assert.Contains(t, ctx.XML, `id="1"`)
|
||||
assert.Contains(t, ctx.XML, `id="2"`)
|
||||
assert.Contains(t, ctx.XML, `id="3"`)
|
||||
|
||||
// Verify different source types are represented
|
||||
assert.Contains(t, ctx.XML, `source="auto"`)
|
||||
|
|
|
|||
|
|
@ -352,7 +352,7 @@ func TestSearcher_BuildReferences(t *testing.T) {
|
|||
Type: types.SearchTypeWeb,
|
||||
Items: []*types.ResultItem{
|
||||
{
|
||||
CitationID: "ref_001",
|
||||
CitationID: "1",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Weight: 0.6,
|
||||
|
|
@ -366,7 +366,7 @@ func TestSearcher_BuildReferences(t *testing.T) {
|
|||
Type: types.SearchTypeKB,
|
||||
Items: []*types.ResultItem{
|
||||
{
|
||||
CitationID: "ref_002",
|
||||
CitationID: "2",
|
||||
Type: types.SearchTypeKB,
|
||||
Source: types.SourceHook,
|
||||
Weight: 0.8,
|
||||
|
|
@ -379,8 +379,8 @@ func TestSearcher_BuildReferences(t *testing.T) {
|
|||
|
||||
refs := s.BuildReferences(results)
|
||||
assert.Equal(t, 2, len(refs))
|
||||
assert.Equal(t, "ref_001", refs[0].ID)
|
||||
assert.Equal(t, "ref_002", refs[1].ID)
|
||||
assert.Equal(t, "1", refs[0].ID)
|
||||
assert.Equal(t, "2", refs[1].ID)
|
||||
}
|
||||
|
||||
func TestSearcher_CitationGeneration(t *testing.T) {
|
||||
|
|
@ -396,7 +396,8 @@ func TestSearcher_CitationGeneration(t *testing.T) {
|
|||
id2 := s.citation.Next()
|
||||
id3 := s.citation.Next()
|
||||
|
||||
assert.Equal(t, "ref_001", id1)
|
||||
assert.Equal(t, "ref_002", id2)
|
||||
assert.Equal(t, "ref_003", id3)
|
||||
// Citation IDs are now simple integers
|
||||
assert.Equal(t, "1", id1)
|
||||
assert.Equal(t, "2", id2)
|
||||
assert.Equal(t, "3", id3)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -157,3 +157,31 @@ func (m *Mongo) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
|
|||
// TODO: implement
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Search Management
|
||||
// =============================================================================
|
||||
|
||||
// SaveSearch saves a search record for a request
|
||||
func (m *Mongo) SaveSearch(search *types.Search) error {
|
||||
// TODO: implement
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetSearches retrieves all search records for a request
|
||||
func (m *Mongo) GetSearches(requestID string) ([]*types.Search, error) {
|
||||
// TODO: implement
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// GetReference retrieves a single reference by request ID and index
|
||||
func (m *Mongo) GetReference(requestID string, index int) (*types.Reference, error) {
|
||||
// TODO: implement
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// DeleteSearches deletes all search records for a chat
|
||||
func (m *Mongo) DeleteSearches(chatID string) error {
|
||||
// TODO: implement
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -157,3 +157,31 @@ func (r *Redis) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
|
|||
// TODO: implement
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Search Management
|
||||
// =============================================================================
|
||||
|
||||
// SaveSearch saves a search record for a request
|
||||
func (r *Redis) SaveSearch(search *types.Search) error {
|
||||
// TODO: implement
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetSearches retrieves all search records for a request
|
||||
func (r *Redis) GetSearches(requestID string) ([]*types.Search, error) {
|
||||
// TODO: implement
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// GetReference retrieves a single reference by request ID and index
|
||||
func (r *Redis) GetReference(requestID string, index int) (*types.Reference, error) {
|
||||
// TODO: implement
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// DeleteSearches deletes all search records for a chat
|
||||
func (r *Redis) DeleteSearches(chatID string) error {
|
||||
// TODO: implement
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -98,6 +98,34 @@ type ChatStore interface {
|
|||
// chatID: Chat ID
|
||||
// Returns: Potential error
|
||||
DeleteResume(chatID string) error
|
||||
|
||||
// ==========================================================================
|
||||
// Search Management
|
||||
// ==========================================================================
|
||||
|
||||
// SaveSearch saves a search record for a request
|
||||
// Used for citation support, debugging, and replay
|
||||
// search: Search record to save
|
||||
// Returns: Potential error
|
||||
SaveSearch(search *Search) error
|
||||
|
||||
// GetSearches retrieves all search records for a request
|
||||
// requestID: Request ID
|
||||
// Returns: Search records and potential error
|
||||
GetSearches(requestID string) ([]*Search, error)
|
||||
|
||||
// GetReference retrieves a single reference by request ID and index
|
||||
// Used for citation click handling
|
||||
// requestID: Request ID
|
||||
// index: Reference index (1-based)
|
||||
// Returns: Reference and potential error
|
||||
GetReference(requestID string, index int) (*Reference, error)
|
||||
|
||||
// DeleteSearches deletes all search records for a chat
|
||||
// Called when deleting a chat
|
||||
// chatID: Chat ID
|
||||
// Returns: Potential error
|
||||
DeleteSearches(chatID string) error
|
||||
}
|
||||
|
||||
// AssistantStore defines the assistant storage interface
|
||||
|
|
|
|||
|
|
@ -447,3 +447,71 @@ type AssistantModel struct {
|
|||
YaoTeamID string `json:"-"` // Team ID for team-based access control (not exposed in JSON)
|
||||
YaoTenantID string `json:"-"` // Tenant ID for multi-tenancy support (not exposed in JSON)
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Search Types (for search result storage)
|
||||
// =============================================================================
|
||||
|
||||
// Search represents stored search results for a request
|
||||
// Stores all intermediate processing results for debugging, replay, and citation
|
||||
type Search struct {
|
||||
ID int64 `json:"id"`
|
||||
RequestID string `json:"request_id"`
|
||||
ChatID string `json:"chat_id"`
|
||||
Query string `json:"query"` // Original query
|
||||
Config map[string]any `json:"config,omitempty"` // Search config used (for tuning)
|
||||
Keywords []string `json:"keywords,omitempty"` // Extracted keywords (Web/NLP)
|
||||
Entities []Entity `json:"entities,omitempty"` // Extracted entities (Graph)
|
||||
Relations []Relation `json:"relations,omitempty"` // Extracted relations (Graph)
|
||||
DSL map[string]any `json:"dsl,omitempty"` // Generated QueryDSL (DB)
|
||||
Source string `json:"source"` // web/kb/db/auto
|
||||
References []Reference `json:"references"` // References with global index
|
||||
Graph []GraphNode `json:"graph,omitempty"` // Graph nodes from KB
|
||||
XML string `json:"xml,omitempty"` // Formatted XML for LLM
|
||||
Prompt string `json:"prompt,omitempty"` // Citation prompt
|
||||
Duration int64 `json:"duration_ms"` // Search duration in ms
|
||||
Error string `json:"error,omitempty"` // Error if failed
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// Reference represents a single reference with global index (for storage)
|
||||
type Reference struct {
|
||||
Index int `json:"index"` // Global index (1-based, unique within request)
|
||||
Type string `json:"type"` // web/kb/db
|
||||
Title string `json:"title"` // Reference title
|
||||
URL string `json:"url,omitempty"` // URL (for web)
|
||||
Snippet string `json:"snippet,omitempty"` // Short snippet
|
||||
Content string `json:"content,omitempty"` // Full content
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// SearchFilter for listing searches
|
||||
type SearchFilter struct {
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
ChatID string `json:"chat_id,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
// Entity represents an extracted entity (for Graph RAG)
|
||||
type Entity struct {
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
// Relation represents an extracted relation (for Graph RAG)
|
||||
type Relation struct {
|
||||
Subject string `json:"subject"`
|
||||
Predicate string `json:"predicate"`
|
||||
Object string `json:"object"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
// GraphNode represents a node from knowledge graph
|
||||
type GraphNode struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Label string `json:"label,omitempty"`
|
||||
Properties map[string]any `json:"properties,omitempty"`
|
||||
Score float64 `json:"score,omitempty"`
|
||||
}
|
||||
|
|
|
|||
301
agent/store/xun/search.go
Normal file
301
agent/store/xun/search.go
Normal file
|
|
@ -0,0 +1,301 @@
|
|||
package xun
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
"github.com/yaoapp/gou/model"
|
||||
"github.com/yaoapp/xun/dbal/query"
|
||||
"github.com/yaoapp/yao/agent/store/types"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// Search Management
|
||||
// =============================================================================
|
||||
|
||||
// SaveSearch saves a search record for a request
|
||||
func (store *Xun) SaveSearch(search *types.Search) error {
|
||||
if search == nil {
|
||||
return fmt.Errorf("search is nil")
|
||||
}
|
||||
if search.RequestID == "" {
|
||||
return fmt.Errorf("request_id is required")
|
||||
}
|
||||
if search.ChatID == "" {
|
||||
return fmt.Errorf("chat_id is required")
|
||||
}
|
||||
if search.Source == "" {
|
||||
return fmt.Errorf("source is required")
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
|
||||
// Build row data
|
||||
row := map[string]interface{}{
|
||||
"request_id": search.RequestID,
|
||||
"chat_id": search.ChatID,
|
||||
"query": search.Query,
|
||||
"source": search.Source,
|
||||
"duration": search.Duration,
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
}
|
||||
|
||||
// Handle JSON fields
|
||||
if search.Config != nil {
|
||||
configJSON, err := jsoniter.MarshalToString(search.Config)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal config: %w", err)
|
||||
}
|
||||
row["config"] = configJSON
|
||||
}
|
||||
|
||||
if len(search.Keywords) > 0 {
|
||||
keywordsJSON, err := jsoniter.MarshalToString(search.Keywords)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal keywords: %w", err)
|
||||
}
|
||||
row["keywords"] = keywordsJSON
|
||||
}
|
||||
|
||||
if len(search.Entities) > 0 {
|
||||
entitiesJSON, err := jsoniter.MarshalToString(search.Entities)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal entities: %w", err)
|
||||
}
|
||||
row["entities"] = entitiesJSON
|
||||
}
|
||||
|
||||
if len(search.Relations) > 0 {
|
||||
relationsJSON, err := jsoniter.MarshalToString(search.Relations)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal relations: %w", err)
|
||||
}
|
||||
row["relations"] = relationsJSON
|
||||
}
|
||||
|
||||
if search.DSL != nil {
|
||||
dslJSON, err := jsoniter.MarshalToString(search.DSL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal dsl: %w", err)
|
||||
}
|
||||
row["dsl"] = dslJSON
|
||||
}
|
||||
|
||||
if len(search.References) > 0 {
|
||||
refsJSON, err := jsoniter.MarshalToString(search.References)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal references: %w", err)
|
||||
}
|
||||
row["references"] = refsJSON
|
||||
}
|
||||
|
||||
if len(search.Graph) > 0 {
|
||||
graphJSON, err := jsoniter.MarshalToString(search.Graph)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal graph: %w", err)
|
||||
}
|
||||
row["graph"] = graphJSON
|
||||
}
|
||||
|
||||
if search.XML != "" {
|
||||
row["xml"] = search.XML
|
||||
}
|
||||
|
||||
if search.Prompt != "" {
|
||||
row["prompt"] = search.Prompt
|
||||
}
|
||||
|
||||
if search.Error != "" {
|
||||
row["error"] = search.Error
|
||||
}
|
||||
|
||||
return store.newQuerySearch().Insert(row)
|
||||
}
|
||||
|
||||
// GetSearches retrieves all search records for a request
|
||||
func (store *Xun) GetSearches(requestID string) ([]*types.Search, error) {
|
||||
if requestID == "" {
|
||||
return nil, fmt.Errorf("request_id is required")
|
||||
}
|
||||
|
||||
rows, err := store.newQuerySearch().
|
||||
Where("request_id", requestID).
|
||||
WhereNull("deleted_at").
|
||||
OrderBy("created_at", "asc").
|
||||
Get()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
searches := make([]*types.Search, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
data := row.ToMap()
|
||||
if data == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
search, err := store.rowToSearch(data)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
searches = append(searches, search)
|
||||
}
|
||||
|
||||
return searches, nil
|
||||
}
|
||||
|
||||
// GetReference retrieves a single reference by request ID and index
|
||||
func (store *Xun) GetReference(requestID string, index int) (*types.Reference, error) {
|
||||
if requestID == "" {
|
||||
return nil, fmt.Errorf("request_id is required")
|
||||
}
|
||||
if index < 1 {
|
||||
return nil, fmt.Errorf("index must be >= 1")
|
||||
}
|
||||
|
||||
// Get all searches for this request
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Find the reference with matching index
|
||||
for _, search := range searches {
|
||||
for _, ref := range search.References {
|
||||
if ref.Index == index {
|
||||
return &ref, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("reference not found: request_id=%s, index=%d", requestID, index)
|
||||
}
|
||||
|
||||
// DeleteSearches deletes all search records for a chat (soft delete)
|
||||
func (store *Xun) DeleteSearches(chatID string) error {
|
||||
if chatID == "" {
|
||||
return fmt.Errorf("chat_id is required")
|
||||
}
|
||||
|
||||
_, err := store.newQuerySearch().
|
||||
Where("chat_id", chatID).
|
||||
WhereNull("deleted_at").
|
||||
Update(map[string]interface{}{
|
||||
"deleted_at": time.Now(),
|
||||
"updated_at": time.Now(),
|
||||
})
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Query Builder
|
||||
// =============================================================================
|
||||
|
||||
// newQuerySearch creates a new query builder for the search table
|
||||
func (store *Xun) newQuerySearch() query.Query {
|
||||
qb := store.query.New()
|
||||
qb.Table(store.getSearchTable())
|
||||
return qb
|
||||
}
|
||||
|
||||
// getSearchTable returns the search table name
|
||||
func (store *Xun) getSearchTable() string {
|
||||
m := model.Select("__yao.agent.search")
|
||||
if m != nil && m.MetaData.Table.Name != "" {
|
||||
return m.MetaData.Table.Name
|
||||
}
|
||||
return "agent_search"
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Helper Functions
|
||||
// =============================================================================
|
||||
|
||||
// rowToSearch converts a database row to a Search struct
|
||||
func (store *Xun) rowToSearch(data map[string]interface{}) (*types.Search, error) {
|
||||
search := &types.Search{
|
||||
ID: getInt64(data, "id"),
|
||||
RequestID: getString(data, "request_id"),
|
||||
ChatID: getString(data, "chat_id"),
|
||||
Query: getString(data, "query"),
|
||||
Source: getString(data, "source"),
|
||||
XML: getString(data, "xml"),
|
||||
Prompt: getString(data, "prompt"),
|
||||
Duration: getInt64(data, "duration"),
|
||||
Error: getString(data, "error"),
|
||||
}
|
||||
|
||||
// Handle timestamps
|
||||
if createdAt := getTime(data, "created_at"); createdAt != nil {
|
||||
search.CreatedAt = *createdAt
|
||||
}
|
||||
|
||||
// Parse JSON fields
|
||||
if config := data["config"]; config != nil {
|
||||
if configStr, ok := config.(string); ok && configStr != "" {
|
||||
var configMap map[string]any
|
||||
if err := jsoniter.UnmarshalFromString(configStr, &configMap); err == nil {
|
||||
search.Config = configMap
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if keywords := data["keywords"]; keywords != nil {
|
||||
if keywordsStr, ok := keywords.(string); ok && keywordsStr != "" {
|
||||
var keywordsList []string
|
||||
if err := jsoniter.UnmarshalFromString(keywordsStr, &keywordsList); err == nil {
|
||||
search.Keywords = keywordsList
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if entities := data["entities"]; entities != nil {
|
||||
if entitiesStr, ok := entities.(string); ok && entitiesStr != "" {
|
||||
var entitiesList []types.Entity
|
||||
if err := jsoniter.UnmarshalFromString(entitiesStr, &entitiesList); err == nil {
|
||||
search.Entities = entitiesList
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if relations := data["relations"]; relations != nil {
|
||||
if relationsStr, ok := relations.(string); ok && relationsStr != "" {
|
||||
var relationsList []types.Relation
|
||||
if err := jsoniter.UnmarshalFromString(relationsStr, &relationsList); err == nil {
|
||||
search.Relations = relationsList
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if dsl := data["dsl"]; dsl != nil {
|
||||
if dslStr, ok := dsl.(string); ok && dslStr != "" {
|
||||
var dslMap map[string]any
|
||||
if err := jsoniter.UnmarshalFromString(dslStr, &dslMap); err == nil {
|
||||
search.DSL = dslMap
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if refs := data["references"]; refs != nil {
|
||||
if refsStr, ok := refs.(string); ok && refsStr != "" {
|
||||
var refsList []types.Reference
|
||||
if err := jsoniter.UnmarshalFromString(refsStr, &refsList); err == nil {
|
||||
search.References = refsList
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if graph := data["graph"]; graph != nil {
|
||||
if graphStr, ok := graph.(string); ok && graphStr != "" {
|
||||
var graphList []types.GraphNode
|
||||
if err := jsoniter.UnmarshalFromString(graphStr, &graphList); err == nil {
|
||||
search.Graph = graphList
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return search, nil
|
||||
}
|
||||
715
agent/store/xun/search_test.go
Normal file
715
agent/store/xun/search_test.go
Normal file
|
|
@ -0,0 +1,715 @@
|
|||
package xun_test
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/yao/agent/store/types"
|
||||
"github.com/yaoapp/yao/agent/store/xun"
|
||||
"github.com/yaoapp/yao/config"
|
||||
"github.com/yaoapp/yao/test"
|
||||
)
|
||||
|
||||
// TestSaveSearch tests saving search records
|
||||
func TestSaveSearch(t *testing.T) {
|
||||
test.Prepare(t, config.Conf)
|
||||
defer test.Clean()
|
||||
|
||||
store, err := xun.NewXun(types.Setting{
|
||||
Connector: "default",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create store: %v", err)
|
||||
}
|
||||
|
||||
// Create a chat first
|
||||
chat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
Title: "Search Test Chat",
|
||||
}
|
||||
err = store.CreateChat(chat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
defer store.DeleteChat(chat.ChatID)
|
||||
|
||||
t.Run("SaveBasicSearch", func(t *testing.T) {
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "What is the weather today?",
|
||||
Source: "web",
|
||||
Duration: 150,
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
// Verify
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 1 {
|
||||
t.Fatalf("Expected 1 search, got %d", len(searches))
|
||||
}
|
||||
|
||||
if searches[0].Query != "What is the weather today?" {
|
||||
t.Errorf("Expected query 'What is the weather today?', got '%s'", searches[0].Query)
|
||||
}
|
||||
if searches[0].Source != "web" {
|
||||
t.Errorf("Expected source 'web', got '%s'", searches[0].Source)
|
||||
}
|
||||
if searches[0].Duration != 150 {
|
||||
t.Errorf("Expected duration 150, got %d", searches[0].Duration)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithKeywords", func(t *testing.T) {
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "Latest news about AI",
|
||||
Keywords: []string{"AI", "news", "latest"},
|
||||
Source: "web",
|
||||
Duration: 200,
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 1 {
|
||||
t.Fatalf("Expected 1 search, got %d", len(searches))
|
||||
}
|
||||
|
||||
if len(searches[0].Keywords) != 3 {
|
||||
t.Errorf("Expected 3 keywords, got %d", len(searches[0].Keywords))
|
||||
}
|
||||
if searches[0].Keywords[0] != "AI" {
|
||||
t.Errorf("Expected first keyword 'AI', got '%s'", searches[0].Keywords[0])
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithReferences", func(t *testing.T) {
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "How to learn Go programming?",
|
||||
Source: "web",
|
||||
References: []types.Reference{
|
||||
{
|
||||
Index: 1,
|
||||
Type: "web",
|
||||
Title: "Go Programming Tutorial",
|
||||
URL: "https://go.dev/tour/",
|
||||
Snippet: "An interactive introduction to Go",
|
||||
},
|
||||
{
|
||||
Index: 2,
|
||||
Type: "web",
|
||||
Title: "Effective Go",
|
||||
URL: "https://go.dev/doc/effective_go",
|
||||
Snippet: "Tips for writing clear, idiomatic Go code",
|
||||
},
|
||||
},
|
||||
XML: "<references>...</references>",
|
||||
Prompt: "Please cite sources using [1], [2]...",
|
||||
Duration: 300,
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 1 {
|
||||
t.Fatalf("Expected 1 search, got %d", len(searches))
|
||||
}
|
||||
|
||||
if len(searches[0].References) != 2 {
|
||||
t.Errorf("Expected 2 references, got %d", len(searches[0].References))
|
||||
}
|
||||
if searches[0].References[0].Title != "Go Programming Tutorial" {
|
||||
t.Errorf("Expected first reference title 'Go Programming Tutorial', got '%s'", searches[0].References[0].Title)
|
||||
}
|
||||
if searches[0].XML != "<references>...</references>" {
|
||||
t.Errorf("Expected XML '<references>...</references>', got '%s'", searches[0].XML)
|
||||
}
|
||||
if searches[0].Prompt != "Please cite sources using [1], [2]..." {
|
||||
t.Errorf("Expected prompt to be set")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithConfig", func(t *testing.T) {
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "Config test",
|
||||
Source: "auto",
|
||||
Config: map[string]any{
|
||||
"uses": map[string]any{
|
||||
"search": "builtin",
|
||||
"web": "builtin",
|
||||
"keyword": "builtin",
|
||||
},
|
||||
"web": map[string]any{
|
||||
"provider": "tavily",
|
||||
"max_results": 5,
|
||||
},
|
||||
},
|
||||
Duration: 100,
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 1 {
|
||||
t.Fatalf("Expected 1 search, got %d", len(searches))
|
||||
}
|
||||
|
||||
if searches[0].Config == nil {
|
||||
t.Fatal("Expected config to be set")
|
||||
}
|
||||
uses, ok := searches[0].Config["uses"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatal("Expected uses in config")
|
||||
}
|
||||
if uses["search"] != "builtin" {
|
||||
t.Errorf("Expected uses.search='builtin', got '%v'", uses["search"])
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithEntitiesAndRelations", func(t *testing.T) {
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "Who is the CEO of Apple?",
|
||||
Source: "kb",
|
||||
Entities: []types.Entity{
|
||||
{Name: "Apple", Type: "Organization"},
|
||||
{Name: "Tim Cook", Type: "Person"},
|
||||
},
|
||||
Relations: []types.Relation{
|
||||
{Subject: "Tim Cook", Predicate: "CEO_of", Object: "Apple"},
|
||||
},
|
||||
Graph: []types.GraphNode{
|
||||
{ID: "node1", Type: "Organization", Label: "Apple", Score: 0.95},
|
||||
{ID: "node2", Type: "Person", Label: "Tim Cook", Score: 0.92},
|
||||
},
|
||||
Duration: 250,
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 1 {
|
||||
t.Fatalf("Expected 1 search, got %d", len(searches))
|
||||
}
|
||||
|
||||
if len(searches[0].Entities) != 2 {
|
||||
t.Errorf("Expected 2 entities, got %d", len(searches[0].Entities))
|
||||
}
|
||||
if searches[0].Entities[0].Name != "Apple" {
|
||||
t.Errorf("Expected first entity 'Apple', got '%s'", searches[0].Entities[0].Name)
|
||||
}
|
||||
|
||||
if len(searches[0].Relations) != 1 {
|
||||
t.Errorf("Expected 1 relation, got %d", len(searches[0].Relations))
|
||||
}
|
||||
if searches[0].Relations[0].Predicate != "CEO_of" {
|
||||
t.Errorf("Expected predicate 'CEO_of', got '%s'", searches[0].Relations[0].Predicate)
|
||||
}
|
||||
|
||||
if len(searches[0].Graph) != 2 {
|
||||
t.Errorf("Expected 2 graph nodes, got %d", len(searches[0].Graph))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithDSL", func(t *testing.T) {
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "Find orders over $1000",
|
||||
Source: "db",
|
||||
DSL: map[string]any{
|
||||
"wheres": []map[string]any{
|
||||
{"column": "amount", "op": ">", "value": 1000},
|
||||
},
|
||||
"orders": []map[string]any{
|
||||
{"column": "created_at", "option": "desc"},
|
||||
},
|
||||
},
|
||||
Duration: 50,
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 1 {
|
||||
t.Fatalf("Expected 1 search, got %d", len(searches))
|
||||
}
|
||||
|
||||
if searches[0].DSL == nil {
|
||||
t.Fatal("Expected DSL to be set")
|
||||
}
|
||||
wheres, ok := searches[0].DSL["wheres"].([]any)
|
||||
if !ok || len(wheres) == 0 {
|
||||
t.Error("Expected wheres in DSL")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithError", func(t *testing.T) {
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "Failed search",
|
||||
Source: "web",
|
||||
Error: "API rate limit exceeded",
|
||||
Duration: 10,
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 1 {
|
||||
t.Fatalf("Expected 1 search, got %d", len(searches))
|
||||
}
|
||||
|
||||
if searches[0].Error != "API rate limit exceeded" {
|
||||
t.Errorf("Expected error 'API rate limit exceeded', got '%s'", searches[0].Error)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithoutRequestID", func(t *testing.T) {
|
||||
search := &types.Search{
|
||||
ChatID: chat.ChatID,
|
||||
Query: "Test",
|
||||
Source: "web",
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err == nil {
|
||||
t.Error("Expected error when saving without request_id")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithoutChatID", func(t *testing.T) {
|
||||
search := &types.Search{
|
||||
RequestID: "req_test",
|
||||
Query: "Test",
|
||||
Source: "web",
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err == nil {
|
||||
t.Error("Expected error when saving without chat_id")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveSearchWithoutSource", func(t *testing.T) {
|
||||
search := &types.Search{
|
||||
RequestID: "req_test",
|
||||
ChatID: chat.ChatID,
|
||||
Query: "Test",
|
||||
}
|
||||
|
||||
err := store.SaveSearch(search)
|
||||
if err == nil {
|
||||
t.Error("Expected error when saving without source")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveNilSearch", func(t *testing.T) {
|
||||
err := store.SaveSearch(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error when saving nil search")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestGetSearches tests retrieving search records
|
||||
func TestGetSearches(t *testing.T) {
|
||||
test.Prepare(t, config.Conf)
|
||||
defer test.Clean()
|
||||
|
||||
store, err := xun.NewXun(types.Setting{
|
||||
Connector: "default",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create store: %v", err)
|
||||
}
|
||||
|
||||
// Create a chat
|
||||
chat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
}
|
||||
err = store.CreateChat(chat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
defer store.DeleteChat(chat.ChatID)
|
||||
|
||||
t.Run("GetMultipleSearches", func(t *testing.T) {
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
|
||||
// Save multiple searches for the same request
|
||||
for i := 1; i <= 3; i++ {
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: fmt.Sprintf("Query %d", i),
|
||||
Source: "web",
|
||||
Duration: int64(i * 100),
|
||||
}
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search %d: %v", i, err)
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond) // Ensure different created_at
|
||||
}
|
||||
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 3 {
|
||||
t.Errorf("Expected 3 searches, got %d", len(searches))
|
||||
}
|
||||
|
||||
// Verify order (by created_at asc)
|
||||
for i := 0; i < len(searches)-1; i++ {
|
||||
if searches[i].CreatedAt.After(searches[i+1].CreatedAt) {
|
||||
t.Error("Searches not ordered by created_at asc")
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("GetSearchesForNonExistentRequest", func(t *testing.T) {
|
||||
searches, err := store.GetSearches("nonexistent_request")
|
||||
if err != nil {
|
||||
t.Fatalf("Unexpected error: %v", err)
|
||||
}
|
||||
if len(searches) != 0 {
|
||||
t.Errorf("Expected 0 searches, got %d", len(searches))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("GetSearchesWithEmptyRequestID", func(t *testing.T) {
|
||||
_, err := store.GetSearches("")
|
||||
if err == nil {
|
||||
t.Error("Expected error when getting searches without request_id")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestGetReference tests retrieving a single reference
|
||||
func TestGetReference(t *testing.T) {
|
||||
test.Prepare(t, config.Conf)
|
||||
defer test.Clean()
|
||||
|
||||
store, err := xun.NewXun(types.Setting{
|
||||
Connector: "default",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create store: %v", err)
|
||||
}
|
||||
|
||||
// Create a chat
|
||||
chat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
}
|
||||
err = store.CreateChat(chat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
defer store.DeleteChat(chat.ChatID)
|
||||
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
|
||||
// Save search with references
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "Test query",
|
||||
Source: "web",
|
||||
References: []types.Reference{
|
||||
{Index: 1, Type: "web", Title: "Reference 1", URL: "https://example.com/1"},
|
||||
{Index: 2, Type: "web", Title: "Reference 2", URL: "https://example.com/2"},
|
||||
{Index: 3, Type: "kb", Title: "Reference 3", Content: "KB content"},
|
||||
},
|
||||
Duration: 100,
|
||||
}
|
||||
err = store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
t.Run("GetExistingReference", func(t *testing.T) {
|
||||
ref, err := store.GetReference(requestID, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get reference: %v", err)
|
||||
}
|
||||
|
||||
if ref.Title != "Reference 1" {
|
||||
t.Errorf("Expected title 'Reference 1', got '%s'", ref.Title)
|
||||
}
|
||||
if ref.URL != "https://example.com/1" {
|
||||
t.Errorf("Expected URL 'https://example.com/1', got '%s'", ref.URL)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("GetReferenceByIndex", func(t *testing.T) {
|
||||
ref, err := store.GetReference(requestID, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get reference: %v", err)
|
||||
}
|
||||
|
||||
if ref.Type != "kb" {
|
||||
t.Errorf("Expected type 'kb', got '%s'", ref.Type)
|
||||
}
|
||||
if ref.Content != "KB content" {
|
||||
t.Errorf("Expected content 'KB content', got '%s'", ref.Content)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("GetNonExistentReference", func(t *testing.T) {
|
||||
_, err := store.GetReference(requestID, 999)
|
||||
if err == nil {
|
||||
t.Error("Expected error when getting non-existent reference")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("GetReferenceWithInvalidIndex", func(t *testing.T) {
|
||||
_, err := store.GetReference(requestID, 0)
|
||||
if err == nil {
|
||||
t.Error("Expected error when getting reference with index 0")
|
||||
}
|
||||
|
||||
_, err = store.GetReference(requestID, -1)
|
||||
if err == nil {
|
||||
t.Error("Expected error when getting reference with negative index")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("GetReferenceWithEmptyRequestID", func(t *testing.T) {
|
||||
_, err := store.GetReference("", 1)
|
||||
if err == nil {
|
||||
t.Error("Expected error when getting reference without request_id")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestDeleteSearches tests deleting search records
|
||||
func TestDeleteSearches(t *testing.T) {
|
||||
test.Prepare(t, config.Conf)
|
||||
defer test.Clean()
|
||||
|
||||
store, err := xun.NewXun(types.Setting{
|
||||
Connector: "default",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create store: %v", err)
|
||||
}
|
||||
|
||||
t.Run("DeleteSearchesForChat", func(t *testing.T) {
|
||||
// Create a chat
|
||||
chat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
}
|
||||
err := store.CreateChat(chat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
defer store.DeleteChat(chat.ChatID)
|
||||
|
||||
// Save multiple searches
|
||||
for i := 1; i <= 3; i++ {
|
||||
requestID := fmt.Sprintf("req_%d_%d", time.Now().UnixNano(), i)
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: fmt.Sprintf("Query %d", i),
|
||||
Source: "web",
|
||||
Duration: 100,
|
||||
}
|
||||
err := store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Delete all searches for the chat
|
||||
err = store.DeleteSearches(chat.ChatID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to delete searches: %v", err)
|
||||
}
|
||||
|
||||
// Note: GetSearches filters by request_id, not chat_id
|
||||
// We can't easily verify deletion without a GetSearchesByChatID method
|
||||
// But the soft delete should have been applied
|
||||
})
|
||||
|
||||
t.Run("DeleteSearchesWithEmptyChatID", func(t *testing.T) {
|
||||
err := store.DeleteSearches("")
|
||||
if err == nil {
|
||||
t.Error("Expected error when deleting searches without chat_id")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestSearchCompleteWorkflow tests a complete search workflow
|
||||
func TestSearchCompleteWorkflow(t *testing.T) {
|
||||
test.Prepare(t, config.Conf)
|
||||
defer test.Clean()
|
||||
|
||||
store, err := xun.NewXun(types.Setting{
|
||||
Connector: "default",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create store: %v", err)
|
||||
}
|
||||
|
||||
t.Run("CompleteWorkflow", func(t *testing.T) {
|
||||
// 1. Create chat
|
||||
chat := &types.Chat{
|
||||
AssistantID: "workflow_assistant",
|
||||
Title: "Search Workflow Test",
|
||||
}
|
||||
err := store.CreateChat(chat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
defer store.DeleteChat(chat.ChatID)
|
||||
|
||||
requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
|
||||
|
||||
// 2. Save search with full data
|
||||
search := &types.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chat.ChatID,
|
||||
Query: "What are the best practices for Go programming?",
|
||||
Config: map[string]any{
|
||||
"uses": map[string]any{"search": "builtin", "web": "builtin"},
|
||||
"web": map[string]any{"provider": "tavily", "max_results": 5},
|
||||
},
|
||||
Keywords: []string{"Go", "programming", "best practices"},
|
||||
Source: "auto",
|
||||
References: []types.Reference{
|
||||
{Index: 1, Type: "web", Title: "Effective Go", URL: "https://go.dev/doc/effective_go"},
|
||||
{Index: 2, Type: "web", Title: "Go Proverbs", URL: "https://go-proverbs.github.io/"},
|
||||
{Index: 3, Type: "kb", Title: "Internal Go Guide", Content: "Our team's Go coding standards..."},
|
||||
},
|
||||
XML: "<references><ref index=\"1\">...</ref></references>",
|
||||
Prompt: "When citing, use [1], [2], [3] format.",
|
||||
Duration: 350,
|
||||
}
|
||||
|
||||
err = store.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save search: %v", err)
|
||||
}
|
||||
|
||||
// 3. Retrieve searches
|
||||
searches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches: %v", err)
|
||||
}
|
||||
|
||||
if len(searches) != 1 {
|
||||
t.Fatalf("Expected 1 search, got %d", len(searches))
|
||||
}
|
||||
|
||||
// 4. Verify all fields
|
||||
s := searches[0]
|
||||
if s.Query != "What are the best practices for Go programming?" {
|
||||
t.Errorf("Query mismatch")
|
||||
}
|
||||
if len(s.Keywords) != 3 {
|
||||
t.Errorf("Expected 3 keywords, got %d", len(s.Keywords))
|
||||
}
|
||||
if len(s.References) != 3 {
|
||||
t.Errorf("Expected 3 references, got %d", len(s.References))
|
||||
}
|
||||
if s.Config == nil {
|
||||
t.Error("Config should not be nil")
|
||||
}
|
||||
|
||||
// 5. Get specific reference
|
||||
ref, err := store.GetReference(requestID, 2)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get reference: %v", err)
|
||||
}
|
||||
if ref.Title != "Go Proverbs" {
|
||||
t.Errorf("Expected 'Go Proverbs', got '%s'", ref.Title)
|
||||
}
|
||||
|
||||
// 6. Delete searches
|
||||
err = store.DeleteSearches(chat.ChatID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to delete searches: %v", err)
|
||||
}
|
||||
|
||||
// 7. Verify deletion (soft delete, so GetSearches should return empty)
|
||||
deletedSearches, err := store.GetSearches(requestID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get searches after delete: %v", err)
|
||||
}
|
||||
if len(deletedSearches) != 0 {
|
||||
t.Errorf("Expected 0 searches after delete, got %d", len(deletedSearches))
|
||||
}
|
||||
|
||||
t.Log("Complete search workflow passed!")
|
||||
})
|
||||
}
|
||||
307
data/bindata.go
307
data/bindata.go
File diff suppressed because it is too large
Load diff
|
|
@ -24,6 +24,7 @@ var systemModels = map[string]string{
|
|||
"__yao.agent.chat": "yao/models/agent/chat.mod.yao",
|
||||
"__yao.agent.message": "yao/models/agent/message.mod.yao",
|
||||
"__yao.agent.resume": "yao/models/agent/resume.mod.yao",
|
||||
"__yao.agent.search": "yao/models/agent/search.mod.yao",
|
||||
"__yao.attachment": "yao/models/attachment.mod.yao",
|
||||
"__yao.audit": "yao/models/audit.mod.yao",
|
||||
"__yao.config": "yao/models/config.mod.yao",
|
||||
|
|
|
|||
|
|
@ -59,6 +59,18 @@ func Attach(group *gin.RouterGroup, oauth types.OAuth) {
|
|||
// Query params: request_id, role, block_id, thread_id, type, limit, offset
|
||||
group.GET("/sessions/:chat_id/messages", GetMessages)
|
||||
|
||||
// ==========================================================================
|
||||
// Search References (Citation Support)
|
||||
// ==========================================================================
|
||||
|
||||
// Get all references for a request
|
||||
// Returns all search references for citation support
|
||||
group.GET("/references/:request_id", GetReferences)
|
||||
|
||||
// Get a single reference by request ID and index
|
||||
// Returns a specific reference for citation click handling
|
||||
group.GET("/references/:request_id/:index", GetReference)
|
||||
|
||||
}
|
||||
|
||||
func placeholder(c *gin.Context) {
|
||||
|
|
|
|||
186
openapi/chat/reference.go
Normal file
186
openapi/chat/reference.go
Normal file
|
|
@ -0,0 +1,186 @@
|
|||
package chat
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/yaoapp/yao/agent/assistant"
|
||||
storetypes "github.com/yaoapp/yao/agent/store/types"
|
||||
"github.com/yaoapp/yao/openapi/oauth/authorized"
|
||||
"github.com/yaoapp/yao/openapi/response"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// Search Reference Handlers
|
||||
// =============================================================================
|
||||
|
||||
// GetReferences retrieves all search references for a request
|
||||
// GET /v1/chat/references/:request_id
|
||||
func GetReferences(c *gin.Context) {
|
||||
// Get chat store
|
||||
chatStore := assistant.GetChatStore()
|
||||
if chatStore == nil {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrServerError.Code,
|
||||
ErrorDescription: "Chat storage not initialized",
|
||||
}
|
||||
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
// Get request ID from URL parameter
|
||||
requestID := c.Param("request_id")
|
||||
if requestID == "" {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrInvalidRequest.Code,
|
||||
ErrorDescription: "Request ID is required",
|
||||
}
|
||||
response.RespondWithError(c, response.StatusBadRequest, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
// Get all search records for this request
|
||||
searches, err := chatStore.GetSearches(requestID)
|
||||
if err != nil {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrServerError.Code,
|
||||
ErrorDescription: err.Error(),
|
||||
}
|
||||
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
// If no searches found, return empty result
|
||||
if len(searches) == 0 {
|
||||
response.RespondWithSuccess(c, response.StatusOK, gin.H{
|
||||
"request_id": requestID,
|
||||
"references": []storetypes.Reference{},
|
||||
"total": 0,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Get authorized information and check permission using chat_id from first search
|
||||
authInfo := authorized.GetInfo(c)
|
||||
chatID := searches[0].ChatID
|
||||
if chatID != "" {
|
||||
hasPermission, err := checkChatPermission(chatStore, authInfo, chatID, true)
|
||||
if err != nil {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrServerError.Code,
|
||||
ErrorDescription: err.Error(),
|
||||
}
|
||||
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
if !hasPermission {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrAccessDenied.Code,
|
||||
ErrorDescription: "Forbidden: No permission to access these references",
|
||||
}
|
||||
response.RespondWithError(c, response.StatusForbidden, errorResp)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Collect all references from all searches
|
||||
var allRefs []storetypes.Reference
|
||||
for _, search := range searches {
|
||||
allRefs = append(allRefs, search.References...)
|
||||
}
|
||||
|
||||
response.RespondWithSuccess(c, response.StatusOK, gin.H{
|
||||
"request_id": requestID,
|
||||
"references": allRefs,
|
||||
"total": len(allRefs),
|
||||
})
|
||||
}
|
||||
|
||||
// GetReference retrieves a single reference by request ID and index
|
||||
// GET /v1/chat/references/:request_id/:index
|
||||
func GetReference(c *gin.Context) {
|
||||
// Get chat store
|
||||
chatStore := assistant.GetChatStore()
|
||||
if chatStore == nil {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrServerError.Code,
|
||||
ErrorDescription: "Chat storage not initialized",
|
||||
}
|
||||
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
// Get request ID from URL parameter
|
||||
requestID := c.Param("request_id")
|
||||
if requestID == "" {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrInvalidRequest.Code,
|
||||
ErrorDescription: "Request ID is required",
|
||||
}
|
||||
response.RespondWithError(c, response.StatusBadRequest, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
// Get index from URL parameter
|
||||
indexStr := c.Param("index")
|
||||
index, err := strconv.Atoi(indexStr)
|
||||
if err != nil || index < 1 {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrInvalidRequest.Code,
|
||||
ErrorDescription: "Invalid reference index, must be a positive integer",
|
||||
}
|
||||
response.RespondWithError(c, response.StatusBadRequest, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
// Get all search records to check permission first
|
||||
searches, err := chatStore.GetSearches(requestID)
|
||||
if err != nil {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrServerError.Code,
|
||||
ErrorDescription: err.Error(),
|
||||
}
|
||||
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
// Check permission using chat_id from first search
|
||||
if len(searches) > 0 {
|
||||
authInfo := authorized.GetInfo(c)
|
||||
chatID := searches[0].ChatID
|
||||
if chatID != "" {
|
||||
hasPermission, err := checkChatPermission(chatStore, authInfo, chatID, true)
|
||||
if err != nil {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrServerError.Code,
|
||||
ErrorDescription: err.Error(),
|
||||
}
|
||||
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
if !hasPermission {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrAccessDenied.Code,
|
||||
ErrorDescription: "Forbidden: No permission to access this reference",
|
||||
}
|
||||
response.RespondWithError(c, response.StatusForbidden, errorResp)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get the specific reference
|
||||
ref, err := chatStore.GetReference(requestID, index)
|
||||
if err != nil {
|
||||
errorResp := &response.ErrorResponse{
|
||||
Code: response.ErrServerError.Code,
|
||||
ErrorDescription: err.Error(),
|
||||
}
|
||||
response.RespondWithError(c, response.StatusNotFound, errorResp)
|
||||
return
|
||||
}
|
||||
|
||||
response.RespondWithSuccess(c, response.StatusOK, ref)
|
||||
}
|
||||
429
openapi/tests/chat/reference_test.go
Normal file
429
openapi/tests/chat/reference_test.go
Normal file
|
|
@ -0,0 +1,429 @@
|
|||
package openapi_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/yaoapp/yao/agent/assistant"
|
||||
storetypes "github.com/yaoapp/yao/agent/store/types"
|
||||
"github.com/yaoapp/yao/openapi"
|
||||
"github.com/yaoapp/yao/openapi/tests/testutils"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// Test Setup Helpers
|
||||
// =============================================================================
|
||||
|
||||
// createTestSearch creates a test search record in the database
|
||||
func createTestSearch(t *testing.T, requestID, chatID, query, source string, refs []storetypes.Reference) {
|
||||
chatStore := assistant.GetChatStore()
|
||||
if chatStore == nil {
|
||||
t.Skip("Chat store not initialized")
|
||||
}
|
||||
|
||||
search := &storetypes.Search{
|
||||
RequestID: requestID,
|
||||
ChatID: chatID,
|
||||
Query: query,
|
||||
Source: source,
|
||||
Duration: 100,
|
||||
References: refs,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
|
||||
err := chatStore.SaveSearch(search)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create test search: %v", err)
|
||||
}
|
||||
|
||||
t.Logf("Created test search: request_id=%s, query=%s", requestID, query)
|
||||
}
|
||||
|
||||
// cleanupTestSearches deletes test search records
|
||||
func cleanupTestSearches(t *testing.T, chatID string) {
|
||||
chatStore := assistant.GetChatStore()
|
||||
if chatStore == nil {
|
||||
return
|
||||
}
|
||||
|
||||
err := chatStore.DeleteSearches(chatID)
|
||||
if err != nil {
|
||||
t.Logf("Warning: Failed to cleanup test searches for chat %s: %v", chatID, err)
|
||||
} else {
|
||||
t.Logf("Cleaned up test searches for chat: %s", chatID)
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Get References Tests
|
||||
// =============================================================================
|
||||
|
||||
// TestGetReferences tests the get all references endpoint
|
||||
func TestGetReferences(t *testing.T) {
|
||||
serverURL := testutils.Prepare(t)
|
||||
defer testutils.Clean()
|
||||
|
||||
// Get base URL from server config
|
||||
baseURL := ""
|
||||
if openapi.Server != nil && openapi.Server.Config != nil {
|
||||
baseURL = openapi.Server.Config.BaseURL
|
||||
}
|
||||
|
||||
// Register test client and get token
|
||||
client := testutils.RegisterTestClient(t, "Reference Test Client", []string{"https://localhost/callback"})
|
||||
defer testutils.CleanupTestClient(t, client.ClientID)
|
||||
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
|
||||
|
||||
// Create test chat
|
||||
chatID := createTestChat(t, "Reference Test Chat", "test-assistant")
|
||||
defer cleanupTestChat(t, chatID)
|
||||
|
||||
requestID := fmt.Sprintf("req_%s", uuid.New().String())
|
||||
|
||||
// Create test search with references
|
||||
refs := []storetypes.Reference{
|
||||
{Index: 1, Type: "web", Title: "Go Documentation", URL: "https://golang.org/doc/", Snippet: "Go is an open source programming language", Content: "Full content 1"},
|
||||
{Index: 2, Type: "web", Title: "Go by Example", URL: "https://gobyexample.com/", Snippet: "Go by Example is a hands-on introduction", Content: "Full content 2"},
|
||||
}
|
||||
createTestSearch(t, requestID, chatID, "golang documentation", "web", refs)
|
||||
defer cleanupTestSearches(t, chatID)
|
||||
|
||||
// Create second search with more references
|
||||
refs2 := []storetypes.Reference{
|
||||
{Index: 3, Type: "kb", Title: "Internal Doc", Snippet: "Internal documentation snippet", Content: "Full content 3"},
|
||||
}
|
||||
createTestSearch(t, requestID, chatID, "internal docs", "kb", refs2)
|
||||
|
||||
t.Run("GetAllReferences", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID, nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
var result map[string]interface{}
|
||||
err = json.NewDecoder(resp.Body).Decode(&result)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, requestID, result["request_id"])
|
||||
assert.Equal(t, float64(3), result["total"])
|
||||
|
||||
references := result["references"].([]interface{})
|
||||
assert.Len(t, references, 3)
|
||||
|
||||
// Check first reference
|
||||
ref1 := references[0].(map[string]interface{})
|
||||
assert.Equal(t, float64(1), ref1["index"])
|
||||
assert.Equal(t, "web", ref1["type"])
|
||||
assert.Equal(t, "Go Documentation", ref1["title"])
|
||||
assert.Equal(t, "https://golang.org/doc/", ref1["url"])
|
||||
|
||||
// Check third reference (from second search)
|
||||
ref3 := references[2].(map[string]interface{})
|
||||
assert.Equal(t, float64(3), ref3["index"])
|
||||
assert.Equal(t, "kb", ref3["type"])
|
||||
assert.Equal(t, "Internal Doc", ref3["title"])
|
||||
|
||||
t.Logf("Successfully retrieved %d references", len(references))
|
||||
})
|
||||
|
||||
t.Run("GetReferences_NotFound", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/non_existent_request_id", nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Should return 200 with empty references
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
var result map[string]interface{}
|
||||
err = json.NewDecoder(resp.Body).Decode(&result)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, float64(0), result["total"])
|
||||
t.Log("Non-existent request returns empty references as expected")
|
||||
})
|
||||
|
||||
t.Run("GetReferences_Unauthorized", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID, nil)
|
||||
assert.NoError(t, err)
|
||||
// No Authorization header
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
|
||||
t.Log("Unauthorized request rejected as expected")
|
||||
})
|
||||
}
|
||||
|
||||
// TestGetReference tests the get single reference endpoint
|
||||
func TestGetReference(t *testing.T) {
|
||||
serverURL := testutils.Prepare(t)
|
||||
defer testutils.Clean()
|
||||
|
||||
// Get base URL from server config
|
||||
baseURL := ""
|
||||
if openapi.Server != nil && openapi.Server.Config != nil {
|
||||
baseURL = openapi.Server.Config.BaseURL
|
||||
}
|
||||
|
||||
// Register test client and get token
|
||||
client := testutils.RegisterTestClient(t, "Single Reference Test Client", []string{"https://localhost/callback"})
|
||||
defer testutils.CleanupTestClient(t, client.ClientID)
|
||||
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
|
||||
|
||||
// Create test chat
|
||||
chatID := createTestChat(t, "Single Reference Test Chat", "test-assistant")
|
||||
defer cleanupTestChat(t, chatID)
|
||||
|
||||
requestID := fmt.Sprintf("req_%s", uuid.New().String())
|
||||
|
||||
// Create test search with references
|
||||
refs := []storetypes.Reference{
|
||||
{Index: 1, Type: "web", Title: "First Reference", URL: "https://example.com/1", Snippet: "First snippet", Content: "First content"},
|
||||
{Index: 2, Type: "kb", Title: "Second Reference", Snippet: "Second snippet", Content: "Second content"},
|
||||
{Index: 3, Type: "db", Title: "Third Reference", Snippet: "Third snippet", Content: "Third content"},
|
||||
}
|
||||
createTestSearch(t, requestID, chatID, "test query", "web", refs)
|
||||
defer cleanupTestSearches(t, chatID)
|
||||
|
||||
t.Run("GetSingleReference", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID+"/2", nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
var ref map[string]interface{}
|
||||
err = json.NewDecoder(resp.Body).Decode(&ref)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, float64(2), ref["index"])
|
||||
assert.Equal(t, "kb", ref["type"])
|
||||
assert.Equal(t, "Second Reference", ref["title"])
|
||||
assert.Equal(t, "Second snippet", ref["snippet"])
|
||||
assert.Equal(t, "Second content", ref["content"])
|
||||
|
||||
t.Logf("Successfully retrieved reference at index 2")
|
||||
})
|
||||
|
||||
t.Run("GetReference_FirstIndex", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID+"/1", nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
var ref map[string]interface{}
|
||||
err = json.NewDecoder(resp.Body).Decode(&ref)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, float64(1), ref["index"])
|
||||
assert.Equal(t, "web", ref["type"])
|
||||
assert.Equal(t, "First Reference", ref["title"])
|
||||
|
||||
t.Log("Successfully retrieved first reference")
|
||||
})
|
||||
|
||||
t.Run("GetReference_NotFound", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID+"/999", nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusNotFound, resp.StatusCode)
|
||||
t.Log("Non-existent reference returns 404 as expected")
|
||||
})
|
||||
|
||||
t.Run("GetReference_InvalidIndex", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID+"/invalid", nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
|
||||
t.Log("Invalid index returns 400 as expected")
|
||||
})
|
||||
|
||||
t.Run("GetReference_ZeroIndex", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID+"/0", nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
|
||||
t.Log("Zero index returns 400 as expected")
|
||||
})
|
||||
|
||||
t.Run("GetReference_NegativeIndex", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID+"/-1", nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
|
||||
t.Log("Negative index returns 400 as expected")
|
||||
})
|
||||
|
||||
t.Run("GetReference_Unauthorized", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID+"/1", nil)
|
||||
assert.NoError(t, err)
|
||||
// No Authorization header
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
|
||||
t.Log("Unauthorized request rejected as expected")
|
||||
})
|
||||
}
|
||||
|
||||
// TestGetReferences_MultipleSearches tests references aggregation from multiple searches
|
||||
func TestGetReferences_MultipleSearches(t *testing.T) {
|
||||
serverURL := testutils.Prepare(t)
|
||||
defer testutils.Clean()
|
||||
|
||||
// Get base URL from server config
|
||||
baseURL := ""
|
||||
if openapi.Server != nil && openapi.Server.Config != nil {
|
||||
baseURL = openapi.Server.Config.BaseURL
|
||||
}
|
||||
|
||||
// Register test client and get token
|
||||
client := testutils.RegisterTestClient(t, "Multiple Searches Test Client", []string{"https://localhost/callback"})
|
||||
defer testutils.CleanupTestClient(t, client.ClientID)
|
||||
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
|
||||
|
||||
// Create test chat
|
||||
chatID := createTestChat(t, "Multiple Searches Test Chat", "test-assistant")
|
||||
defer cleanupTestChat(t, chatID)
|
||||
|
||||
requestID := fmt.Sprintf("req_%s", uuid.New().String())
|
||||
|
||||
// Create first search (web)
|
||||
refs1 := []storetypes.Reference{
|
||||
{Index: 1, Type: "web", Title: "Web Result 1", URL: "https://example.com/1"},
|
||||
{Index: 2, Type: "web", Title: "Web Result 2", URL: "https://example.com/2"},
|
||||
}
|
||||
createTestSearch(t, requestID, chatID, "web search query", "web", refs1)
|
||||
|
||||
// Create second search (kb)
|
||||
refs2 := []storetypes.Reference{
|
||||
{Index: 3, Type: "kb", Title: "KB Result 1"},
|
||||
{Index: 4, Type: "kb", Title: "KB Result 2"},
|
||||
}
|
||||
createTestSearch(t, requestID, chatID, "kb search query", "kb", refs2)
|
||||
|
||||
// Create third search (db)
|
||||
refs3 := []storetypes.Reference{
|
||||
{Index: 5, Type: "db", Title: "DB Result 1"},
|
||||
}
|
||||
createTestSearch(t, requestID, chatID, "db search query", "db", refs3)
|
||||
|
||||
defer cleanupTestSearches(t, chatID)
|
||||
|
||||
t.Run("AggregatedReferences", func(t *testing.T) {
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID, nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
var result map[string]interface{}
|
||||
err = json.NewDecoder(resp.Body).Decode(&result)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, float64(5), result["total"])
|
||||
|
||||
references := result["references"].([]interface{})
|
||||
assert.Len(t, references, 5)
|
||||
|
||||
// Verify all types are present
|
||||
types := make(map[string]int)
|
||||
for _, r := range references {
|
||||
ref := r.(map[string]interface{})
|
||||
refType := ref["type"].(string)
|
||||
types[refType]++
|
||||
}
|
||||
|
||||
assert.Equal(t, 2, types["web"])
|
||||
assert.Equal(t, 2, types["kb"])
|
||||
assert.Equal(t, 1, types["db"])
|
||||
|
||||
t.Logf("Successfully aggregated references: web=%d, kb=%d, db=%d", types["web"], types["kb"], types["db"])
|
||||
})
|
||||
|
||||
t.Run("GetSpecificReference", func(t *testing.T) {
|
||||
// Get reference from second search
|
||||
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/references/"+requestID+"/4", nil)
|
||||
assert.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, resp)
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
var ref map[string]interface{}
|
||||
err = json.NewDecoder(resp.Body).Decode(&ref)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, float64(4), ref["index"])
|
||||
assert.Equal(t, "kb", ref["type"])
|
||||
assert.Equal(t, "KB Result 2", ref["title"])
|
||||
|
||||
t.Log("Successfully retrieved specific reference from aggregated searches")
|
||||
})
|
||||
}
|
||||
|
|
@ -199,6 +199,7 @@ var testSystemModels = map[string]string{
|
|||
"__yao.agent.chat": "yao/models/agent/chat.mod.yao",
|
||||
"__yao.agent.message": "yao/models/agent/message.mod.yao",
|
||||
"__yao.agent.resume": "yao/models/agent/resume.mod.yao",
|
||||
"__yao.agent.search": "yao/models/agent/search.mod.yao",
|
||||
"__yao.attachment": "yao/models/attachment.mod.yao",
|
||||
"__yao.audit": "yao/models/audit.mod.yao",
|
||||
"__yao.config": "yao/models/config.mod.yao",
|
||||
|
|
|
|||
138
yao/models/agent/search.mod.yao
Normal file
138
yao/models/agent/search.mod.yao
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
{
|
||||
"name": "Search",
|
||||
"label": "Search",
|
||||
"description": "Search records for citation support and debugging",
|
||||
"tags": ["agent", "system"],
|
||||
"builtin": true,
|
||||
"readonly": true,
|
||||
"sort": 9999,
|
||||
"table": { "name": "agent_search", "comment": "Agent search table" },
|
||||
"columns": [
|
||||
{
|
||||
"name": "id",
|
||||
"type": "ID",
|
||||
"label": "ID",
|
||||
"comment": "Auto-increment primary key"
|
||||
},
|
||||
{
|
||||
"name": "request_id",
|
||||
"type": "string",
|
||||
"label": "Request ID",
|
||||
"comment": "Associated request ID",
|
||||
"length": 64,
|
||||
"nullable": false,
|
||||
"index": true
|
||||
},
|
||||
{
|
||||
"name": "chat_id",
|
||||
"type": "string",
|
||||
"label": "Chat ID",
|
||||
"comment": "Associated chat ID",
|
||||
"length": 64,
|
||||
"nullable": false,
|
||||
"index": true
|
||||
},
|
||||
{
|
||||
"name": "query",
|
||||
"type": "text",
|
||||
"label": "Query",
|
||||
"comment": "Original search query",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "config",
|
||||
"type": "json",
|
||||
"label": "Config",
|
||||
"comment": "Search config used (for tuning)",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "keywords",
|
||||
"type": "json",
|
||||
"label": "Keywords",
|
||||
"comment": "Extracted keywords (from NLP)",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "entities",
|
||||
"type": "json",
|
||||
"label": "Entities",
|
||||
"comment": "Extracted entities (for Graph search)",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "relations",
|
||||
"type": "json",
|
||||
"label": "Relations",
|
||||
"comment": "Extracted relations (for Graph search)",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "dsl",
|
||||
"type": "json",
|
||||
"label": "DSL",
|
||||
"comment": "Generated QueryDSL (for DB search)",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "source",
|
||||
"type": "string",
|
||||
"label": "Source",
|
||||
"comment": "Search source: web/kb/db/auto",
|
||||
"length": 32,
|
||||
"nullable": false
|
||||
},
|
||||
{
|
||||
"name": "references",
|
||||
"type": "json",
|
||||
"label": "References",
|
||||
"comment": "Reference[] with global index",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "graph",
|
||||
"type": "json",
|
||||
"label": "Graph",
|
||||
"comment": "GraphNode[] from knowledge graph",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "xml",
|
||||
"type": "text",
|
||||
"label": "XML",
|
||||
"comment": "Formatted XML for LLM context",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "prompt",
|
||||
"type": "text",
|
||||
"label": "Prompt",
|
||||
"comment": "Citation instruction prompt",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "duration",
|
||||
"type": "integer",
|
||||
"label": "Duration",
|
||||
"comment": "Search duration in milliseconds",
|
||||
"nullable": true
|
||||
},
|
||||
{
|
||||
"name": "error",
|
||||
"type": "text",
|
||||
"label": "Error",
|
||||
"comment": "Error message if failed",
|
||||
"nullable": true
|
||||
}
|
||||
],
|
||||
"relations": {
|
||||
"chat": {
|
||||
"type": "hasOne",
|
||||
"model": "__yao.agent.chat",
|
||||
"key": "chat_id",
|
||||
"foreign": "chat_id"
|
||||
}
|
||||
},
|
||||
"option": { "timestamps": true, "soft_deletes": true }
|
||||
}
|
||||
|
||||
Loading…
Add table
Reference in a new issue