github.com/weaviate/weaviate@v1.24.6/modules/text2vec-contextionary/additional/nearestneighbors/extender.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package nearestneighbors
    13  
    14  import (
    15  	"context"
    16  	"fmt"
    17  
    18  	"github.com/weaviate/weaviate/entities/moduletools"
    19  
    20  	"github.com/pkg/errors"
    21  	"github.com/tailor-inc/graphql/language/ast"
    22  	"github.com/weaviate/weaviate/entities/models"
    23  	"github.com/weaviate/weaviate/entities/search"
    24  	txt2vecmodels "github.com/weaviate/weaviate/modules/text2vec-contextionary/additional/models"
    25  )
    26  
    27  const (
    28  	DefaultLimit = 10
    29  	DefaultK     = 32
    30  )
    31  
    32  type Extender struct {
    33  	searcher contextionary
    34  }
    35  
    36  type contextionary interface {
    37  	MultiNearestWordsByVector(ctx context.Context, vectors [][]float32, k, n int) ([]*txt2vecmodels.NearestNeighbors, error)
    38  }
    39  
    40  func (e *Extender) AdditionalPropertyDefaultValue() interface{} {
    41  	return true
    42  }
    43  
    44  func (e *Extender) AdditionalPropertyFn(ctx context.Context,
    45  	in []search.Result, params interface{}, limit *int,
    46  	argumentModuleParams map[string]interface{}, cfg moduletools.ClassConfig,
    47  ) ([]search.Result, error) {
    48  	return e.Multi(ctx, in, limit)
    49  }
    50  
    51  func (e *Extender) ExtractAdditionalFn(param []*ast.Argument) interface{} {
    52  	return true
    53  }
    54  
    55  func (e *Extender) Single(ctx context.Context, in *search.Result, limit *int) (*search.Result, error) {
    56  	if in == nil {
    57  		return nil, nil
    58  	}
    59  
    60  	multiRes, err := e.Multi(ctx, []search.Result{*in}, limit) // safe to deref, as we did a nil check before
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  
    65  	return &multiRes[0], nil
    66  }
    67  
    68  func (e *Extender) Multi(ctx context.Context, in []search.Result, limit *int) ([]search.Result, error) {
    69  	if in == nil {
    70  		return nil, nil
    71  	}
    72  
    73  	vectors := make([][]float32, len(in))
    74  	for i, res := range in {
    75  		if res.Vector == nil || len(res.Vector) == 0 {
    76  			return nil, fmt.Errorf("item %d has no vector", i)
    77  		}
    78  		vectors[i] = res.Vector
    79  	}
    80  
    81  	neighbors, err := e.searcher.MultiNearestWordsByVector(ctx, vectors, DefaultK, limitOrDefault(limit))
    82  	if err != nil {
    83  		return nil, errors.Wrap(err, "get neighbors for search results")
    84  	}
    85  
    86  	if len(neighbors) != len(in) {
    87  		return nil, fmt.Errorf("inconsistent results: input=%d neighbors=%d", len(in), len(neighbors))
    88  	}
    89  
    90  	for i, res := range in {
    91  		up := res.AdditionalProperties
    92  		if up == nil {
    93  			up = models.AdditionalProperties{}
    94  		}
    95  
    96  		up["nearestNeighbors"] = removeDollarElements(neighbors[i])
    97  		in[i].AdditionalProperties = up
    98  	}
    99  
   100  	return in, nil
   101  }
   102  
   103  func NewExtender(searcher contextionary) *Extender {
   104  	return &Extender{searcher: searcher}
   105  }
   106  
   107  func limitOrDefault(user *int) int {
   108  	if user == nil || *user == 0 {
   109  		return DefaultLimit
   110  	}
   111  
   112  	return *user
   113  }
   114  
   115  func removeDollarElements(in *txt2vecmodels.NearestNeighbors) *txt2vecmodels.NearestNeighbors {
   116  	neighbors := make([]*txt2vecmodels.NearestNeighbor, len(in.Neighbors))
   117  	i := 0
   118  	for _, elem := range in.Neighbors {
   119  		if elem.Concept[0] == '$' {
   120  			continue
   121  		}
   122  
   123  		neighbors[i] = elem
   124  		i++
   125  	}
   126  
   127  	return &txt2vecmodels.NearestNeighbors{
   128  		Neighbors: neighbors[:i],
   129  	}
   130  }