mirror of
https://github.com/mudler/LocalAI.git
synced 2026-11-06 12:33:16 -05:00
feat(addons): move connector file bytes by handle
A connector add-on that declares files.transfer no longer carries file bytes in the connector envelope. The host puts a random handle in the request of the op media (write) or send_file (read), and the add-on streams the file with two new host operations, connector.file.put and connector.file.get, in chunks of at most 256 KiB. The process that runs the add-on registers the handle when the call arrives and ends it when the call returns. A handle is bound to the add-on, so another add-on gets the same answer as for an unknown handle. An unfinished inbound file is dropped with its partial bytes, and a handle with no call for two minutes expires. Inbound bytes go to the blob store of the workspace, outbound bytes come from it, so a cluster needs the shared object store and no shared disk. The file is never held in memory. The host overwrites the size and hash of an inbound file with its own record. The change is additive: the field files.transfer needs manifest API 2, a kind without it keeps the inline path and its 1 MiB cap, and the host never sends it a handle. A call that carries a handle may run up to 15 minutes. The ceiling of a file also stays under the largest file of a workspace. Assisted-by: Claude:claude-sonnet-5-5 go test
This commit is contained in:
21 files changed
+1453
-14
No files matched your search
@@ -11,6 +11,7 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/addons/contract"
|
||||
@@ -111,6 +112,13 @@ type Splitter interface {
|
||||
// means. A request for another kind, without a token or with an unknown op is
|
||||
// refused.
|
||||
func Serve(ctx context.Context, p Platform, req contract.ConnectorRequest) (contract.ConnectorResponse, error) {
|
||||
return ServeWith(ctx, p, nil, req)
|
||||
}
|
||||
|
||||
// ServeWith is Serve for an add-on that has the host operation channel, so it
|
||||
// can move file bytes by handle (see ConnectorTransfer). A request with a
|
||||
// handle that the Platform cannot serve is refused.
|
||||
func ServeWith(ctx context.Context, p Platform, host Host, req contract.ConnectorRequest) (contract.ConnectorResponse, error) {
|
||||
if req.Kind != p.Kind() || req.Token == "" {
|
||||
return contract.ConnectorResponse{}, fmt.Errorf("%w: unknown kind or no token", contract.ErrRefused)
|
||||
}
|
||||
@@ -149,6 +157,18 @@ func Serve(ctx context.Context, p Platform, req contract.ConnectorRequest) (cont
|
||||
out.BotID, out.BotUsername = bot.ID, bot.Username
|
||||
}
|
||||
case contract.ConnectorOpSendFile:
|
||||
if req.Transfer != nil {
|
||||
ss, ok := p.(StreamSender)
|
||||
if !ok || req.File != nil || req.Transfer.Handle == "" || req.Transfer.Size <= 0 || req.Transfer.Name != contract.SafeFileName(req.Transfer.Name) {
|
||||
return contract.ConnectorResponse{}, fmt.Errorf("%w: no file, or one this platform cannot send by handle", contract.ErrRefused)
|
||||
}
|
||||
err = Retry(ctx, func() error {
|
||||
id, err := ss.SendStream(ctx, req, Get(ctx, host, req.Transfer))
|
||||
out.MessageID = id
|
||||
return err
|
||||
})
|
||||
break
|
||||
}
|
||||
fs, ok := p.(FileSender)
|
||||
if !ok || !validFile(req.File) {
|
||||
return contract.ConnectorResponse{}, fmt.Errorf("%w: no file, or one this platform cannot send", contract.ErrRefused)
|
||||
@@ -159,6 +179,10 @@ func Serve(ctx context.Context, p Platform, req contract.ConnectorRequest) (cont
|
||||
return err
|
||||
})
|
||||
case contract.ConnectorOpMedia:
|
||||
if req.Transfer != nil {
|
||||
out.Received, err = receive(ctx, p, host, req)
|
||||
break
|
||||
}
|
||||
m, ok := p.(MediaReader)
|
||||
if !ok || req.Media == "" {
|
||||
return contract.ConnectorResponse{}, fmt.Errorf("%w: no media", contract.ErrRefused)
|
||||
@@ -188,12 +212,36 @@ func Serve(ctx context.Context, p Platform, req contract.ConnectorRequest) (cont
|
||||
default:
|
||||
return contract.ConnectorResponse{}, fmt.Errorf("%w: unknown op", contract.ErrRefused)
|
||||
}
|
||||
if errors.Is(err, contract.ErrRefused) {
|
||||
return contract.ConnectorResponse{}, err
|
||||
}
|
||||
if err != nil {
|
||||
return contract.ConnectorResponse{Error: Typed(err)}, nil
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// receive streams the file of a reference to the host by the handle.
|
||||
func receive(ctx context.Context, p Platform, host Host, req contract.ConnectorRequest) (*contract.ConnectorReceived, error) {
|
||||
sr, ok := p.(StreamReader)
|
||||
if !ok || req.Media == "" || req.Transfer.Handle == "" || req.Transfer.MaxBytes <= 0 {
|
||||
return nil, fmt.Errorf("%w: no media, or a platform that cannot stream it", contract.ErrRefused)
|
||||
}
|
||||
m, err := sr.MediaStream(ctx, req.Token, req.Media, req.Transfer.MaxBytes)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = m.Body.Close() }()
|
||||
if m.Size > req.Transfer.MaxBytes {
|
||||
return nil, ErrTooLarge
|
||||
}
|
||||
// One byte past the bound shows a file that is larger than it said.
|
||||
if _, err := Put(ctx, host, req.Transfer, io.LimitReader(m.Body, req.Transfer.MaxBytes+1)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &contract.ConnectorReceived{Name: contract.SafeFileName(m.Name), MIME: m.MIME}, nil
|
||||
}
|
||||
|
||||
type optionsKey struct{}
|
||||
|
||||
// OptionsFrom returns the options in effect that the host sent with the op
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package connector
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/addons/contract"
|
||||
)
|
||||
|
||||
// Host is the host operation channel of the add-on. *hostclient.Client
|
||||
// implements it.
|
||||
type Host interface {
|
||||
Do(ctx context.Context, operationID, op string, args, out any) error
|
||||
}
|
||||
|
||||
// transferTries is how often one chunk is sent before the transfer fails. A
|
||||
// repeat is safe: the host answers where it is, and a chunk it already has is
|
||||
// not written twice.
|
||||
const transferTries = 3
|
||||
|
||||
// ErrNoHost means the add-on got a handle but has no host operation channel.
|
||||
var ErrNoHost = errors.New("the host gave this add-on no channel for file transfer")
|
||||
|
||||
func hostError(err error) error {
|
||||
var he *contract.HostError
|
||||
if errors.As(err, &he) && he.Code == contract.HostTooLarge {
|
||||
return fmt.Errorf("%w: %v", ErrTooLarge, err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// Put streams r to the host through the handle of an inbound file and ends
|
||||
// the file. It holds one chunk in memory, whatever the size of the file, and
|
||||
// returns the bytes sent. The host refuses the file when it passes the bound
|
||||
// of the handle, and Put then returns ErrTooLarge.
|
||||
func Put(ctx context.Context, h Host, t *contract.ConnectorTransfer, r io.Reader) (int64, error) {
|
||||
if h == nil {
|
||||
return 0, ErrNoHost
|
||||
}
|
||||
buf := make([]byte, contract.TransferChunkBytes)
|
||||
var off int64
|
||||
for {
|
||||
n, rerr := io.ReadFull(r, buf)
|
||||
last := rerr != nil // a short read or the end
|
||||
if rerr != nil && !errors.Is(rerr, io.EOF) && !errors.Is(rerr, io.ErrUnexpectedEOF) {
|
||||
return off, rerr
|
||||
}
|
||||
if n == 0 && !last {
|
||||
continue
|
||||
}
|
||||
var res contract.ConnectorFilePutResult
|
||||
var err error
|
||||
for try := 0; try < transferTries; try++ {
|
||||
err = h.Do(ctx, t.Handle, contract.ConnectorFilePutOp, contract.ConnectorFilePut{Handle: t.Handle, Offset: off, Data: buf[:n], Final: last}, &res)
|
||||
if err == nil || ctx.Err() != nil || errors.Is(hostError(err), ErrTooLarge) {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return off, hostError(err)
|
||||
}
|
||||
if res.Received != off+int64(n) {
|
||||
// The host is not where the add-on is, and a stream cannot be
|
||||
// rewound; the platform download starts again.
|
||||
return off, fmt.Errorf("the host has %d bytes and the add-on sent %d", res.Received, off+int64(n))
|
||||
}
|
||||
off += int64(n)
|
||||
if last {
|
||||
return off, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns a reader of the outbound file of the handle. Each Read pulls at
|
||||
// most one chunk, so a platform upload that reads from it streams the file
|
||||
// and holds one chunk. A repeated read of a chunk is safe.
|
||||
func Get(ctx context.Context, h Host, t *contract.ConnectorTransfer) io.Reader {
|
||||
return &pull{ctx: ctx, h: h, t: t}
|
||||
}
|
||||
|
||||
type pull struct {
|
||||
ctx context.Context
|
||||
h Host
|
||||
t *contract.ConnectorTransfer
|
||||
off int64
|
||||
buf []byte
|
||||
eof bool
|
||||
}
|
||||
|
||||
func (p *pull) Read(b []byte) (int, error) {
|
||||
if p.h == nil {
|
||||
return 0, ErrNoHost
|
||||
}
|
||||
for len(p.buf) == 0 {
|
||||
if p.eof {
|
||||
return 0, io.EOF
|
||||
}
|
||||
var res contract.ConnectorFileGetResult
|
||||
var err error
|
||||
for try := 0; try < transferTries; try++ {
|
||||
err = p.h.Do(p.ctx, p.t.Handle, contract.ConnectorFileGetOp, contract.ConnectorFileGet{Handle: p.t.Handle, Offset: p.off, Length: contract.TransferChunkBytes}, &res)
|
||||
if err == nil || p.ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
p.buf, p.eof = res.Data, res.EOF
|
||||
p.off += int64(len(res.Data))
|
||||
if len(res.Data) == 0 && !res.EOF {
|
||||
return 0, errors.New("the host sent an empty chunk before the end of the file")
|
||||
}
|
||||
}
|
||||
n := copy(b, p.buf)
|
||||
p.buf = p.buf[n:]
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// Sniff returns a reader that gives the same bytes as r and the first bytes
|
||||
// of r, at most n, for a platform that decides the type of a file from them.
|
||||
func Sniff(r io.Reader, n int) (io.Reader, []byte) {
|
||||
br := bufio.NewReaderSize(r, max(n, 16))
|
||||
head, _ := br.Peek(n)
|
||||
return br, head
|
||||
}
|
||||
|
||||
// MediaStream is a file that a platform streams to the add-on.
|
||||
type MediaStream struct {
|
||||
Name, MIME string
|
||||
// Size is the size the platform states, or zero when it states none.
|
||||
Size int64
|
||||
// Body is the stream of bytes; Serve closes it.
|
||||
Body io.ReadCloser
|
||||
}
|
||||
|
||||
// StreamReader is a Platform that streams the files of messages for the op
|
||||
// media when the host gives a handle. max is the bound of the handle: the
|
||||
// platform should refuse a larger file as soon as it knows the size. A
|
||||
// Platform that does not implement it keeps the inline path.
|
||||
type StreamReader interface {
|
||||
MediaStream(ctx context.Context, token, ref string, max int64) (MediaStream, error)
|
||||
}
|
||||
|
||||
// StreamSender is a Platform that sends a file for the op send_file from a
|
||||
// stream. body is the file of req.Transfer: t.Size bytes. Serve builds a new
|
||||
// body for every attempt.
|
||||
type StreamSender interface {
|
||||
SendStream(ctx context.Context, req contract.ConnectorRequest, body io.Reader) (string, error)
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package connector_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/addons/internal/connector"
|
||||
"github.com/mudler/LocalAI/pkg/addons/contract"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// memHost is the host side of a handle in memory: it keeps what was put, and
|
||||
// serves what is to be got. fail makes the next calls fail.
|
||||
type memHost struct {
|
||||
got bytes.Buffer
|
||||
final bool
|
||||
src []byte
|
||||
fail int
|
||||
max int64
|
||||
tooLarge bool
|
||||
calls int
|
||||
}
|
||||
|
||||
func (h *memHost) Do(_ context.Context, _, op string, args, out any) error {
|
||||
h.calls++
|
||||
if h.fail > 0 {
|
||||
h.fail--
|
||||
return errors.New("network")
|
||||
}
|
||||
switch a := args.(type) {
|
||||
case contract.ConnectorFilePut:
|
||||
if h.tooLarge || (h.max > 0 && int64(h.got.Len()+len(a.Data)) > h.max) {
|
||||
return &contract.HostError{Code: contract.HostTooLarge, Message: "too large"}
|
||||
}
|
||||
if a.Offset == int64(h.got.Len()) {
|
||||
h.got.Write(a.Data)
|
||||
h.final = a.Final
|
||||
}
|
||||
*(out.(*contract.ConnectorFilePutResult)) = contract.ConnectorFilePutResult{Received: int64(h.got.Len())}
|
||||
case contract.ConnectorFileGet:
|
||||
end := min(a.Offset+a.Length, int64(len(h.src)))
|
||||
*(out.(*contract.ConnectorFileGetResult)) = contract.ConnectorFileGetResult{Data: h.src[a.Offset:end], EOF: end == int64(len(h.src))}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var _ = Describe("file transfer by handle", func() {
|
||||
tr := &contract.ConnectorTransfer{Handle: "h", MaxBytes: 1 << 30}
|
||||
|
||||
It("puts a stream in chunks and ends it, also when its size is a multiple of the chunk", func() {
|
||||
for _, n := range []int{1, contract.TransferChunkBytes - 1, contract.TransferChunkBytes, 3*contract.TransferChunkBytes + 7} {
|
||||
h := &memHost{}
|
||||
data := bytes.Repeat([]byte{'x'}, n)
|
||||
sent, err := connector.Put(context.Background(), h, tr, bytes.NewReader(data))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(sent).To(BeEquivalentTo(n))
|
||||
Expect(h.got.Bytes()).To(Equal(data))
|
||||
Expect(h.final).To(BeTrue())
|
||||
}
|
||||
})
|
||||
|
||||
It("sends a chunk again after a failed call, once the host says where it is", func() {
|
||||
h := &memHost{fail: 2}
|
||||
_, err := connector.Put(context.Background(), h, tr, strings.NewReader("hello"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(h.got.String()).To(Equal("hello"))
|
||||
h = &memHost{fail: 3}
|
||||
_, err = connector.Put(context.Background(), h, tr, strings.NewReader("hello"))
|
||||
Expect(err).To(HaveOccurred(), "three tries, then the transfer fails")
|
||||
})
|
||||
|
||||
It("reports a file over the bound of the handle as too large", func() {
|
||||
_, err := connector.Put(context.Background(), &memHost{tooLarge: true}, tr, strings.NewReader("hello"))
|
||||
Expect(err).To(MatchError(connector.ErrTooLarge))
|
||||
})
|
||||
|
||||
It("needs a host", func() {
|
||||
_, err := connector.Put(context.Background(), nil, tr, strings.NewReader("x"))
|
||||
Expect(err).To(MatchError(connector.ErrNoHost))
|
||||
_, err = io.ReadAll(connector.Get(context.Background(), nil, tr))
|
||||
Expect(err).To(MatchError(connector.ErrNoHost))
|
||||
})
|
||||
|
||||
It("gets a file chunk by chunk as a reader", func() {
|
||||
src := bytes.Repeat([]byte("abc"), contract.TransferChunkBytes)
|
||||
h := &memHost{src: src}
|
||||
got, err := io.ReadAll(connector.Get(context.Background(), h, tr))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(got).To(Equal(src))
|
||||
Expect(h.calls).To(Equal(3), "one call per chunk")
|
||||
})
|
||||
})
|
||||
|
||||
// streaming is a Platform that streams.
|
||||
type streaming struct {
|
||||
script
|
||||
body string
|
||||
size int64
|
||||
upload []byte
|
||||
gotName string
|
||||
}
|
||||
|
||||
func (s *streaming) MediaStream(_ context.Context, _, _ string, _ int64) (connector.MediaStream, error) {
|
||||
return connector.MediaStream{Name: "../a.txt", MIME: "text/plain", Size: s.size, Body: io.NopCloser(strings.NewReader(s.body))}, nil
|
||||
}
|
||||
|
||||
func (s *streaming) SendStream(_ context.Context, req contract.ConnectorRequest, body io.Reader) (string, error) {
|
||||
s.upload, _ = io.ReadAll(body)
|
||||
s.gotName = req.Transfer.Name
|
||||
return "f1", nil
|
||||
}
|
||||
|
||||
var _ = Describe("the connector route with a handle", func() {
|
||||
serve := func(p connector.Platform, h connector.Host, req contract.ConnectorRequest) (contract.ConnectorResponse, error) {
|
||||
req.Kind, req.Token = "chat", "t"
|
||||
return connector.ServeWith(context.Background(), p, h, req)
|
||||
}
|
||||
|
||||
It("streams the file of an op media to the host and answers with the name only", func() {
|
||||
h := &memHost{}
|
||||
out, err := serve(&streaming{body: "hello", size: 5}, h, contract.ConnectorRequest{Op: contract.ConnectorOpMedia, Media: "m", Transfer: &contract.ConnectorTransfer{Handle: "h", MaxBytes: 100}})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(out.Error).To(BeNil())
|
||||
Expect(out.File).To(BeNil())
|
||||
Expect(out.Received).To(Equal(&contract.ConnectorReceived{Name: "a.txt", MIME: "text/plain"}))
|
||||
Expect(h.got.String()).To(Equal("hello"))
|
||||
})
|
||||
|
||||
It("refuses a file the platform says is over the bound, before any byte moves", func() {
|
||||
h := &memHost{}
|
||||
out, _ := serve(&streaming{body: "hello", size: 500}, h, contract.ConnectorRequest{Op: contract.ConnectorOpMedia, Media: "m", Transfer: &contract.ConnectorTransfer{Handle: "h", MaxBytes: 100}})
|
||||
Expect(out.Error.Code).To(Equal(contract.ConnectorTooLarge))
|
||||
Expect(h.calls).To(BeZero())
|
||||
})
|
||||
|
||||
It("stops a platform that states a small size and sends more", func() {
|
||||
h := &memHost{max: 100}
|
||||
out, _ := serve(&streaming{body: strings.Repeat("x", 500), size: 5}, h, contract.ConnectorRequest{Op: contract.ConnectorOpMedia, Media: "m", Transfer: &contract.ConnectorTransfer{Handle: "h", MaxBytes: 100}})
|
||||
Expect(out.Error.Code).To(Equal(contract.ConnectorTooLarge))
|
||||
})
|
||||
|
||||
It("sends the file of an op send_file from the host's chunks", func() {
|
||||
p := &streaming{}
|
||||
h := &memHost{src: []byte("report body")}
|
||||
out, err := serve(p, h, contract.ConnectorRequest{Op: contract.ConnectorOpSendFile, Chat: "c", Transfer: &contract.ConnectorTransfer{Handle: "h", Name: "r.txt", MIME: "text/plain", Size: 11}})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(out.MessageID).To(Equal("f1"))
|
||||
Expect(string(p.upload)).To(Equal("report body"))
|
||||
Expect(p.gotName).To(Equal("r.txt"))
|
||||
})
|
||||
|
||||
It("refuses a handle on a platform that cannot stream, and a name that is not safe", func() {
|
||||
_, err := serve(&script{}, &memHost{}, contract.ConnectorRequest{Op: contract.ConnectorOpMedia, Media: "m", Transfer: &contract.ConnectorTransfer{Handle: "h", MaxBytes: 1}})
|
||||
Expect(err).To(MatchError(contract.ErrRefused))
|
||||
_, err = serve(&streaming{}, &memHost{}, contract.ConnectorRequest{Op: contract.ConnectorOpSendFile, Transfer: &contract.ConnectorTransfer{Handle: "h", Name: "../x", Size: 1}})
|
||||
Expect(err).To(MatchError(contract.ErrRefused))
|
||||
})
|
||||
})
|
||||
Reference in new issue
Block a user