6714ef9813
- 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.
486 lines
11 KiB
Go
486 lines
11 KiB
Go
package search
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/yaoapp/yao/agent/search/types"
|
|
)
|
|
|
|
func TestBuildReferences(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
results []*types.Result
|
|
expected int
|
|
}{
|
|
{
|
|
name: "nil results",
|
|
results: nil,
|
|
expected: 0,
|
|
},
|
|
{
|
|
name: "empty results",
|
|
results: []*types.Result{},
|
|
expected: 0,
|
|
},
|
|
{
|
|
name: "single result with items",
|
|
results: []*types.Result{
|
|
{
|
|
Type: types.SearchTypeWeb,
|
|
Query: "test query",
|
|
Items: []*types.ResultItem{
|
|
{
|
|
CitationID: "1",
|
|
Type: types.SearchTypeWeb,
|
|
Source: types.SourceAuto,
|
|
Weight: 0.6,
|
|
Score: 0.9,
|
|
Title: "Test Title",
|
|
Content: "Test content",
|
|
URL: "https://example.com",
|
|
},
|
|
{
|
|
CitationID: "2",
|
|
Type: types.SearchTypeWeb,
|
|
Source: types.SourceAuto,
|
|
Weight: 0.6,
|
|
Score: 0.8,
|
|
Title: "Test Title 2",
|
|
Content: "Test content 2",
|
|
URL: "https://example2.com",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
expected: 2,
|
|
},
|
|
{
|
|
name: "multiple results",
|
|
results: []*types.Result{
|
|
{
|
|
Type: types.SearchTypeWeb,
|
|
Items: []*types.ResultItem{
|
|
{CitationID: "1", Type: types.SearchTypeWeb, Content: "Web content"},
|
|
},
|
|
},
|
|
{
|
|
Type: types.SearchTypeKB,
|
|
Items: []*types.ResultItem{
|
|
{CitationID: "2", Type: types.SearchTypeKB, Content: "KB content"},
|
|
},
|
|
},
|
|
{
|
|
Type: types.SearchTypeDB,
|
|
Items: []*types.ResultItem{
|
|
{CitationID: "3", Type: types.SearchTypeDB, Content: "DB content"},
|
|
},
|
|
},
|
|
},
|
|
expected: 3,
|
|
},
|
|
{
|
|
name: "result with nil items",
|
|
results: []*types.Result{
|
|
{
|
|
Type: types.SearchTypeWeb,
|
|
Items: []*types.ResultItem{
|
|
{CitationID: "1", Content: "Content 1"},
|
|
nil,
|
|
{CitationID: "2", Content: "Content 2"},
|
|
},
|
|
},
|
|
},
|
|
expected: 2,
|
|
},
|
|
{
|
|
name: "nil result in slice",
|
|
results: []*types.Result{
|
|
{
|
|
Type: types.SearchTypeWeb,
|
|
Items: []*types.ResultItem{
|
|
{CitationID: "1", Content: "Content"},
|
|
},
|
|
},
|
|
nil,
|
|
{
|
|
Type: types.SearchTypeKB,
|
|
Items: []*types.ResultItem{
|
|
{CitationID: "2", Content: "Content 2"},
|
|
},
|
|
},
|
|
},
|
|
expected: 2,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
refs := BuildReferences(tt.results)
|
|
assert.Equal(t, tt.expected, len(refs))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuildReferences_FieldMapping(t *testing.T) {
|
|
item := &types.ResultItem{
|
|
CitationID: "1",
|
|
Type: types.SearchTypeWeb,
|
|
Source: types.SourceHook,
|
|
Weight: 0.8,
|
|
Score: 0.95,
|
|
Title: "Test Title",
|
|
Content: "Test Content",
|
|
URL: "https://example.com",
|
|
}
|
|
|
|
results := []*types.Result{
|
|
{Items: []*types.ResultItem{item}},
|
|
}
|
|
|
|
refs := BuildReferences(results)
|
|
assert.Equal(t, 1, len(refs))
|
|
|
|
ref := refs[0]
|
|
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)
|
|
assert.Equal(t, 0.95, ref.Score)
|
|
assert.Equal(t, "Test Title", ref.Title)
|
|
assert.Equal(t, "Test Content", ref.Content)
|
|
assert.Equal(t, "https://example.com", ref.URL)
|
|
}
|
|
|
|
func TestFormatReferencesXML(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
refs []*types.Reference
|
|
contains []string
|
|
excludes []string
|
|
}{
|
|
{
|
|
name: "nil refs",
|
|
refs: nil,
|
|
contains: []string{},
|
|
excludes: []string{"<references>"},
|
|
},
|
|
{
|
|
name: "empty refs",
|
|
refs: []*types.Reference{},
|
|
contains: []string{},
|
|
excludes: []string{"<references>"},
|
|
},
|
|
{
|
|
name: "single ref with all fields",
|
|
refs: []*types.Reference{
|
|
{
|
|
ID: "1",
|
|
Type: types.SearchTypeWeb,
|
|
Source: types.SourceUser,
|
|
Weight: 1.0,
|
|
Score: 0.9,
|
|
Title: "Test Title",
|
|
Content: "Test Content",
|
|
URL: "https://example.com",
|
|
},
|
|
},
|
|
contains: []string{
|
|
"<references>",
|
|
"</references>",
|
|
`<ref id="1" type="web" weight="1.0" source="user">`,
|
|
"</ref>",
|
|
"Test Title",
|
|
"Test Content",
|
|
"URL: https://example.com",
|
|
},
|
|
},
|
|
{
|
|
name: "ref without title",
|
|
refs: []*types.Reference{
|
|
{
|
|
ID: "1",
|
|
Type: types.SearchTypeKB,
|
|
Source: types.SourceHook,
|
|
Weight: 0.8,
|
|
Content: "Content without title",
|
|
},
|
|
},
|
|
contains: []string{
|
|
`<ref id="1" type="kb" weight="0.8" source="hook">`,
|
|
"Content without title",
|
|
},
|
|
excludes: []string{
|
|
"URL:",
|
|
},
|
|
},
|
|
{
|
|
name: "ref without URL",
|
|
refs: []*types.Reference{
|
|
{
|
|
ID: "1",
|
|
Type: types.SearchTypeDB,
|
|
Source: types.SourceAuto,
|
|
Weight: 0.6,
|
|
Title: "DB Record",
|
|
Content: "Database content",
|
|
},
|
|
},
|
|
contains: []string{
|
|
`<ref id="1" type="db" weight="0.6" source="auto">`,
|
|
"DB Record",
|
|
"Database content",
|
|
},
|
|
excludes: []string{
|
|
"URL:",
|
|
},
|
|
},
|
|
{
|
|
name: "multiple refs",
|
|
refs: []*types.Reference{
|
|
{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="1"`,
|
|
`id="2"`,
|
|
`id="3"`,
|
|
"Content 1",
|
|
"Content 2",
|
|
"Content 3",
|
|
},
|
|
},
|
|
{
|
|
name: "nil ref in slice",
|
|
refs: []*types.Reference{
|
|
{ID: "1", Type: types.SearchTypeWeb, Weight: 1.0, Content: "Content 1"},
|
|
nil,
|
|
{ID: "2", Type: types.SearchTypeKB, Weight: 0.8, Content: "Content 2"},
|
|
},
|
|
contains: []string{
|
|
`id="1"`,
|
|
`id="2"`,
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
xml := FormatReferencesXML(tt.refs)
|
|
|
|
for _, s := range tt.contains {
|
|
assert.Contains(t, xml, s, "expected XML to contain: %s", s)
|
|
}
|
|
|
|
for _, s := range tt.excludes {
|
|
assert.NotContains(t, xml, s, "expected XML to not contain: %s", s)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFormatReferencesXML_Structure(t *testing.T) {
|
|
refs := []*types.Reference{
|
|
{
|
|
ID: "1",
|
|
Type: types.SearchTypeWeb,
|
|
Source: types.SourceUser,
|
|
Weight: 1.0,
|
|
Title: "Title",
|
|
Content: "Content",
|
|
URL: "https://example.com",
|
|
},
|
|
}
|
|
|
|
xml := FormatReferencesXML(refs)
|
|
|
|
// Check structure
|
|
assert.True(t, strings.HasPrefix(xml, "<references>\n"))
|
|
assert.True(t, strings.HasSuffix(xml, "</references>"))
|
|
assert.Contains(t, xml, "</ref>\n")
|
|
}
|
|
|
|
func TestGetCitationPrompt(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
cfg *types.CitationConfig
|
|
expected string
|
|
}{
|
|
{
|
|
name: "nil config",
|
|
cfg: nil,
|
|
expected: DefaultCitationPrompt,
|
|
},
|
|
{
|
|
name: "empty config",
|
|
cfg: &types.CitationConfig{},
|
|
expected: DefaultCitationPrompt,
|
|
},
|
|
{
|
|
name: "config with custom prompt",
|
|
cfg: &types.CitationConfig{
|
|
CustomPrompt: "Custom citation instructions",
|
|
},
|
|
expected: "Custom citation instructions",
|
|
},
|
|
{
|
|
name: "config with empty custom prompt",
|
|
cfg: &types.CitationConfig{
|
|
CustomPrompt: "",
|
|
},
|
|
expected: DefaultCitationPrompt,
|
|
},
|
|
{
|
|
name: "config with format but no custom prompt",
|
|
cfg: &types.CitationConfig{
|
|
Format: "[{id}]",
|
|
},
|
|
expected: DefaultCitationPrompt,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
prompt := GetCitationPrompt(tt.cfg)
|
|
assert.Equal(t, tt.expected, prompt)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestDefaultCitationPrompt(t *testing.T) {
|
|
// Verify default prompt contains key instructions
|
|
assert.Contains(t, DefaultCitationPrompt, "<references>")
|
|
assert.Contains(t, DefaultCitationPrompt, "id: Citation identifier")
|
|
assert.Contains(t, DefaultCitationPrompt, "type: Data type")
|
|
assert.Contains(t, DefaultCitationPrompt, "weight: Relevance weight")
|
|
assert.Contains(t, DefaultCitationPrompt, "source: Origin")
|
|
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) {
|
|
results := []*types.Result{
|
|
{
|
|
Type: types.SearchTypeWeb,
|
|
Items: []*types.ResultItem{
|
|
{
|
|
CitationID: "1",
|
|
Type: types.SearchTypeWeb,
|
|
Source: types.SourceAuto,
|
|
Weight: 0.6,
|
|
Title: "Test",
|
|
Content: "Content",
|
|
URL: "https://example.com",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
t.Run("with nil config", func(t *testing.T) {
|
|
ctx := BuildReferenceContext(results, nil)
|
|
|
|
assert.NotNil(t, ctx)
|
|
assert.Equal(t, 1, len(ctx.References))
|
|
assert.Contains(t, ctx.XML, "<references>")
|
|
assert.Contains(t, ctx.XML, `id="1"`)
|
|
assert.Equal(t, DefaultCitationPrompt, ctx.Prompt)
|
|
})
|
|
|
|
t.Run("with custom prompt config", func(t *testing.T) {
|
|
cfg := &types.CitationConfig{
|
|
CustomPrompt: "Custom prompt",
|
|
}
|
|
ctx := BuildReferenceContext(results, cfg)
|
|
|
|
assert.NotNil(t, ctx)
|
|
assert.Equal(t, "Custom prompt", ctx.Prompt)
|
|
})
|
|
|
|
t.Run("with empty results", func(t *testing.T) {
|
|
ctx := BuildReferenceContext([]*types.Result{}, nil)
|
|
|
|
assert.NotNil(t, ctx)
|
|
assert.Equal(t, 0, len(ctx.References))
|
|
assert.Equal(t, "", ctx.XML)
|
|
assert.Equal(t, DefaultCitationPrompt, ctx.Prompt)
|
|
})
|
|
}
|
|
|
|
func TestBuildReferenceContext_Integration(t *testing.T) {
|
|
// Simulate a real-world scenario with multiple search types
|
|
results := []*types.Result{
|
|
{
|
|
Type: types.SearchTypeWeb,
|
|
Query: "AI developments",
|
|
Items: []*types.ResultItem{
|
|
{
|
|
CitationID: "1",
|
|
Type: types.SearchTypeWeb,
|
|
Source: types.SourceAuto,
|
|
Weight: 0.6,
|
|
Score: 0.95,
|
|
Title: "OpenAI Announces GPT-5",
|
|
Content: "OpenAI has announced the development of GPT-5...",
|
|
URL: "https://news.example.com/gpt5",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Type: types.SearchTypeKB,
|
|
Query: "AI developments",
|
|
Items: []*types.ResultItem{
|
|
{
|
|
CitationID: "2",
|
|
Type: types.SearchTypeKB,
|
|
Source: types.SourceHook,
|
|
Weight: 0.8,
|
|
Score: 0.88,
|
|
Title: "Internal AI Research Notes",
|
|
Content: "Our internal research on AI capabilities...",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Type: types.SearchTypeDB,
|
|
Query: "AI developments",
|
|
Items: []*types.ResultItem{
|
|
{
|
|
CitationID: "3",
|
|
Type: types.SearchTypeDB,
|
|
Source: types.SourceUser,
|
|
Weight: 1.0,
|
|
Score: 0.92,
|
|
Title: "Product: AI Assistant",
|
|
Content: "Name: AI Assistant\nPrice: $99\nCategory: Software",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
ctx := BuildReferenceContext(results, nil)
|
|
|
|
// Verify all references are included
|
|
assert.Equal(t, 3, len(ctx.References))
|
|
|
|
// Verify XML contains all references
|
|
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"`)
|
|
assert.Contains(t, ctx.XML, `source="hook"`)
|
|
assert.Contains(t, ctx.XML, `source="user"`)
|
|
|
|
// Verify different search types are represented
|
|
assert.Contains(t, ctx.XML, `type="web"`)
|
|
assert.Contains(t, ctx.XML, `type="kb"`)
|
|
assert.Contains(t, ctx.XML, `type="db"`)
|
|
}
|