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.
556 lines
14 KiB
Go
556 lines
14 KiB
Go
package api
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sort"
|
|
"sync"
|
|
|
|
graphragtypes "github.com/yaoapp/gou/graphrag/types"
|
|
"github.com/yaoapp/kun/log"
|
|
"github.com/yaoapp/yao/kb/providers/factory"
|
|
)
|
|
|
|
// Default search parameters
|
|
const (
|
|
DefaultSearchK = 10
|
|
DefaultMaxDepth = 2
|
|
DefaultMinScore = 0.0
|
|
MaxSearchK = 100
|
|
DefaultSearchPageSize = 20
|
|
)
|
|
|
|
// Search performs batch search operations on the knowledge base
|
|
// Queries can span multiple collections; implementation groups by CollectionID
|
|
// Mode, providers (embedding/extraction/reranker) are read from each collection's config
|
|
// All results are merged and deduplicated
|
|
func (kb *KBInstance) Search(ctx context.Context, queries []Query) (*SearchResult, error) {
|
|
if len(queries) == 0 {
|
|
return &SearchResult{
|
|
Segments: []graphragtypes.Segment{},
|
|
Total: 0,
|
|
}, nil
|
|
}
|
|
|
|
// 1. Validate queries
|
|
if err := kb.validateQueries(queries); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 2. Group queries by CollectionID
|
|
groupedQueries := kb.groupQueriesByCollection(queries)
|
|
|
|
// 3. Process each collection group in parallel
|
|
var (
|
|
allSegments []graphragtypes.Segment
|
|
allGraph *GraphData
|
|
mu sync.Mutex
|
|
wg sync.WaitGroup
|
|
errChan = make(chan error, len(groupedQueries))
|
|
)
|
|
|
|
for collectionID, collQueries := range groupedQueries {
|
|
wg.Add(1)
|
|
go func(collID string, qs []Query) {
|
|
defer wg.Done()
|
|
|
|
segments, graph, err := kb.searchCollection(ctx, collID, qs)
|
|
if err != nil {
|
|
errChan <- fmt.Errorf("search in collection %s failed: %w", collID, err)
|
|
return
|
|
}
|
|
|
|
mu.Lock()
|
|
allSegments = append(allSegments, segments...)
|
|
if graph != nil {
|
|
allGraph = mergeGraphData(allGraph, graph)
|
|
}
|
|
mu.Unlock()
|
|
}(collectionID, collQueries)
|
|
}
|
|
|
|
wg.Wait()
|
|
close(errChan)
|
|
|
|
// Collect errors
|
|
var errs []error
|
|
for err := range errChan {
|
|
errs = append(errs, err)
|
|
}
|
|
if len(errs) > 0 {
|
|
// Combine all error messages
|
|
errMsgs := make([]string, len(errs))
|
|
for i, e := range errs {
|
|
errMsgs[i] = e.Error()
|
|
}
|
|
return nil, fmt.Errorf("search failed: %v", errMsgs)
|
|
}
|
|
|
|
// 4. Merge and deduplicate results
|
|
mergedSegments := kb.deduplicateSegments(allSegments)
|
|
|
|
// 5. Sort by score (descending)
|
|
sort.Slice(mergedSegments, func(i, j int) bool {
|
|
return mergedSegments[i].Score > mergedSegments[j].Score
|
|
})
|
|
|
|
// 6. Apply pagination from first query (if specified)
|
|
result := kb.applyPagination(mergedSegments, queries[0])
|
|
result.Graph = allGraph
|
|
|
|
return result, nil
|
|
}
|
|
|
|
// ========== Validation ==========
|
|
|
|
// validateQueries validates all queries
|
|
func (kb *KBInstance) validateQueries(queries []Query) error {
|
|
for i, q := range queries {
|
|
if q.CollectionID == "" {
|
|
return fmt.Errorf("query %d: collection_id is required", i)
|
|
}
|
|
if q.Input == "" && len(q.Messages) == 0 {
|
|
return fmt.Errorf("query %d: either input or messages is required", i)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ========== Query Grouping ==========
|
|
|
|
// groupQueriesByCollection groups queries by their CollectionID
|
|
func (kb *KBInstance) groupQueriesByCollection(queries []Query) map[string][]Query {
|
|
grouped := make(map[string][]Query)
|
|
for _, q := range queries {
|
|
grouped[q.CollectionID] = append(grouped[q.CollectionID], q)
|
|
}
|
|
return grouped
|
|
}
|
|
|
|
// ========== Collection Search ==========
|
|
|
|
// searchCollection processes all queries for a single collection
|
|
func (kb *KBInstance) searchCollection(ctx context.Context, collectionID string, queries []Query) ([]graphragtypes.Segment, *GraphData, error) {
|
|
// Get collection config
|
|
collection, err := kb.GetCollection(ctx, collectionID)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("failed to get collection: %w", err)
|
|
}
|
|
|
|
// Get embedding provider from collection config
|
|
embeddingProviderID, _ := collection["embedding_provider_id"].(string)
|
|
embeddingOptionID, _ := collection["embedding_option_id"].(string)
|
|
if embeddingProviderID == "" || embeddingOptionID == "" {
|
|
return nil, nil, fmt.Errorf("collection %s missing embedding configuration", collectionID)
|
|
}
|
|
|
|
// Create embedding function
|
|
embedding, err := kb.createEmbedding(embeddingProviderID, embeddingOptionID, "en")
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("failed to create embedding: %w", err)
|
|
}
|
|
|
|
var (
|
|
allSegments []graphragtypes.Segment
|
|
allGraph *GraphData
|
|
mu sync.Mutex
|
|
wg sync.WaitGroup
|
|
errChan = make(chan error, len(queries))
|
|
)
|
|
|
|
// Process queries in parallel
|
|
for _, query := range queries {
|
|
wg.Add(1)
|
|
go func(q Query) {
|
|
defer wg.Done()
|
|
|
|
segments, graph, err := kb.executeQuery(ctx, collectionID, q, embedding, collection)
|
|
if err != nil {
|
|
errChan <- err
|
|
return
|
|
}
|
|
|
|
mu.Lock()
|
|
allSegments = append(allSegments, segments...)
|
|
if graph != nil {
|
|
allGraph = mergeGraphData(allGraph, graph)
|
|
}
|
|
mu.Unlock()
|
|
}(query)
|
|
}
|
|
|
|
wg.Wait()
|
|
close(errChan)
|
|
|
|
// Collect errors
|
|
var errors []error
|
|
for err := range errChan {
|
|
errors = append(errors, err)
|
|
}
|
|
if len(errors) > 0 {
|
|
return allSegments, allGraph, errors[0]
|
|
}
|
|
|
|
return allSegments, allGraph, nil
|
|
}
|
|
|
|
// executeQuery executes a single query based on its mode
|
|
func (kb *KBInstance) executeQuery(ctx context.Context, collectionID string, query Query, embedding graphragtypes.Embedding, collection map[string]interface{}) ([]graphragtypes.Segment, *GraphData, error) {
|
|
// Determine search mode
|
|
mode := query.Mode
|
|
if mode == "" {
|
|
// Default to expand mode
|
|
mode = SearchModeExpand
|
|
}
|
|
|
|
// Get query text
|
|
queryText := kb.getQueryText(query)
|
|
if queryText == "" {
|
|
return nil, nil, fmt.Errorf("no query text found")
|
|
}
|
|
|
|
// Execute based on mode
|
|
switch mode {
|
|
case SearchModeVector:
|
|
return kb.searchVector(ctx, collectionID, queryText, query, embedding)
|
|
case SearchModeGraph:
|
|
return kb.searchGraph(ctx, collectionID, queryText, query, collection)
|
|
case SearchModeExpand:
|
|
return kb.searchExpand(ctx, collectionID, queryText, query, embedding, collection)
|
|
default:
|
|
return nil, nil, fmt.Errorf("unknown search mode: %s", mode)
|
|
}
|
|
}
|
|
|
|
// getQueryText extracts query text from Input or Messages
|
|
func (kb *KBInstance) getQueryText(query Query) string {
|
|
// Input takes precedence
|
|
if query.Input != "" {
|
|
return query.Input
|
|
}
|
|
|
|
// Extract from last user message
|
|
for i := len(query.Messages) - 1; i >= 0; i-- {
|
|
if query.Messages[i].Role == "user" {
|
|
return query.Messages[i].Content
|
|
}
|
|
}
|
|
|
|
return ""
|
|
}
|
|
|
|
// ========== Vector Search ==========
|
|
|
|
// searchVector performs pure vector similarity search
|
|
func (kb *KBInstance) searchVector(ctx context.Context, collectionID string, queryText string, query Query, embedding graphragtypes.Embedding) ([]graphragtypes.Segment, *GraphData, error) {
|
|
// Build search options
|
|
k := query.PageSize
|
|
if k <= 0 {
|
|
k = DefaultSearchK
|
|
}
|
|
if k > MaxSearchK {
|
|
k = MaxSearchK
|
|
}
|
|
|
|
options := &graphragtypes.VectorSearchOptions{
|
|
CollectionID: collectionID,
|
|
DocumentID: query.DocumentID,
|
|
Query: queryText,
|
|
K: k,
|
|
MinScore: query.MinScore,
|
|
Embedding: embedding,
|
|
}
|
|
|
|
// Add metadata filter
|
|
if len(query.Metadata) > 0 {
|
|
options.Filter = query.Metadata
|
|
}
|
|
|
|
// Execute search
|
|
result, err := kb.GraphRag.SearchVector(ctx, options)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("vector search failed: %w", err)
|
|
}
|
|
return result.Segments, nil, nil
|
|
}
|
|
|
|
// ========== Graph Search ==========
|
|
|
|
// searchGraph performs pure graph traversal search
|
|
func (kb *KBInstance) searchGraph(ctx context.Context, collectionID string, queryText string, query Query, collection map[string]interface{}) ([]graphragtypes.Segment, *GraphData, error) {
|
|
// Get extraction provider for entity extraction
|
|
extraction, err := kb.createExtraction(collection)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("failed to create extraction: %w", err)
|
|
}
|
|
|
|
// Build graph search options
|
|
maxDepth := query.MaxDepth
|
|
if maxDepth <= 0 {
|
|
maxDepth = DefaultMaxDepth
|
|
}
|
|
|
|
options := &graphragtypes.GraphSearchOptions{
|
|
CollectionID: collectionID,
|
|
DocumentID: query.DocumentID,
|
|
Query: queryText,
|
|
MaxDepth: maxDepth,
|
|
Extraction: extraction,
|
|
}
|
|
|
|
// Execute search
|
|
result, err := kb.GraphRag.SearchGraph(ctx, options)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("graph search failed: %w", err)
|
|
}
|
|
|
|
// Convert to GraphData
|
|
graph := &GraphData{
|
|
Nodes: result.Nodes,
|
|
Relationships: result.Relationships,
|
|
}
|
|
|
|
return result.Segments, graph, nil
|
|
}
|
|
|
|
// ========== Expand Search (Graph + Vector) ==========
|
|
|
|
// searchExpand performs graph-based entity expansion + vector search
|
|
// This mode uses graph to find related entities, then enhances vector search
|
|
func (kb *KBInstance) searchExpand(ctx context.Context, collectionID string, queryText string, query Query, embedding graphragtypes.Embedding, collection map[string]interface{}) ([]graphragtypes.Segment, *GraphData, error) {
|
|
// Step 1: Extract entities from query using graph search
|
|
extraction, err := kb.createExtraction(collection)
|
|
if err != nil {
|
|
// Fall back to pure vector search if extraction is not available
|
|
log.Warn("Extraction not available, falling back to vector search: %v", err)
|
|
return kb.searchVector(ctx, collectionID, queryText, query, embedding)
|
|
}
|
|
|
|
maxDepth := query.MaxDepth
|
|
if maxDepth <= 0 {
|
|
maxDepth = DefaultMaxDepth
|
|
}
|
|
|
|
graphOptions := &graphragtypes.GraphSearchOptions{
|
|
CollectionID: collectionID,
|
|
DocumentID: query.DocumentID,
|
|
Query: queryText,
|
|
MaxDepth: maxDepth,
|
|
Extraction: extraction,
|
|
}
|
|
|
|
// Execute graph search to find related entities
|
|
graphResult, graphErr := kb.GraphRag.SearchGraph(ctx, graphOptions)
|
|
|
|
// Step 2: Perform vector search
|
|
k := query.PageSize
|
|
if k <= 0 {
|
|
k = DefaultSearchK
|
|
}
|
|
if k > MaxSearchK {
|
|
k = MaxSearchK
|
|
}
|
|
|
|
vectorOptions := &graphragtypes.VectorSearchOptions{
|
|
CollectionID: collectionID,
|
|
DocumentID: query.DocumentID,
|
|
Query: queryText,
|
|
K: k,
|
|
MinScore: query.MinScore,
|
|
Embedding: embedding,
|
|
}
|
|
|
|
if len(query.Metadata) > 0 {
|
|
vectorOptions.Filter = query.Metadata
|
|
}
|
|
|
|
vectorResult, err := kb.GraphRag.SearchVector(ctx, vectorOptions)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("vector search failed: %w", err)
|
|
}
|
|
|
|
// Step 3: Merge results
|
|
segments := vectorResult.Segments
|
|
|
|
var graph *GraphData
|
|
if graphErr == nil && graphResult != nil {
|
|
// Add graph segments (deduplicated later)
|
|
segments = append(segments, graphResult.Segments...)
|
|
|
|
// Include graph data
|
|
graph = &GraphData{
|
|
Nodes: graphResult.Nodes,
|
|
Relationships: graphResult.Relationships,
|
|
}
|
|
}
|
|
|
|
return segments, graph, nil
|
|
}
|
|
|
|
// ========== Helper Functions ==========
|
|
|
|
// createEmbedding creates an embedding function from provider config
|
|
func (kb *KBInstance) createEmbedding(providerID, optionID, locale string) (graphragtypes.Embedding, error) {
|
|
if locale == "" {
|
|
locale = "en"
|
|
}
|
|
|
|
// Get provider option
|
|
option, err := kb.getProviderOption("embedding", providerID, optionID, locale)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get embedding option: %w", err)
|
|
}
|
|
|
|
// Create embedding provider
|
|
return factory.MakeEmbedding(providerID, option)
|
|
}
|
|
|
|
// createExtraction creates an extraction function from collection config
|
|
func (kb *KBInstance) createExtraction(collection map[string]interface{}) (graphragtypes.Extraction, error) {
|
|
// Try to get extraction provider from collection metadata
|
|
metadata, _ := collection["metadata"].(map[string]interface{})
|
|
if metadata == nil {
|
|
metadata = collection
|
|
}
|
|
|
|
extractionProviderID, _ := metadata["__extraction_provider"].(string)
|
|
extractionOptionID, _ := metadata["__extraction_option"].(string)
|
|
|
|
// Fall back to default extraction provider
|
|
if extractionProviderID == "" {
|
|
extractionProviderID = "__yao.openai"
|
|
extractionOptionID = "gpt-4o-mini"
|
|
}
|
|
|
|
// Get provider option
|
|
option, err := kb.getProviderOption("extraction", extractionProviderID, extractionOptionID, "en")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get extraction option: %w", err)
|
|
}
|
|
|
|
// Create extraction provider
|
|
return factory.MakeExtraction(extractionProviderID, option)
|
|
}
|
|
|
|
// deduplicateSegments removes duplicate segments by ID, keeping highest score
|
|
func (kb *KBInstance) deduplicateSegments(segments []graphragtypes.Segment) []graphragtypes.Segment {
|
|
seen := make(map[string]int) // ID -> index in result
|
|
result := make([]graphragtypes.Segment, 0, len(segments))
|
|
|
|
for _, seg := range segments {
|
|
if idx, exists := seen[seg.ID]; exists {
|
|
// Keep the one with higher score
|
|
if seg.Score > result[idx].Score {
|
|
result[idx] = seg
|
|
}
|
|
} else {
|
|
seen[seg.ID] = len(result)
|
|
result = append(result, seg)
|
|
}
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
// mergeGraphData merges two GraphData objects
|
|
func mergeGraphData(a, b *GraphData) *GraphData {
|
|
if a == nil {
|
|
return b
|
|
}
|
|
if b == nil {
|
|
return a
|
|
}
|
|
|
|
// Merge nodes (deduplicate by ID)
|
|
nodeMap := make(map[string]graphragtypes.GraphNode)
|
|
for _, n := range a.Nodes {
|
|
nodeMap[n.ID] = n
|
|
}
|
|
for _, n := range b.Nodes {
|
|
nodeMap[n.ID] = n
|
|
}
|
|
|
|
nodes := make([]graphragtypes.GraphNode, 0, len(nodeMap))
|
|
for _, n := range nodeMap {
|
|
nodes = append(nodes, n)
|
|
}
|
|
|
|
// Merge relationships (deduplicate by ID)
|
|
relMap := make(map[string]graphragtypes.GraphRelationship)
|
|
for _, r := range a.Relationships {
|
|
relMap[r.ID] = r
|
|
}
|
|
for _, r := range b.Relationships {
|
|
relMap[r.ID] = r
|
|
}
|
|
|
|
relationships := make([]graphragtypes.GraphRelationship, 0, len(relMap))
|
|
for _, r := range relMap {
|
|
relationships = append(relationships, r)
|
|
}
|
|
|
|
return &GraphData{
|
|
Nodes: nodes,
|
|
Relationships: relationships,
|
|
}
|
|
}
|
|
|
|
// applyPagination applies pagination to segments
|
|
func (kb *KBInstance) applyPagination(segments []graphragtypes.Segment, query Query) *SearchResult {
|
|
total := len(segments)
|
|
|
|
// If no pagination requested, return all
|
|
if query.Page <= 0 && query.PageSize <= 0 {
|
|
return &SearchResult{
|
|
Segments: segments,
|
|
Total: total,
|
|
}
|
|
}
|
|
|
|
page := query.Page
|
|
if page <= 0 {
|
|
page = 1
|
|
}
|
|
|
|
pageSize := query.PageSize
|
|
if pageSize <= 0 {
|
|
pageSize = DefaultSearchPageSize
|
|
}
|
|
|
|
// Calculate pagination
|
|
totalPages := (total + pageSize - 1) / pageSize
|
|
start := (page - 1) * pageSize
|
|
end := start + pageSize
|
|
|
|
if start >= total {
|
|
return &SearchResult{
|
|
Segments: []graphragtypes.Segment{},
|
|
Total: total,
|
|
Page: page,
|
|
PageSize: pageSize,
|
|
TotalPages: totalPages,
|
|
}
|
|
}
|
|
|
|
if end > total {
|
|
end = total
|
|
}
|
|
|
|
result := &SearchResult{
|
|
Segments: segments[start:end],
|
|
Total: total,
|
|
Page: page,
|
|
PageSize: pageSize,
|
|
TotalPages: totalPages,
|
|
}
|
|
|
|
// Set next/prev page
|
|
if page < totalPages {
|
|
result.Next = page + 1
|
|
}
|
|
if page > 1 {
|
|
result.Prev = page - 1
|
|
}
|
|
|
|
return result
|
|
}
|