09d368dae7
- Introduced a new unit test target for KB tests in the Makefile, allowing for dedicated testing of the KB module. - Updated the test folder selection logic to exclude additional AI-related components, ensuring focused testing. - Enhanced the GitHub Actions workflows to include KB tests, setting up necessary services like Qdrant, Neo4j, and MongoDB for a comprehensive testing environment. - Refactored search test functions to improve data existence checks and streamline test setup processes, enhancing test reliability and maintainability.
436 lines
11 KiB
Go
436 lines
11 KiB
Go
package api_test
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
graphragtypes "github.com/yaoapp/gou/graphrag/types"
|
|
"github.com/yaoapp/yao/kb"
|
|
"github.com/yaoapp/yao/kb/api"
|
|
)
|
|
|
|
// Note: TestMain is defined in collection_test.go
|
|
// Note: Test data setup is in search_setup_test.go
|
|
|
|
// ========== Search Query Tests ==========
|
|
|
|
// ensureTestDataExists ensures test collections exist by running setup if needed
|
|
// Setup will skip creation if data already exists
|
|
func ensureTestDataExists(t *testing.T, ctx context.Context) {
|
|
// Run setup - it checks if data exists and skips if already complete
|
|
TestSearchSetup(t)
|
|
}
|
|
|
|
func TestSearchQuery(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
ensureTestDataExists(t, ctx)
|
|
|
|
t.Run("VectorSearch_SingleCollection", func(t *testing.T) {
|
|
// Test: Simple vector search in science collection
|
|
// Query about Einstein should find Einstein-related documents
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "Who is Albert Einstein and what did he discover?",
|
|
Mode: api.SearchModeVector,
|
|
PageSize: 5,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Search error (may be expected if not implemented): %v", err)
|
|
return
|
|
}
|
|
|
|
if result == nil {
|
|
t.Skip("Search not implemented yet (returned nil)")
|
|
}
|
|
|
|
assert.Greater(t, len(result.Segments), 0, "Should find segments about Einstein")
|
|
t.Logf("Vector search returned %d segments", len(result.Segments))
|
|
|
|
// Verify relevance - top results should mention Einstein
|
|
for i, seg := range result.Segments {
|
|
t.Logf(" Segment %d (score: %.4f): %s...", i, seg.Score, truncateText(seg.Text, 100))
|
|
}
|
|
})
|
|
|
|
t.Run("VectorSearch_MultipleQueries", func(t *testing.T) {
|
|
// Test: Multiple queries in same collection, results should be merged
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "relativity theory",
|
|
Mode: api.SearchModeVector,
|
|
PageSize: 3,
|
|
},
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "Nobel Prize physics",
|
|
Mode: api.SearchModeVector,
|
|
PageSize: 3,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("Multi-query search returned %d merged segments", len(result.Segments))
|
|
})
|
|
|
|
t.Run("VectorSearch_CrossCollection", func(t *testing.T) {
|
|
// Test: Search across both collections
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "innovation and discovery",
|
|
Mode: api.SearchModeVector,
|
|
PageSize: 3,
|
|
},
|
|
{
|
|
CollectionID: SearchTestTechCollection,
|
|
Input: "technology innovation",
|
|
Mode: api.SearchModeVector,
|
|
PageSize: 3,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("Cross-collection search returned %d segments", len(result.Segments))
|
|
})
|
|
|
|
t.Run("ExpandSearch_EntityExpansion", func(t *testing.T) {
|
|
// Test: Expand mode should find related entities through graph
|
|
// Query: "photoelectric effect" should expand to find:
|
|
// - Einstein (discovered it)
|
|
// - Nobel Prize (awarded for it)
|
|
// - Quantum mechanics (built upon it)
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "photoelectric effect",
|
|
Mode: api.SearchModeExpand,
|
|
MaxDepth: 2,
|
|
PageSize: 5,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Expand search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("Expand search returned %d segments", len(result.Segments))
|
|
|
|
// Check if graph data is returned
|
|
if result.Graph != nil {
|
|
t.Logf(" Graph nodes: %d, relationships: %d",
|
|
len(result.Graph.Nodes), len(result.Graph.Relationships))
|
|
}
|
|
|
|
// Verify expanded results include related entities
|
|
for i, seg := range result.Segments {
|
|
t.Logf(" Segment %d (score: %.4f): %s...", i, seg.Score, truncateText(seg.Text, 100))
|
|
}
|
|
})
|
|
|
|
t.Run("ExpandSearch_DeepAssociation", func(t *testing.T) {
|
|
// Test: Deep association through entity relationships
|
|
// Query: "Germany physics" should expand to find:
|
|
// - Einstein (born in Germany, physicist)
|
|
// - Relativity (Einstein's theory)
|
|
// - Planck (German physicist, quantum theory)
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "German physicist contributions",
|
|
Mode: api.SearchModeExpand,
|
|
MaxDepth: 3,
|
|
PageSize: 5,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Deep expand search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("Deep expand search returned %d segments", len(result.Segments))
|
|
})
|
|
|
|
t.Run("GraphSearch_EntityTraversal", func(t *testing.T) {
|
|
// Test: Pure graph search - find segments through entity relationships
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestTechCollection,
|
|
Input: "Steve Jobs",
|
|
Mode: api.SearchModeGraph,
|
|
MaxDepth: 2,
|
|
PageSize: 5,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Graph search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("Graph search returned %d segments", len(result.Segments))
|
|
|
|
if result.Graph != nil {
|
|
t.Logf(" Found %d nodes, %d relationships",
|
|
len(result.Graph.Nodes), len(result.Graph.Relationships))
|
|
for _, node := range result.Graph.Nodes {
|
|
t.Logf(" Node: %s (%s)", node.ID, node.EntityType)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("Search_WithMessages", func(t *testing.T) {
|
|
// Test: Search using conversation history instead of direct input
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Messages: []graphragtypes.ChatMessage{
|
|
{Role: "user", Content: "Tell me about famous physicists"},
|
|
{Role: "assistant", Content: "There are many famous physicists throughout history..."},
|
|
{Role: "user", Content: "What about Einstein specifically?"},
|
|
},
|
|
Mode: api.SearchModeVector,
|
|
PageSize: 5,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Message-based search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("Message-based search returned %d segments", len(result.Segments))
|
|
})
|
|
|
|
t.Run("Search_WithDocumentFilter", func(t *testing.T) {
|
|
// Test: Search within a specific document
|
|
// First, get a document ID
|
|
filter := &api.ListDocumentsFilter{
|
|
Page: 1,
|
|
PageSize: 1,
|
|
CollectionID: SearchTestScienceCollection,
|
|
}
|
|
listResult, err := kb.API.ListDocuments(ctx, filter)
|
|
if err != nil || len(listResult.Data) == 0 {
|
|
t.Skip("No documents available for filter test")
|
|
}
|
|
|
|
docID, ok := listResult.Data[0]["document_id"].(string)
|
|
if !ok {
|
|
t.Skip("Could not get document ID")
|
|
}
|
|
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
DocumentID: docID,
|
|
Input: "physics discovery",
|
|
Mode: api.SearchModeVector,
|
|
PageSize: 5,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Document-filtered search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("Document-filtered search returned %d segments", len(result.Segments))
|
|
|
|
// Verify all results are from the specified document
|
|
for _, seg := range result.Segments {
|
|
if seg.DocumentID != "" {
|
|
assert.Equal(t, docID, seg.DocumentID, "All segments should be from filtered document")
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("Search_WithPagination", func(t *testing.T) {
|
|
// Test: Pagination
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "physics",
|
|
Mode: api.SearchModeVector,
|
|
Page: 1,
|
|
PageSize: 2,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Paginated search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
assert.LessOrEqual(t, len(result.Segments), 2, "Should respect page size")
|
|
t.Logf("Page 1: %d segments, Total: %d, TotalPages: %d",
|
|
len(result.Segments), result.Total, result.TotalPages)
|
|
|
|
// Get page 2
|
|
queries[0].Page = 2
|
|
result2, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Page 2 search error: %v", err)
|
|
return
|
|
}
|
|
|
|
if result2 != nil && len(result2.Segments) > 0 {
|
|
t.Logf("Page 2: %d segments", len(result2.Segments))
|
|
}
|
|
})
|
|
|
|
t.Run("Search_WithMinScore", func(t *testing.T) {
|
|
// Test: Filter by minimum score
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "Einstein relativity",
|
|
Mode: api.SearchModeVector,
|
|
MinScore: 0.5,
|
|
PageSize: 10,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("MinScore search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("MinScore search returned %d segments", len(result.Segments))
|
|
|
|
// Verify all results meet minimum score
|
|
for _, seg := range result.Segments {
|
|
assert.GreaterOrEqual(t, seg.Score, 0.5, "All segments should meet minimum score")
|
|
}
|
|
})
|
|
|
|
t.Run("Search_WithMetadataFilter", func(t *testing.T) {
|
|
// Test: Filter by metadata
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: SearchTestScienceCollection,
|
|
Input: "physics",
|
|
Mode: api.SearchModeVector,
|
|
Metadata: map[string]interface{}{
|
|
"title": "Albert Einstein Biography",
|
|
},
|
|
PageSize: 10,
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.Search(ctx, queries)
|
|
if err != nil {
|
|
t.Logf("Metadata filter search error: %v", err)
|
|
return
|
|
}
|
|
|
|
assert.NotNil(t, result)
|
|
t.Logf("Metadata-filtered search returned %d segments", len(result.Segments))
|
|
})
|
|
}
|
|
|
|
// ========== Error Handling Tests ==========
|
|
|
|
func TestSearchErrorHandling(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
|
|
t.Run("EmptyQueries", func(t *testing.T) {
|
|
result, err := kb.API.Search(ctx, []api.Query{})
|
|
// Empty queries should return empty result or error
|
|
if err != nil {
|
|
assert.Contains(t, err.Error(), "required")
|
|
} else {
|
|
assert.NotNil(t, result)
|
|
assert.Equal(t, 0, len(result.Segments))
|
|
}
|
|
})
|
|
|
|
t.Run("MissingCollectionID", func(t *testing.T) {
|
|
queries := []api.Query{
|
|
{
|
|
Input: "test query",
|
|
Mode: api.SearchModeVector,
|
|
},
|
|
}
|
|
|
|
_, err := kb.API.Search(ctx, queries)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "collection")
|
|
})
|
|
|
|
t.Run("MissingInputAndMessages", func(t *testing.T) {
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: "some_collection",
|
|
Mode: api.SearchModeVector,
|
|
},
|
|
}
|
|
|
|
_, err := kb.API.Search(ctx, queries)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "input")
|
|
})
|
|
|
|
t.Run("NonexistentCollection", func(t *testing.T) {
|
|
queries := []api.Query{
|
|
{
|
|
CollectionID: "nonexistent_collection_xyz",
|
|
Input: "test query",
|
|
Mode: api.SearchModeVector,
|
|
},
|
|
}
|
|
|
|
_, err := kb.API.Search(ctx, queries)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
// ========== Helper Functions ==========
|
|
|
|
func truncateText(text string, maxLen int) string {
|
|
if len(text) <= maxLen {
|
|
return text
|
|
}
|
|
return text[:maxLen] + "..."
|
|
}
|