wehub-resource-sync
104 lines
3.4 KiB
Go
104 lines
3.4 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
|
"github.com/mudler/LocalAI/pkg/grpc/proto"
|
|
)
|
|
|
|
// Wrapper for the GRPC client so that simple use cases are handled without verbosity
|
|
|
|
// NamespacePrefix marks a load request as a store namespace rather than a
|
|
// model artefact. Core's StoreBackend prepends it to the namespace it sends
|
|
// as the gRPC model name, and the local-store backend refuses loads without
|
|
// it — otherwise the model loader's greedy autoload (which probes every
|
|
// installed backend with LLM model names) would bind arbitrary models to
|
|
// the vector store, since store loads have no artefact to validate.
|
|
const NamespacePrefix = "store://"
|
|
|
|
// 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 {
|
|
res, err := c.StoresSet(ctx, &proto.StoresSetOptions{
|
|
Keys: WrapKeys(keys),
|
|
Values: WrapValues(values),
|
|
})
|
|
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 {
|
|
res, err := c.StoresDelete(ctx, &proto.StoresDeleteOptions{Keys: WrapKeys(keys)})
|
|
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) {
|
|
res, err := c.StoresGet(ctx, &proto.StoresGetOptions{Keys: WrapKeys(keys)})
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
return UnwrapKeys(res.Keys), UnwrapValues(res.Values), 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) {
|
|
res, err := c.StoresFind(ctx, &proto.StoresFindOptions{
|
|
Key: &proto.StoresKey{Floats: key},
|
|
TopK: int32(topk),
|
|
})
|
|
if err != nil {
|
|
return nil, nil, nil, err
|
|
}
|
|
return UnwrapKeys(res.Keys), UnwrapValues(res.Values), res.Similarities, nil
|
|
}
|