feat(stores): Vector store backend (#1795)

Add simple vector store backend

Signed-off-by: Richard Palethorpe <io@richiejp.com>

pkg/store/client.go

@@ -0,0 +1,155 @@
+package store
+
+import (
+	"context"
+	"fmt"
+
+	grpc "github.com/go-skynet/LocalAI/pkg/grpc"
+	"github.com/go-skynet/LocalAI/pkg/grpc/proto"
+)
+
+// Wrapper for the GRPC client so that simple use cases are handled without verbosity
+
+// SetCols sets multiple key-value pairs in the store
+// It's in columnar format so that keys[i] is associated with values[i]
+func SetCols(ctx context.Context, c grpc.Backend, keys [][]float32, values [][]byte) error {
+	protoKeys := make([]*proto.StoresKey, len(keys))
+	for i, k := range keys {
+		protoKeys[i] = &proto.StoresKey{
+			Floats: k,
+		}
+	}
+	protoValues := make([]*proto.StoresValue, len(values))
+	for i, v := range values {
+		protoValues[i] = &proto.StoresValue{
+			Bytes: v,
+		}
+	}
+	setOpts := &proto.StoresSetOptions{
+		Keys:   protoKeys,
+		Values: protoValues,
+	}
+
+	res, err := c.StoresSet(ctx, setOpts)
+	if err != nil {
+		return err
+	}
+
+	if res.Success {
+		return nil
+	}
+
+	return fmt.Errorf("failed to set keys: %v", res.Message)
+}
+
+// SetSingle sets a single key-value pair in the store
+// Don't call this in a tight loop, instead use SetCols
+func SetSingle(ctx context.Context, c grpc.Backend, key []float32, value []byte) error {
+	return SetCols(ctx, c, [][]float32{key}, [][]byte{value})
+}
+
+// DeleteCols deletes multiple key-value pairs from the store
+// It's in columnar format so that keys[i] is associated with values[i]
+func DeleteCols(ctx context.Context, c grpc.Backend, keys [][]float32) error {
+	protoKeys := make([]*proto.StoresKey, len(keys))
+	for i, k := range keys {
+		protoKeys[i] = &proto.StoresKey{
+			Floats: k,
+		}
+	}
+	deleteOpts := &proto.StoresDeleteOptions{
+		Keys: protoKeys,
+	}
+
+	res, err := c.StoresDelete(ctx, deleteOpts)
+	if err != nil {
+		return err
+	}
+
+	if res.Success {
+		return nil
+	}
+
+	return fmt.Errorf("failed to delete keys: %v", res.Message)
+}
+
+// DeleteSingle deletes a single key-value pair from the store
+// Don't call this in a tight loop, instead use DeleteCols
+func DeleteSingle(ctx context.Context, c grpc.Backend, key []float32) error {
+	return DeleteCols(ctx, c, [][]float32{key})
+}
+
+// GetCols gets multiple key-value pairs from the store
+// It's in columnar format so that keys[i] is associated with values[i]
+// Be warned the keys are sorted and will be returned in a different order than they were input
+// There is no guarantee as to how the keys are sorted
+func GetCols(ctx context.Context, c grpc.Backend, keys [][]float32) ([][]float32, [][]byte, error) {
+	protoKeys := make([]*proto.StoresKey, len(keys))
+	for i, k := range keys {
+		protoKeys[i] = &proto.StoresKey{
+			Floats: k,
+		}
+	}
+	getOpts := &proto.StoresGetOptions{
+		Keys: protoKeys,
+	}
+
+	res, err := c.StoresGet(ctx, getOpts)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	ks := make([][]float32, len(res.Keys))
+	for i, k := range res.Keys {
+		ks[i] = k.Floats
+	}
+	vs := make([][]byte, len(res.Values))
+	for i, v := range res.Values {
+		vs[i] = v.Bytes
+	}
+
+	return ks, vs, nil
+}
+
+// GetSingle gets a single key-value pair from the store
+// Don't call this in a tight loop, instead use GetCols
+func GetSingle(ctx context.Context, c grpc.Backend, key []float32) ([]byte, error) {
+	_, values, err := GetCols(ctx, c, [][]float32{key})
+	if err != nil {
+		return nil, err
+	}
+
+	if len(values) > 0 {
+		return values[0], nil
+	}
+
+	return nil, fmt.Errorf("failed to get key")
+}
+
+// Find similar keys to the given key. Returns the keys, values, and similarities
+func Find(ctx context.Context, c grpc.Backend, key []float32, topk int) ([][]float32, [][]byte, []float32, error) {
+	findOpts := &proto.StoresFindOptions{
+		Key: &proto.StoresKey{
+			Floats: key,
+		},
+		TopK: int32(topk),
+	}
+
+	res, err := c.StoresFind(ctx, findOpts)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	ks := make([][]float32, len(res.Keys))
+	vs := make([][]byte, len(res.Values))
+
+	for i, k := range res.Keys {
+		ks[i] = k.Floats
+	}
+
+	for i, v := range res.Values {
+		vs[i] = v.Bytes
+	}
+
+	return ks, vs, res.Similarities, nil
+}