b63f3fe246
- Enhanced the KB search handler to utilize the KB API for executing search queries, improving search accuracy and performance. - Implemented authorization checks for collections in the search requests, ensuring only accessible collections are queried. - Updated the search request structure to include metadata filtering capabilities, allowing for more refined search results. - Refactored unit tests to validate new search functionalities, including threshold handling and collection initialization checks, ensuring robust test coverage. - Adjusted the Makefile to streamline test coverage reporting and updated GitHub Actions workflows to include Codecov integration for better visibility on test coverage metrics.
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
|
|
DefaultThreshold = 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.Threshold,
|
|
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.Threshold,
|
|
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
|
|
}
|