230 lines · 6.0 KB · Go
RawDownload
package vector import (	"context"	"errors"	"fmt"	"math"	"strings" 	crawlvector "github.com/openclaw/crawlkit/vector") type Item struct {	ThreadID int64	Vector   []float64} type Neighbor struct {	ThreadID int64   `json:"thread_id"`	Score    float64 `json:"score"`} type QueryOptions struct {	Backend         string	Limit           int	ExcludeThreadID int64	TurboVec        crawlvector.TurboVecOptions} const maxTurboVecTieCandidates = 16 * 1024 func Query(items []Item, query []float64, limit int, excludeThreadID int64) []Neighbor {	neighbors, _ := queryExact(context.Background(), items, query, limit, excludeThreadID)	return neighbors} func QueryWithOptions(ctx context.Context, items []Item, query []float64, opts QueryOptions) ([]Neighbor, error) {	backend := crawlvector.SearchBackend(strings.ToLower(strings.TrimSpace(opts.Backend)))	if backend == "" {		backend = crawlvector.BackendExact	}	switch backend {	case crawlvector.BackendExact:		return queryExact(ctx, items, query, opts.Limit, opts.ExcludeThreadID)	case crawlvector.BackendTurboVec:		return queryTurboVec(ctx, items, query, opts)	default:		return nil, fmt.Errorf("unsupported vector backend %q", opts.Backend)	}} func queryExact(ctx context.Context, items []Item, query []float64, limit int, excludeThreadID int64) ([]Neighbor, error) {	if limit <= 0 {		limit = 20	}	if err := validateExactQuery(query); err != nil {		return nil, err	}	scored := make([]crawlvector.Scored[Neighbor], 0, len(items))	for _, item := range items {		if err := ctx.Err(); err != nil {			return nil, err		}		if item.ThreadID == excludeThreadID {			continue		}		score := Cosine(query, item.Vector)		if math.IsNaN(score) || math.IsInf(score, 0) || score <= 0 {			continue		}		neighbor := Neighbor{ThreadID: item.ThreadID, Score: score}		scored = append(scored, crawlvector.Scored[Neighbor]{Item: neighbor, Score: score})	}	if err := ctx.Err(); err != nil {		return nil, err	}	top := crawlvector.TopK(scored, limit, func(left, right Neighbor) bool {		return left.ThreadID < right.ThreadID	})	out := make([]Neighbor, len(top))	for i, item := range top {		out[i] = item.Item	}	return out, nil} func queryTurboVec(ctx context.Context, items []Item, query []float64, opts QueryOptions) ([]Neighbor, error) {	if opts.Limit <= 0 {		opts.Limit = 20	}	if err := validateExactQuery(query); err != nil {		return nil, err	}	candidateCount := 0	for _, item := range items {		if item.ThreadID == opts.ExcludeThreadID {			continue		}		candidateCount++	}	if shouldFallbackToExact(len(query), candidateCount, opts.TurboVec.MaxInputBytes) {		return queryExact(ctx, items, query, opts.Limit, opts.ExcludeThreadID)	}	candidates := make([]crawlvector.SearchCandidate[Neighbor], 0, len(items))	for _, item := range items {		if err := ctx.Err(); err != nil {			return nil, err		}		if item.ThreadID == opts.ExcludeThreadID {			continue		}		candidates = append(candidates, crawlvector.SearchCandidate[Neighbor]{			Item:   Neighbor{ThreadID: item.ThreadID},			Vector: float64To32(item.Vector),		})	}	if len(candidates) == 0 {		return nil, nil	}	results, err := crawlvector.Search(ctx, float64To32(query), candidates, crawlvector.SearchOptions[Neighbor]{		Backend:       crawlvector.BackendTurboVec,		Limit:         opts.Limit,		InvalidVector: crawlvector.InvalidVectorSkip,		TurboVec:      opts.TurboVec,		TieLess: func(left, right Neighbor) bool {			return left.ThreadID < right.ThreadID		},	})	if err != nil {		return nil, err	}	out := make([]Neighbor, 0, len(results))	for _, result := range results {		if math.IsNaN(result.Score) || math.IsInf(result.Score, 0) || result.Score <= 0 {			continue		}		neighbor := result.Item		neighbor.Score = result.Score		out = append(out, neighbor)	}	return out, nil} func float64To32(values []float64) []float32 {	out := make([]float32, len(values))	for i, value := range values {		out[i] = float32(value)	}	return out} func validateExactQuery(query []float64) error {	if len(query) == 0 {		return errors.New("query vector is empty")	}	var maxAbs float64	for i, value := range query {		if math.IsNaN(value) || math.IsInf(value, 0) {			return fmt.Errorf("query vector contains non-finite value at index %d", i)		}		maxAbs = max(maxAbs, math.Abs(value))	}	if maxAbs == 0 {		return errors.New("query vector is zero")	}	var magnitude float64	for _, value := range query {		scaled := value / maxAbs		magnitude += scaled * scaled	}	if magnitude == 0 {		return errors.New("query vector is zero")	}	return nil} func shouldFallbackToExact(dimensions, candidateCount int, maxInputBytes int64) bool {	if candidateCount > maxTurboVecTieCandidates {		return true	}	effectiveMaxInput := maxInputBytes	if effectiveMaxInput == 0 {		effectiveMaxInput = crawlvector.DefaultTurboVecMaxInputSize	}	if effectiveMaxInput <= 0 || dimensions <= 0 || candidateCount <= 0 {		return false	}	const baseJSONBudget = int64(256)	const floatJSONBudget = int64(16)	estimate := baseJSONBudget + int64(dimensions)*int64(candidateCount+1)*floatJSONBudget + int64(candidateCount*4)	return estimate > effectiveMaxInput} func Cosine(left, right []float64) float64 {	if len(left) == 0 || len(left) != len(right) {		return 0	}	var leftMaxAbs, rightMaxAbs float64	for index := range left {		leftValue := left[index]		rightValue := right[index]		if math.IsNaN(leftValue) || math.IsNaN(rightValue) || math.IsInf(leftValue, 0) || math.IsInf(rightValue, 0) {			return 0		}		leftMaxAbs = max(leftMaxAbs, math.Abs(leftValue))		rightMaxAbs = max(rightMaxAbs, math.Abs(rightValue))	}	if leftMaxAbs == 0 || rightMaxAbs == 0 {		return 0	}	var dot, leftMag, rightMag float64	for index := range left {		leftValue := left[index] / leftMaxAbs		rightValue := right[index] / rightMaxAbs		dot += leftValue * rightValue		leftMag += leftValue * leftValue		rightMag += rightValue * rightValue	}	if leftMag == 0 || rightMag == 0 {		return 0	}	score := dot / (math.Sqrt(leftMag) * math.Sqrt(rightMag))	if score > 1 {		return 1	}	if score < -1 {		return -1	}	return score}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t