diff --git a/agent/assistant/search.go b/agent/assistant/search.go index da9c725c..fd29451d 100644 --- a/agent/assistant/search.go +++ b/agent/assistant/search.go @@ -3,6 +3,7 @@ package assistant import ( "fmt" "strings" + "time" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" @@ -10,6 +11,7 @@ import ( "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" ) @@ -102,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 @@ -134,8 +137,6 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context query = optimizedQuery } } - // extractedKeywords will be used for storage in saveSearch() - TODO: Phase 1.9.5 - _ = extractedKeywords // Build search requests based on configuration requests := ast.buildSearchRequests(query, searchConfig) @@ -153,7 +154,10 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context // 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) @@ -163,6 +167,17 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context // === 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 } @@ -182,6 +197,16 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context // === 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 @@ -483,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 +} diff --git a/agent/search/citation.go b/agent/search/citation.go index 22a735ed..f286f96e 100644 --- a/agent/search/citation.go +++ b/agent/search/citation.go @@ -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:]) +} diff --git a/agent/search/citation_test.go b/agent/search/citation_test.go index 8fc855c1..3d43dbee 100644 --- a/agent/search/citation_test.go +++ b/agent/search/citation_test.go @@ -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) + } +} diff --git a/agent/search/reference.go b/agent/search/reference.go index 19a3ab5b..e969de15 100644 --- a/agent/search/reference.go +++ b/agent/search/reference.go @@ -9,7 +9,7 @@ import ( // DefaultCitationPrompt is the default prompt for citation instructions const DefaultCitationPrompt = `You have access to reference data in tags. Each 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: [{id}] -Example: According to the product data[ref_001], the price is $999.` +Example: According to the product data[1], the price is $999.` // BuildReferences converts search results to unified Reference format func BuildReferences(results []*types.Result) []*types.Reference { diff --git a/agent/search/reference_test.go b/agent/search/reference_test.go index 0c4182dd..177ef5b6 100644 --- a/agent/search/reference_test.go +++ b/agent/search/reference_test.go @@ -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{ "", "", - ``, + ``, "", "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{ - ``, + ``, "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{ - ``, + ``, "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{ "", "", - `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, `") - 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"`) diff --git a/agent/search/search_test.go b/agent/search/search_test.go index 8f730d66..1aa7b26c 100644 --- a/agent/search/search_test.go +++ b/agent/search/search_test.go @@ -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) } diff --git a/openapi/chat/chat.go b/openapi/chat/chat.go index 2121f2f5..df528a45 100644 --- a/openapi/chat/chat.go +++ b/openapi/chat/chat.go @@ -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) { diff --git a/openapi/chat/reference.go b/openapi/chat/reference.go new file mode 100644 index 00000000..26c3f178 --- /dev/null +++ b/openapi/chat/reference.go @@ -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) +} diff --git a/openapi/tests/chat/reference_test.go b/openapi/tests/chat/reference_test.go new file mode 100644 index 00000000..5daaf8ec --- /dev/null +++ b/openapi/tests/chat/reference_test.go @@ -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") + }) +}