2024-01-23 07:56:36 +00:00
|
|
|
package grpc
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
2024-02-21 01:21:19 +00:00
|
|
|
|
2024-06-23 08:24:36 +00:00
|
|
|
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
2024-01-23 07:56:36 +00:00
|
|
|
"google.golang.org/grpc"
|
|
|
|
)
|
|
|
|
|
|
|
|
var embeds = map[string]*embedBackend{}
|
|
|
|
|
|
|
|
func Provide(addr string, llm LLM) {
|
|
|
|
embeds[addr] = &embedBackend{s: &server{llm: llm}}
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewClient(address string, parallel bool, wd WatchDog, enableWatchDog bool) Backend {
|
|
|
|
if bc, ok := embeds[address]; ok {
|
|
|
|
return bc
|
|
|
|
}
|
2024-08-25 12:36:09 +00:00
|
|
|
return buildClient(address, parallel, wd, enableWatchDog)
|
2024-01-23 07:56:36 +00:00
|
|
|
}
|
|
|
|
|
2024-08-25 12:36:09 +00:00
|
|
|
func buildClient(address string, parallel bool, wd WatchDog, enableWatchDog bool) Backend {
|
2024-01-23 07:56:36 +00:00
|
|
|
if !enableWatchDog {
|
|
|
|
wd = nil
|
|
|
|
}
|
|
|
|
return &Client{
|
|
|
|
address: address,
|
|
|
|
parallel: parallel,
|
|
|
|
wd: wd,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
type Backend interface {
|
|
|
|
IsBusy() bool
|
|
|
|
HealthCheck(ctx context.Context) (bool, error)
|
|
|
|
Embeddings(ctx context.Context, in *pb.PredictOptions, opts ...grpc.CallOption) (*pb.EmbeddingResult, error)
|
|
|
|
Predict(ctx context.Context, in *pb.PredictOptions, opts ...grpc.CallOption) (*pb.Reply, error)
|
|
|
|
LoadModel(ctx context.Context, in *pb.ModelOptions, opts ...grpc.CallOption) (*pb.Result, error)
|
2024-12-08 16:53:36 +00:00
|
|
|
PredictStream(ctx context.Context, in *pb.PredictOptions, f func(s []byte), opts ...grpc.CallOption) error
|
2024-01-23 07:56:36 +00:00
|
|
|
GenerateImage(ctx context.Context, in *pb.GenerateImageRequest, opts ...grpc.CallOption) (*pb.Result, error)
|
|
|
|
TTS(ctx context.Context, in *pb.TTSRequest, opts ...grpc.CallOption) (*pb.Result, error)
|
2024-08-24 00:20:28 +00:00
|
|
|
SoundGeneration(ctx context.Context, in *pb.SoundGenerationRequest, opts ...grpc.CallOption) (*pb.Result, error)
|
2024-09-02 13:48:53 +00:00
|
|
|
AudioTranscription(ctx context.Context, in *pb.TranscriptRequest, opts ...grpc.CallOption) (*pb.TranscriptResult, error)
|
2024-01-23 07:56:36 +00:00
|
|
|
TokenizeString(ctx context.Context, in *pb.PredictOptions, opts ...grpc.CallOption) (*pb.TokenizationResponse, error)
|
|
|
|
Status(ctx context.Context) (*pb.StatusResponse, error)
|
2024-03-22 20:14:04 +00:00
|
|
|
|
|
|
|
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)
|
|
|
|
StoresFind(ctx context.Context, in *pb.StoresFindOptions, opts ...grpc.CallOption) (*pb.StoresFindResult, error)
|
2024-04-24 22:19:02 +00:00
|
|
|
|
|
|
|
Rerank(ctx context.Context, in *pb.RerankRequest, opts ...grpc.CallOption) (*pb.RerankResult, error)
|
2024-10-01 12:41:20 +00:00
|
|
|
|
|
|
|
GetTokenMetrics(ctx context.Context, in *pb.MetricsRequest, opts ...grpc.CallOption) (*pb.MetricsResponse, error)
|
2024-11-20 13:48:40 +00:00
|
|
|
|
|
|
|
VAD(ctx context.Context, in *pb.VADRequest, opts ...grpc.CallOption) (*pb.VADResponse, error)
|
2024-01-23 07:56:36 +00:00
|
|
|
}
|