This commit is contained in:
Ettore Di Giacinto 2025-02-08 09:28:56 +08:00 committed by GitHub
commit 02e97f775e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 50 additions and 489 deletions

View file

@ -46,8 +46,7 @@ type Backend interface {
Status(ctx context.Context) (*pb.StatusResponse, error)
StoresSet(ctx context.Context, in *pb.StoresSetOptions, opts ...grpc.CallOption) (*pb.Result, error)
StoresDelete(ctx context.Context, in *pb.StoresDeleteOptions, opts ...grpc.CallOption) (*pb.Result, error)
StoresGet(ctx context.Context, in *pb.StoresGetOptions, opts ...grpc.CallOption) (*pb.StoresGetResult, error)
StoresReset(ctx context.Context, in *pb.StoresResetOptions, opts ...grpc.CallOption) (*pb.Result, error)
StoresFind(ctx context.Context, in *pb.StoresFindOptions, opts ...grpc.CallOption) (*pb.StoresFindResult, error)
Rerank(ctx context.Context, in *pb.RerankRequest, opts ...grpc.CallOption) (*pb.RerankResult, error)

View file

@ -80,11 +80,7 @@ func (llm *Base) StoresSet(*pb.StoresSetOptions) error {
return fmt.Errorf("unimplemented")
}
func (llm *Base) StoresGet(*pb.StoresGetOptions) (pb.StoresGetResult, error) {
return pb.StoresGetResult{}, fmt.Errorf("unimplemented")
}
func (llm *Base) StoresDelete(*pb.StoresDeleteOptions) error {
func (llm *Base) StoresReset(*pb.StoresResetOptions) error {
return fmt.Errorf("unimplemented")
}

View file

@ -303,7 +303,7 @@ func (c *Client) StoresSet(ctx context.Context, in *pb.StoresSetOptions, opts ..
return client.StoresSet(ctx, in, opts...)
}
func (c *Client) StoresDelete(ctx context.Context, in *pb.StoresDeleteOptions, opts ...grpc.CallOption) (*pb.Result, error) {
func (c *Client) StoreReset(ctx context.Context, in *pb.StoresResetOptions, opts ...grpc.CallOption) (*pb.Result, error) {
if !c.parallel {
c.opMutex.Lock()
defer c.opMutex.Unlock()
@ -318,25 +318,7 @@ func (c *Client) StoresDelete(ctx context.Context, in *pb.StoresDeleteOptions, o
}
defer conn.Close()
client := pb.NewBackendClient(conn)
return client.StoresDelete(ctx, in, opts...)
}
func (c *Client) StoresGet(ctx context.Context, in *pb.StoresGetOptions, opts ...grpc.CallOption) (*pb.StoresGetResult, error) {
if !c.parallel {
c.opMutex.Lock()
defer c.opMutex.Unlock()
}
c.setBusy(true)
defer c.setBusy(false)
c.wdMark()
defer c.wdUnMark()
conn, err := grpc.Dial(c.address, grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
return nil, err
}
defer conn.Close()
client := pb.NewBackendClient(conn)
return client.StoresGet(ctx, in, opts...)
return client.StoresReset(ctx, in, opts...)
}
func (c *Client) StoresFind(ctx context.Context, in *pb.StoresFindOptions, opts ...grpc.CallOption) (*pb.StoresFindResult, error) {

View file

@ -71,12 +71,8 @@ func (e *embedBackend) StoresSet(ctx context.Context, in *pb.StoresSetOptions, o
return e.s.StoresSet(ctx, in)
}
func (e *embedBackend) StoresDelete(ctx context.Context, in *pb.StoresDeleteOptions, opts ...grpc.CallOption) (*pb.Result, error) {
return e.s.StoresDelete(ctx, in)
}
func (e *embedBackend) StoresGet(ctx context.Context, in *pb.StoresGetOptions, opts ...grpc.CallOption) (*pb.StoresGetResult, error) {
return e.s.StoresGet(ctx, in)
func (e *embedBackend) StoresReset(ctx context.Context, in *pb.StoresResetOptions, opts ...grpc.CallOption) (*pb.Result, error) {
return e.s.StoresReset(ctx, in)
}
func (e *embedBackend) StoresFind(ctx context.Context, in *pb.StoresFindOptions, opts ...grpc.CallOption) (*pb.StoresFindResult, error) {

View file

@ -21,8 +21,7 @@ type LLM interface {
Status() (pb.StatusResponse, error)
StoresSet(*pb.StoresSetOptions) error
StoresDelete(*pb.StoresDeleteOptions) error
StoresGet(*pb.StoresGetOptions) (pb.StoresGetResult, error)
StoresReset(*pb.StoresResetOptions) error
StoresFind(*pb.StoresFindOptions) (pb.StoresFindResult, error)
VAD(*pb.VADRequest) (pb.VADResponse, error)

View file

@ -191,28 +191,16 @@ func (s *server) StoresSet(ctx context.Context, in *pb.StoresSetOptions) (*pb.Re
return &pb.Result{Message: "Set key", Success: true}, nil
}
func (s *server) StoresDelete(ctx context.Context, in *pb.StoresDeleteOptions) (*pb.Result, error) {
func (s *server) StoresReset(ctx context.Context, in *pb.StoresResetOptions) (*pb.Result, error) {
if s.llm.Locking() {
s.llm.Lock()
defer s.llm.Unlock()
}
err := s.llm.StoresDelete(in)
err := s.llm.StoresReset(in)
if err != nil {
return &pb.Result{Message: fmt.Sprintf("Error deleting entry: %s", err.Error()), Success: false}, err
}
return &pb.Result{Message: "Deleted key", Success: true}, nil
}
func (s *server) StoresGet(ctx context.Context, in *pb.StoresGetOptions) (*pb.StoresGetResult, error) {
if s.llm.Locking() {
s.llm.Lock()
defer s.llm.Unlock()
}
res, err := s.llm.StoresGet(in)
if err != nil {
return nil, err
}
return &res, nil
return &pb.Result{Message: "Deleted mem db", Success: true}, nil
}
func (s *server) StoresFind(ctx context.Context, in *pb.StoresFindOptions) (*pb.StoresFindResult, error) {

View file

@ -1,155 +0,0 @@
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
// 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
}