Refactor Citation Handling and Enhance Search Result Storage
- Updated the CitationGenerator to produce simple integer IDs instead of formatted strings, improving clarity and consistency in citation references. - Enhanced the executeAutoSearch method to save both successful and failed search results, capturing detailed execution data for better traceability. - Introduced a new SearchExecutionResult type to structure search result data, including query, keywords, configuration, duration, and error information. - Updated related tests to reflect changes in citation ID format and ensure proper functionality of the new storage mechanisms. - Revised documentation to clarify the new citation format and search result handling processes.
This commit is contained in:
Max
parent
00f9d86788
commit
6714ef9813
9 changed files with 928 additions and 66 deletions
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
})
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue