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}
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
}