mirror of
https://github.com/mudler/LocalAI.git
synced 2025-06-27 21:24:59 +00:00
Merge a1d5462ad0
into 230fe0098f
This commit is contained in:
commit
02e97f775e
16 changed files with 50 additions and 489 deletions
|
@ -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)
|
||||
|
|
|
@ -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")
|
||||
}
|
||||
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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
|
||||
}
|
Loading…
Add table
Add a link
Reference in a new issue