vectorset_commands.go

Documentation: github.com/redis/go-redis/v9

     1  package redis
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"strconv"
     7  )
     8  
     9  // note: the APIs is experimental and may be subject to change.
    10  type VectorSetCmdable interface {
    11  	VAdd(ctx context.Context, key, element string, val Vector) *BoolCmd
    12  	VAddWithArgs(ctx context.Context, key, element string, val Vector, addArgs *VAddArgs) *BoolCmd
    13  	VCard(ctx context.Context, key string) *IntCmd
    14  	VDim(ctx context.Context, key string) *IntCmd
    15  	VEmb(ctx context.Context, key, element string, raw bool) *SliceCmd
    16  	VGetAttr(ctx context.Context, key, element string) *StringCmd
    17  	VInfo(ctx context.Context, key string) *MapStringInterfaceCmd
    18  	VLinks(ctx context.Context, key, element string) *StringSliceCmd
    19  	VLinksWithScores(ctx context.Context, key, element string) *VectorScoreSliceCmd
    20  	VRandMember(ctx context.Context, key string) *StringCmd
    21  	VRandMemberCount(ctx context.Context, key string, count int) *StringSliceCmd
    22  	VRem(ctx context.Context, key, element string) *BoolCmd
    23  	VSetAttr(ctx context.Context, key, element string, attr interface{}) *BoolCmd
    24  	VClearAttributes(ctx context.Context, key, element string) *BoolCmd
    25  	VSim(ctx context.Context, key string, val Vector) *StringSliceCmd
    26  	VSimWithScores(ctx context.Context, key string, val Vector) *VectorScoreSliceCmd
    27  	VSimWithArgs(ctx context.Context, key string, val Vector, args *VSimArgs) *StringSliceCmd
    28  	VSimWithArgsWithScores(ctx context.Context, key string, val Vector, args *VSimArgs) *VectorScoreSliceCmd
    29  }
    30  
    31  type Vector interface {
    32  	Value() []any
    33  }
    34  
    35  const (
    36  	vectorFormatFP32   string = "FP32"
    37  	vectorFormatValues string = "Values"
    38  )
    39  
    40  type VectorFP32 struct {
    41  	Val []byte
    42  }
    43  
    44  func (v *VectorFP32) Value() []any {
    45  	return []any{vectorFormatFP32, v.Val}
    46  }
    47  
    48  var _ Vector = (*VectorFP32)(nil)
    49  
    50  type VectorValues struct {
    51  	Val []float64
    52  }
    53  
    54  func (v *VectorValues) Value() []any {
    55  	res := make([]any, 2+len(v.Val))
    56  	res[0] = vectorFormatValues
    57  	res[1] = len(v.Val)
    58  	for i, v := range v.Val {
    59  		res[2+i] = v
    60  	}
    61  	return res
    62  }
    63  
    64  var _ Vector = (*VectorValues)(nil)
    65  
    66  type VectorRef struct {
    67  	Name string // the name of the referent vector
    68  }
    69  
    70  func (v *VectorRef) Value() []any {
    71  	return []any{"ele", v.Name}
    72  }
    73  
    74  var _ Vector = (*VectorRef)(nil)
    75  
    76  type VectorScore struct {
    77  	Name  string
    78  	Score float64
    79  }
    80  
    81  // `VADD key (FP32 | VALUES num) vector element`
    82  // note: the API is experimental and may be subject to change.
    83  func (c cmdable) VAdd(ctx context.Context, key, element string, val Vector) *BoolCmd {
    84  	return c.VAddWithArgs(ctx, key, element, val, &VAddArgs{})
    85  }
    86  
    87  type VAddArgs struct {
    88  	// the REDUCE option must be passed immediately after the key
    89  	Reduce int64
    90  	Cas    bool
    91  
    92  	// The NoQuant, Q8 and Bin options are mutually exclusive.
    93  	NoQuant bool
    94  	Q8      bool
    95  	Bin     bool
    96  
    97  	EF      int64
    98  	SetAttr string
    99  	M       int64
   100  }
   101  
   102  func (v VAddArgs) reduce() int64 {
   103  	return v.Reduce
   104  }
   105  
   106  func (v VAddArgs) appendArgs(args []any) []any {
   107  	if v.Cas {
   108  		args = append(args, "cas")
   109  	}
   110  
   111  	if v.NoQuant {
   112  		args = append(args, "noquant")
   113  	} else if v.Q8 {
   114  		args = append(args, "q8")
   115  	} else if v.Bin {
   116  		args = append(args, "bin")
   117  	}
   118  
   119  	if v.EF > 0 {
   120  		args = append(args, "ef", strconv.FormatInt(v.EF, 10))
   121  	}
   122  	if len(v.SetAttr) > 0 {
   123  		args = append(args, "setattr", v.SetAttr)
   124  	}
   125  	if v.M > 0 {
   126  		args = append(args, "m", strconv.FormatInt(v.M, 10))
   127  	}
   128  	return args
   129  }
   130  
   131  // `VADD key [REDUCE dim] (FP32 | VALUES num) vector element [CAS] [NOQUANT | Q8 | BIN] [EF build-exploration-factor] [SETATTR attributes] [M numlinks]`
   132  // note: the API is experimental and may be subject to change.
   133  func (c cmdable) VAddWithArgs(ctx context.Context, key, element string, val Vector, addArgs *VAddArgs) *BoolCmd {
   134  	if addArgs == nil {
   135  		addArgs = &VAddArgs{}
   136  	}
   137  	args := []any{"vadd", key}
   138  	if addArgs.reduce() > 0 {
   139  		args = append(args, "reduce", addArgs.reduce())
   140  	}
   141  	args = append(args, val.Value()...)
   142  	args = append(args, element)
   143  	args = addArgs.appendArgs(args)
   144  	cmd := NewBoolCmd(ctx, args...)
   145  	_ = c(ctx, cmd)
   146  	return cmd
   147  }
   148  
   149  // `VCARD key`
   150  // note: the API is experimental and may be subject to change.
   151  func (c cmdable) VCard(ctx context.Context, key string) *IntCmd {
   152  	cmd := NewIntCmd(ctx, "vcard", key)
   153  	_ = c(ctx, cmd)
   154  	return cmd
   155  }
   156  
   157  // `VDIM key`
   158  // note: the API is experimental and may be subject to change.
   159  func (c cmdable) VDim(ctx context.Context, key string) *IntCmd {
   160  	cmd := NewIntCmd(ctx, "vdim", key)
   161  	_ = c(ctx, cmd)
   162  	return cmd
   163  }
   164  
   165  // `VEMB key element [RAW]`
   166  // note: the API is experimental and may be subject to change.
   167  func (c cmdable) VEmb(ctx context.Context, key, element string, raw bool) *SliceCmd {
   168  	args := []any{"vemb", key, element}
   169  	if raw {
   170  		args = append(args, "raw")
   171  	}
   172  	cmd := NewSliceCmd(ctx, args...)
   173  	_ = c(ctx, cmd)
   174  	return cmd
   175  }
   176  
   177  // `VGETATTR key element`
   178  // note: the API is experimental and may be subject to change.
   179  func (c cmdable) VGetAttr(ctx context.Context, key, element string) *StringCmd {
   180  	cmd := NewStringCmd(ctx, "vgetattr", key, element)
   181  	_ = c(ctx, cmd)
   182  	return cmd
   183  }
   184  
   185  // `VINFO key`
   186  // note: the API is experimental and may be subject to change.
   187  func (c cmdable) VInfo(ctx context.Context, key string) *MapStringInterfaceCmd {
   188  	cmd := NewMapStringInterfaceCmd(ctx, "vinfo", key)
   189  	_ = c(ctx, cmd)
   190  	return cmd
   191  }
   192  
   193  // `VLINKS key element`
   194  // note: the API is experimental and may be subject to change.
   195  func (c cmdable) VLinks(ctx context.Context, key, element string) *StringSliceCmd {
   196  	cmd := NewStringSliceCmd(ctx, "vlinks", key, element)
   197  	_ = c(ctx, cmd)
   198  	return cmd
   199  }
   200  
   201  // `VLINKS key element WITHSCORES`
   202  // note: the API is experimental and may be subject to change.
   203  func (c cmdable) VLinksWithScores(ctx context.Context, key, element string) *VectorScoreSliceCmd {
   204  	cmd := NewVectorInfoSliceCmd(ctx, "vlinks", key, element, "withscores")
   205  	_ = c(ctx, cmd)
   206  	return cmd
   207  }
   208  
   209  // `VRANDMEMBER key`
   210  // note: the API is experimental and may be subject to change.
   211  func (c cmdable) VRandMember(ctx context.Context, key string) *StringCmd {
   212  	cmd := NewStringCmd(ctx, "vrandmember", key)
   213  	_ = c(ctx, cmd)
   214  	return cmd
   215  }
   216  
   217  // `VRANDMEMBER key [count]`
   218  // note: the API is experimental and may be subject to change.
   219  func (c cmdable) VRandMemberCount(ctx context.Context, key string, count int) *StringSliceCmd {
   220  	cmd := NewStringSliceCmd(ctx, "vrandmember", key, count)
   221  	_ = c(ctx, cmd)
   222  	return cmd
   223  }
   224  
   225  // `VREM key element`
   226  // note: the API is experimental and may be subject to change.
   227  func (c cmdable) VRem(ctx context.Context, key, element string) *BoolCmd {
   228  	cmd := NewBoolCmd(ctx, "vrem", key, element)
   229  	_ = c(ctx, cmd)
   230  	return cmd
   231  }
   232  
   233  // `VSETATTR key element "{ JSON obj }"`
   234  // The `attr` must be something that can be marshaled to JSON (using encoding/JSON) unless
   235  // the argument is a string or []byte when we assume that it can be passed directly as JSON.
   236  //
   237  // note: the API is experimental and may be subject to change.
   238  func (c cmdable) VSetAttr(ctx context.Context, key, element string, attr interface{}) *BoolCmd {
   239  	var attrStr string
   240  	var err error
   241  	switch v := attr.(type) {
   242  	case string:
   243  		attrStr = v
   244  	case []byte:
   245  		attrStr = string(v)
   246  	default:
   247  		var bytes []byte
   248  		bytes, err = json.Marshal(v)
   249  		if err != nil {
   250  			// If marshalling fails, create the command and set the error; this command won't be executed.
   251  			cmd := NewBoolCmd(ctx, "vsetattr", key, element, "")
   252  			cmd.SetErr(err)
   253  			return cmd
   254  		}
   255  		attrStr = string(bytes)
   256  	}
   257  	cmd := NewBoolCmd(ctx, "vsetattr", key, element, attrStr)
   258  	_ = c(ctx, cmd)
   259  	return cmd
   260  }
   261  
   262  // `VClearAttributes` clear attributes on a vector set element.
   263  // The implementation of `VClearAttributes` is execute command `VSETATTR key element ""`.
   264  // note: the API is experimental and may be subject to change.
   265  func (c cmdable) VClearAttributes(ctx context.Context, key, element string) *BoolCmd {
   266  	cmd := NewBoolCmd(ctx, "vsetattr", key, element, "")
   267  	_ = c(ctx, cmd)
   268  	return cmd
   269  }
   270  
   271  // `VSIM key (ELE | FP32 | VALUES num) (vector | element)`
   272  // note: the API is experimental and may be subject to change.
   273  func (c cmdable) VSim(ctx context.Context, key string, val Vector) *StringSliceCmd {
   274  	return c.VSimWithArgs(ctx, key, val, &VSimArgs{})
   275  }
   276  
   277  // `VSIM key (ELE | FP32 | VALUES num) (vector | element) WITHSCORES`
   278  // note: the API is experimental and may be subject to change.
   279  func (c cmdable) VSimWithScores(ctx context.Context, key string, val Vector) *VectorScoreSliceCmd {
   280  	return c.VSimWithArgsWithScores(ctx, key, val, &VSimArgs{})
   281  }
   282  
   283  type VSimArgs struct {
   284  	Count    int64
   285  	EF       int64
   286  	Filter   string
   287  	FilterEF int64
   288  	Truth    bool
   289  	NoThread bool
   290  	Epsilon  float64
   291  }
   292  
   293  func (v VSimArgs) appendArgs(args []any) []any {
   294  	if v.Count > 0 {
   295  		args = append(args, "count", v.Count)
   296  	}
   297  	if v.EF > 0 {
   298  		args = append(args, "ef", v.EF)
   299  	}
   300  	if len(v.Filter) > 0 {
   301  		args = append(args, "filter", v.Filter)
   302  	}
   303  	if v.FilterEF > 0 {
   304  		args = append(args, "filter-ef", v.FilterEF)
   305  	}
   306  	if v.Truth {
   307  		args = append(args, "truth")
   308  	}
   309  	if v.NoThread {
   310  		args = append(args, "nothread")
   311  	}
   312  	if v.Epsilon > 0 {
   313  		args = append(args, "Epsilon", v.Epsilon)
   314  	}
   315  	return args
   316  }
   317  
   318  // `VSIM key (ELE | FP32 | VALUES num) (vector | element) [COUNT num] [EPSILON delta]
   319  // [EF search-exploration-factor] [FILTER expression] [FILTER-EF max-filtering-effort] [TRUTH] [NOTHREAD]`
   320  // note: the API is experimental and may be subject to change.
   321  func (c cmdable) VSimWithArgs(ctx context.Context, key string, val Vector, simArgs *VSimArgs) *StringSliceCmd {
   322  	if simArgs == nil {
   323  		simArgs = &VSimArgs{}
   324  	}
   325  	args := []any{"vsim", key}
   326  	args = append(args, val.Value()...)
   327  	args = simArgs.appendArgs(args)
   328  	cmd := NewStringSliceCmd(ctx, args...)
   329  	_ = c(ctx, cmd)
   330  	return cmd
   331  }
   332  
   333  // `VSIM key (ELE | FP32 | VALUES num) (vector | element) [WITHSCORES] [COUNT num] [EPSILON delta]
   334  // [EF search-exploration-factor] [FILTER expression] [FILTER-EF max-filtering-effort] [TRUTH] [NOTHREAD]`
   335  // note: the API is experimental and may be subject to change.
   336  func (c cmdable) VSimWithArgsWithScores(ctx context.Context, key string, val Vector, simArgs *VSimArgs) *VectorScoreSliceCmd {
   337  	if simArgs == nil {
   338  		simArgs = &VSimArgs{}
   339  	}
   340  	args := []any{"vsim", key}
   341  	args = append(args, val.Value()...)
   342  	args = append(args, "withscores")
   343  	args = simArgs.appendArgs(args)
   344  	cmd := NewVectorInfoSliceCmd(ctx, args...)
   345  	_ = c(ctx, cmd)
   346  	return cmd
   347  }
   348
View as plain text