diff --git a/provider/llamaswap/audio.go b/provider/llamaswap/audio.go index 22ba06f..d4e168d 100644 --- a/provider/llamaswap/audio.go +++ b/provider/llamaswap/audio.go @@ -63,7 +63,7 @@ func (m *speechModel) Speak(ctx context.Context, req audio.SpeechRequest, opts . if err != nil { return nil, fmt.Errorf("llama-swap: encode speech request: %w", err) } - audioBytes, contentType, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/speech", m.id, "application/json", bytes.NewReader(body)) + audioBytes, contentType, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/speech", m.id, "application/json", bytes.NewReader(body), maxResponseBytes) if err != nil { return nil, err } @@ -140,7 +140,7 @@ func (m *transcriptionModel) Transcribe(ctx context.Context, req audio.Transcrip return nil, fmt.Errorf("llama-swap: build transcription form: %w", err) } - raw, _, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/transcriptions", m.id, w.FormDataContentType(), &buf) + raw, _, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/transcriptions", m.id, w.FormDataContentType(), &buf, maxResponseBytes) if err != nil { return nil, err } @@ -201,7 +201,7 @@ func (p *Provider) ListVoices(ctx context.Context, model string) ([]string, erro if model == "" { return nil, fmt.Errorf("llama-swap: ListVoices requires a model id") } - raw, _, err := p.doRaw(ctx, http.MethodGet, "/v1/audio/voices?model="+url.QueryEscape(model), model, "", nil) + raw, _, err := p.doRaw(ctx, http.MethodGet, "/v1/audio/voices?model="+url.QueryEscape(model), model, "", nil, maxResponseBytes) if err != nil { return nil, err } @@ -256,10 +256,10 @@ func parseVoices(raw []byte) ([]string, error) { // doRaw performs a request to a llama-swap endpoint and returns the raw // response body and its Content-Type — the sibling of doJSON for endpoints -// whose success payload is not JSON (audio bytes) or whose shape varies. -// contentType sets the request Content-Type when body is non-nil. A response -// larger than maxResponseBytes is an error, never a silent truncation. -func (p *Provider) doRaw(ctx context.Context, method, path, model, contentType string, body io.Reader) ([]byte, string, error) { +// whose success payload is not JSON (audio/video bytes) or whose shape +// varies. contentType sets the request Content-Type when body is non-nil. +// A response larger than maxBytes is an error, never a silent truncation. +func (p *Provider) doRaw(ctx context.Context, method, path, model, contentType string, body io.Reader, maxBytes int64) ([]byte, string, error) { if err := p.requireBaseURL(); err != nil { return nil, "", err } @@ -275,12 +275,12 @@ func (p *Provider) doRaw(ctx context.Context, method, path, model, contentType s if resp.StatusCode/100 != 2 { return nil, "", p.apiError(resp, model) } - data, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1)) + data, err := io.ReadAll(io.LimitReader(resp.Body, maxBytes+1)) if err != nil { return nil, "", fmt.Errorf("llama-swap: read response: %w", err) } - if len(data) > maxResponseBytes { - return nil, "", fmt.Errorf("llama-swap: response exceeds %d bytes", maxResponseBytes) + if int64(len(data)) > maxBytes { + return nil, "", fmt.Errorf("llama-swap: response exceeds %d bytes", maxBytes) } return data, resp.Header.Get("Content-Type"), nil } diff --git a/provider/llamaswap/llamaswap.go b/provider/llamaswap/llamaswap.go index 8dedece..7713c02 100644 --- a/provider/llamaswap/llamaswap.go +++ b/provider/llamaswap/llamaswap.go @@ -47,6 +47,12 @@ const DefaultName = "llama-swap" // can't make a decode allocate without limit. const maxResponseBytes = 64 << 20 +// maxVideoResponseBytes caps the /v1/videos/sync body — the response IS an +// encoded clip, which legitimately dwarfs any JSON/audio payload (a long +// high-bitrate generation can pass 64MB). Still bounded so a buggy upstream +// can't allocate without limit. +const maxVideoResponseBytes = 512 << 20 + // Provider is a llama-swap client. It satisfies llm.Provider (chat, delegated // to provider/openai) and imagegen.Provider (image generation), and exposes // llama-swap's management endpoints as concrete methods. diff --git a/provider/llamaswap/video.go b/provider/llamaswap/video.go index b0e893b..caa584e 100644 --- a/provider/llamaswap/video.go +++ b/provider/llamaswap/video.go @@ -96,7 +96,7 @@ func (m *videoModel) Generate(ctx context.Context, req videogen.Request, opts .. return nil, fmt.Errorf("llama-swap: build video form: %w", err) } - videoBytes, contentType, err := m.p.doRaw(ctx, http.MethodPost, "/v1/videos/sync", m.id, w.FormDataContentType(), &buf) + videoBytes, contentType, err := m.p.doRaw(ctx, http.MethodPost, "/v1/videos/sync", m.id, w.FormDataContentType(), &buf, maxVideoResponseBytes) if err != nil { return nil, err } @@ -145,7 +145,7 @@ func initImageFilename(mimeType string) string { return "frame.jpg" case "image/webp": return "frame.webp" - default: // stable-diffusion outputs and unknowns: PNG is the safe hint + default: // unknown MIME — PNG is the safe hint return "frame.png" } } diff --git a/provider/llamaswap/video_test.go b/provider/llamaswap/video_test.go index 2684279..9958b00 100644 --- a/provider/llamaswap/video_test.go +++ b/provider/llamaswap/video_test.go @@ -23,6 +23,7 @@ func TestVideoGenerate(t *testing.T) { gotContentType = r.Header.Get("Content-Type") if err := r.ParseMultipartForm(32 << 20); err != nil { t.Errorf("parse form: %v", err) + return } gotForm = map[string]string{} for k, v := range r.MultipartForm.Value {