Revert "feat: include tokens usage for streamed output (#4282)"

This reverts commit 0d6c3a7d57.
This commit is contained in:
Ettore Di Giacinto 2024-12-08 16:31:48 +01:00
parent 87b7648591
commit 184fbc26bf
5 changed files with 10 additions and 25 deletions

View file

@ -35,7 +35,7 @@ func (e *embedBackend) LoadModel(ctx context.Context, in *pb.ModelOptions, opts
return e.s.LoadModel(ctx, in)
}
func (e *embedBackend) PredictStream(ctx context.Context, in *pb.PredictOptions, f func(reply *pb.Reply), opts ...grpc.CallOption) error {
func (e *embedBackend) PredictStream(ctx context.Context, in *pb.PredictOptions, f func(s []byte), opts ...grpc.CallOption) error {
bs := &embedBackendServerStream{
ctx: ctx,
fn: f,
@ -97,11 +97,11 @@ func (e *embedBackend) GetTokenMetrics(ctx context.Context, in *pb.MetricsReques
type embedBackendServerStream struct {
ctx context.Context
fn func(reply *pb.Reply)
fn func(s []byte)
}
func (e *embedBackendServerStream) Send(reply *pb.Reply) error {
e.fn(reply)
e.fn(reply.GetMessage())
return nil
}