feat: musicgen + embeddings/rerank surfaces (ADR-0021, ADR-0022) #15
@@ -0,0 +1,35 @@
|
||||
# ADR-0021: musicgen interface (blocking Generate over an async job queue)
|
||||
|
||||
Status: Accepted (2026-07-12)
|
||||
|
||||
## Context
|
||||
|
||||
The llama-swap host gained ACE-Step 1.5 (full songs with vocals, seconds per
|
||||
clip warm on the reference GPU). Its API is an async job queue —
|
||||
`POST /release_task` → poll `POST /query_result` → `GET` the result file —
|
||||
unlike every other majordomo media backend, which is one blocking call.
|
||||
|
||||
## Decision
|
||||
|
||||
- New `musicgen` leaf package, ADR-0016→0020 conventions: `Request{Prompt,
|
||||
Lyrics, DurationSeconds, Format, Steps, Seed}`, `Result{Audio{Data, MIME},
|
||||
Raw}`, functional options, zero value = backend default, bytes-only.
|
||||
- **`Model.Generate` blocks, polling internally** (2s interval, ctx-bounded).
|
||||
The one-call contract is the package's value: callers budget the whole job
|
||||
with a context deadline exactly like imagegen/videogen, and the async
|
||||
mechanics stay a provider detail. An async job surface (mirroring the
|
||||
deliberately-deferred videogen one, ADR-0019) can come later if a caller
|
||||
ever needs progress.
|
||||
- provider/llamaswap reaches ACE-Step through the `/upstream/<model>/`
|
||||
passthrough (ADR-0020). Envelope parsing is tolerant (`data` wrapper or
|
||||
bare array; `result` arrives as a double-encoded JSON string) and the
|
||||
result-file URL is routed back through the same upstream. `audio_duration`
|
||||
is the v1 param name — an unknown field degrades to default clip length
|
||||
upstream, never an error; verify at smoke.
|
||||
|
||||
## Consequences
|
||||
|
||||
- Music jobs occupy the exclusive GPU group like video; callers own the
|
||||
timeout budget (mort keeps the tool timeout under its agent ceiling).
|
||||
- The double-encoded `result` string and envelope shapes are pinned by the
|
||||
netherstorm image build; host smoke tests are the drift defence.
|
||||
@@ -0,0 +1,41 @@
|
||||
# ADR-0022: embeddings + rerank interface
|
||||
|
||||
Status: Accepted (2026-07-12)
|
||||
|
||||
## Context
|
||||
|
||||
majordomo had no embedding or reranking surface at all. The llama-swap host
|
||||
now runs two persistent CPU-only llama-server members (Qwen3-Embedding-0.6B
|
||||
via `/v1/embeddings`, bge-reranker-v2-m3 via `/v1/rerank`), and mort wants a
|
||||
reranking stage in memory retrieval with embedding-backed retrieval as a
|
||||
later step.
|
||||
|
||||
## Decision
|
||||
|
||||
- New `embeddings` leaf package with TWO half-surfaces, split like audio's
|
||||
Speech/Transcription: `EmbedModel`/`EmbedProvider` and
|
||||
`RerankModel`/`RerankProvider`. They are separate mints because on the
|
||||
reference host they are two DIFFERENT server instances — llama-server with
|
||||
`--embeddings` and `--rerank` together returns all-zero embeddings
|
||||
(llama.cpp #20085) — and because rerankers are cross-encoders, not
|
||||
embedders.
|
||||
- `EmbedResult.Vectors [][]float32` in input order (provider must order by
|
||||
the response's `index`, never trust wire order). No `dimensions` param:
|
||||
llama-server doesn't implement it; Matryoshka truncation is caller-side.
|
||||
- `InstructedQuery(task, query)` helper encodes the instruction-aware
|
||||
asymmetry (queries wrapped, documents bare) so call sites can't silently
|
||||
degrade retrieval by forgetting the prefix.
|
||||
- `RerankResult` sorted by descending score; parser reads ONLY
|
||||
`results[].index` and `results[].relevance_score` because llama-server
|
||||
documents the shape as subject to change. Scores are model-specific —
|
||||
comparable within one response only.
|
||||
|
||||
## Consequences
|
||||
|
||||
- Callers get vectors/scores with strict validation (count mismatch, index
|
||||
out of range, empty vector are hard errors — a silently missing vector is
|
||||
a retrieval bug factory).
|
||||
- llama-server's rerank scoring has open correctness issues for some models
|
||||
(llama.cpp #16407); consumers must validate against a fixture before
|
||||
trusting scores in production (mort gates its memory-rerank convar on
|
||||
exactly that).
|
||||
@@ -0,0 +1,166 @@
|
||||
// Package embeddings is majordomo's canonical text-embedding and reranking
|
||||
// surface (ADR-0022, following the ADR-0016→0021 leaf-contract lineage).
|
||||
// Two small halves, split like audio's Speech/Transcription so backends can
|
||||
// implement either:
|
||||
//
|
||||
// - EmbedModel turns texts into dense vectors (/v1/embeddings-style).
|
||||
// - RerankModel scores documents against a query with a cross-encoder
|
||||
// (/v1/rerank-style) — usually a DIFFERENT backend model than the
|
||||
// embedder, hence a separate mint.
|
||||
//
|
||||
// Instruction-aware embedders (Qwen3-Embedding et al.) want queries wrapped
|
||||
// as "Instruct: {task}\nQuery: {query}" while documents go in bare;
|
||||
// InstructedQuery encodes that so callers don't hand-roll (and silently
|
||||
// degrade retrieval) at each site.
|
||||
package embeddings
|
||||
|
||||
import "context"
|
||||
|
||||
// EmbedRequest is a batch embedding request. Inputs are embedded
|
||||
// independently; the result vector order matches the input order.
|
||||
type EmbedRequest struct {
|
||||
// Inputs are the texts to embed. Required (at least one).
|
||||
Inputs []string
|
||||
}
|
||||
|
||||
// EmbedResult is the canonical embedding result.
|
||||
type EmbedResult struct {
|
||||
// Vectors holds one embedding per input, in input order. Backends
|
||||
// normalize per their own convention (llama-server: Euclidean-normalized).
|
||||
Vectors [][]float32
|
||||
|
||||
// Raw is the provider-native response object. May be nil.
|
||||
Raw any
|
||||
}
|
||||
|
||||
// EmbedOption mutates an EmbedRequest before it is sent. Reserved: the
|
||||
// request shape is deliberately minimal today.
|
||||
type EmbedOption func(*EmbedRequest)
|
||||
|
||||
// Apply returns a copy of the request with all options applied.
|
||||
func (r EmbedRequest) Apply(opts ...EmbedOption) EmbedRequest {
|
||||
for _, opt := range opts {
|
||||
opt(&r)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// InstructedQuery wraps a retrieval QUERY for instruction-aware embedding
|
||||
// models. Documents must NOT be wrapped — the asymmetry is the point, and
|
||||
// getting it wrong silently costs retrieval quality. An empty task uses the
|
||||
// generic web-search instruction the reference model was trained with.
|
||||
func InstructedQuery(task, query string) string {
|
||||
if task == "" {
|
||||
task = "Given a web search query, retrieve relevant passages that answer the query"
|
||||
}
|
||||
return "Instruct: " + task + "\nQuery: " + query
|
||||
}
|
||||
|
||||
// EmbedModel embeds texts as dense vectors.
|
||||
type EmbedModel interface {
|
||||
// Embed returns one vector per input, in input order.
|
||||
Embed(ctx context.Context, req EmbedRequest, opts ...EmbedOption) (*EmbedResult, error)
|
||||
}
|
||||
|
||||
// EmbedModelOption configures an EmbedModel at construction time. Reserved
|
||||
// for future per-model settings.
|
||||
type EmbedModelOption func(*EmbedModelConfig)
|
||||
|
||||
// EmbedModelConfig carries per-model construction settings.
|
||||
type EmbedModelConfig struct{}
|
||||
|
||||
// ApplyEmbedModelOptions folds options into a config.
|
||||
func ApplyEmbedModelOptions(opts []EmbedModelOption) EmbedModelConfig {
|
||||
var cfg EmbedModelConfig
|
||||
for _, opt := range opts {
|
||||
opt(&cfg)
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// EmbedProvider mints embedding models bound to one backend.
|
||||
type EmbedProvider interface {
|
||||
// Name is the registry identifier for the provider.
|
||||
Name() string
|
||||
|
||||
// EmbedModel returns an EmbedModel bound to the given id (passed through
|
||||
// to the backend verbatim; no catalog validation).
|
||||
EmbedModel(id string, opts ...EmbedModelOption) (EmbedModel, error)
|
||||
}
|
||||
|
||||
// RerankRequest scores documents against a query.
|
||||
type RerankRequest struct {
|
||||
// Query is the search query. Required.
|
||||
Query string
|
||||
|
||||
// Documents are the candidate texts to score. Required (at least one).
|
||||
Documents []string
|
||||
|
||||
// TopN limits how many results the backend returns; 0 = all.
|
||||
TopN int
|
||||
}
|
||||
|
||||
// RerankItem is one scored document.
|
||||
type RerankItem struct {
|
||||
// Index is the document's position in the request's Documents slice.
|
||||
Index int
|
||||
|
||||
// Score is the backend's relevance score (higher = more relevant).
|
||||
// Scales are model-specific — compare within one response only.
|
||||
Score float64
|
||||
}
|
||||
|
||||
// RerankResult is the canonical rerank result, sorted by descending Score.
|
||||
type RerankResult struct {
|
||||
// Results are the scored documents (top-N when the request bounded it).
|
||||
Results []RerankItem
|
||||
|
||||
// Raw is the provider-native response object. May be nil.
|
||||
Raw any
|
||||
}
|
||||
|
||||
// RerankOption mutates a RerankRequest before it is sent.
|
||||
type RerankOption func(*RerankRequest)
|
||||
|
||||
// WithTopN limits how many results the backend returns.
|
||||
func WithTopN(n int) RerankOption { return func(r *RerankRequest) { r.TopN = n } }
|
||||
|
||||
// Apply returns a copy of the request with all options applied.
|
||||
func (r RerankRequest) Apply(opts ...RerankOption) RerankRequest {
|
||||
for _, opt := range opts {
|
||||
opt(&r)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// RerankModel scores documents against a query with a cross-encoder.
|
||||
type RerankModel interface {
|
||||
// Rerank returns scored documents sorted by descending relevance.
|
||||
Rerank(ctx context.Context, req RerankRequest, opts ...RerankOption) (*RerankResult, error)
|
||||
}
|
||||
|
||||
// RerankModelOption configures a RerankModel at construction time. Reserved
|
||||
// for future per-model settings.
|
||||
type RerankModelOption func(*RerankModelConfig)
|
||||
|
||||
// RerankModelConfig carries per-model construction settings.
|
||||
type RerankModelConfig struct{}
|
||||
|
||||
// ApplyRerankModelOptions folds options into a config.
|
||||
func ApplyRerankModelOptions(opts []RerankModelOption) RerankModelConfig {
|
||||
var cfg RerankModelConfig
|
||||
for _, opt := range opts {
|
||||
opt(&cfg)
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// RerankProvider mints rerank models bound to one backend.
|
||||
type RerankProvider interface {
|
||||
// Name is the registry identifier for the provider.
|
||||
Name() string
|
||||
|
||||
// RerankModel returns a RerankModel bound to the given id (passed
|
||||
// through to the backend verbatim; no catalog validation).
|
||||
RerankModel(id string, opts ...RerankModelOption) (RerankModel, error)
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
// Package musicgen is majordomo's canonical music-generation surface (full
|
||||
// songs from a text prompt, optionally with lyrics). Like imagegen/audio/
|
||||
// videogen/meshgen, it is a deliberately separate leaf contract from the llm
|
||||
// package (ADR-0021, following the ADR-0016→0020 lineage: functional
|
||||
// options, zero values = backend default, bytes-only I/O, Raw escape hatch).
|
||||
//
|
||||
// The first implementation is provider/llamaswap, which targets an
|
||||
// ACE-Step-1.5-style async job API reached through the /upstream
|
||||
// passthrough; Generate blocks (polling internally) so callers get the
|
||||
// imagegen-style one-call contract — bound it with a context deadline.
|
||||
package musicgen
|
||||
|
||||
import "context"
|
||||
|
||||
// Audio is one generated piece: raw encoded bytes plus a MIME type.
|
||||
type Audio struct {
|
||||
// Data is the encoded audio container.
|
||||
Data []byte
|
||||
|
||||
// MIME is the audio MIME type, e.g. "audio/mpeg".
|
||||
MIME string
|
||||
}
|
||||
|
||||
// Request is a music-generation request. Zero values mean "backend default".
|
||||
type Request struct {
|
||||
// Prompt describes the music (genre, mood, instrumentation, tempo...).
|
||||
Prompt string
|
||||
|
||||
// Lyrics are optional song lyrics for models that sing; "" =
|
||||
// instrumental or model-written lyrics, per backend behavior.
|
||||
Lyrics string
|
||||
|
||||
// DurationSeconds is the requested clip length; 0 = backend default.
|
||||
DurationSeconds int
|
||||
|
||||
// Format is the audio container ("mp3", "wav", "flac", "opus", "aac");
|
||||
// "" = backend default (mp3 on ACE-Step).
|
||||
Format string
|
||||
|
||||
// Steps is the number of inference steps; nil = backend default.
|
||||
Steps *int
|
||||
|
||||
// Seed fixes the RNG seed for reproducible output; nil = backend
|
||||
// default (random).
|
||||
Seed *int64
|
||||
}
|
||||
|
||||
// Result is the canonical music-generation result.
|
||||
type Result struct {
|
||||
// Audio is the generated piece.
|
||||
Audio Audio
|
||||
|
||||
// Raw is the provider-native response object (e.g. the job result with
|
||||
// bpm/keyscale metadata). May be nil.
|
||||
Raw any
|
||||
}
|
||||
|
||||
// Option mutates a Request before it is sent. Options passed to Generate are
|
||||
// applied to a copy of the request, so a Request value can be reused.
|
||||
type Option func(*Request)
|
||||
|
||||
// WithLyrics sets song lyrics.
|
||||
func WithLyrics(l string) Option { return func(r *Request) { r.Lyrics = l } }
|
||||
|
||||
// WithDuration sets the requested clip length in seconds.
|
||||
func WithDuration(seconds int) Option {
|
||||
return func(r *Request) { r.DurationSeconds = seconds }
|
||||
}
|
||||
|
||||
// WithFormat sets the audio container format.
|
||||
func WithFormat(f string) Option { return func(r *Request) { r.Format = f } }
|
||||
|
||||
// WithSteps overrides the number of inference steps.
|
||||
func WithSteps(n int) Option { return func(r *Request) { r.Steps = &n } }
|
||||
|
||||
// WithSeed fixes the RNG seed.
|
||||
func WithSeed(seed int64) Option { return func(r *Request) { r.Seed = &seed } }
|
||||
|
||||
// Apply returns a copy of the request with all options applied. Providers
|
||||
// call this once at the top of Generate.
|
||||
func (r Request) Apply(opts ...Option) Request {
|
||||
for _, opt := range opts {
|
||||
opt(&r)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// Model generates music from text.
|
||||
type Model interface {
|
||||
// Generate renders the request as one audio clip. It blocks until the
|
||||
// backend finishes (or ctx expires) even when the backend is an async
|
||||
// job queue — polling is the provider's business, not the caller's.
|
||||
Generate(ctx context.Context, req Request, opts ...Option) (*Result, error)
|
||||
}
|
||||
|
||||
// ModelOption configures a Model at construction time. Reserved for future
|
||||
// per-model settings.
|
||||
type ModelOption func(*ModelConfig)
|
||||
|
||||
// ModelConfig carries per-model construction settings.
|
||||
type ModelConfig struct{}
|
||||
|
||||
// ApplyModelOptions folds options into a config.
|
||||
func ApplyModelOptions(opts []ModelOption) ModelConfig {
|
||||
var cfg ModelConfig
|
||||
for _, opt := range opts {
|
||||
opt(&cfg)
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// Provider mints music models bound to one backend.
|
||||
type Provider interface {
|
||||
// Name is the registry identifier for the provider.
|
||||
Name() string
|
||||
|
||||
// MusicModel returns a Model bound to the given id (passed through to
|
||||
// the backend verbatim; no catalog validation).
|
||||
MusicModel(id string, opts ...ModelOption) (Model, error)
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
// embed.go implements embeddings.EmbedProvider and embeddings.RerankProvider
|
||||
// against llama-server instances behind llama-swap (ADR-0022):
|
||||
//
|
||||
// POST /v1/embeddings {model, input: [...]} (OpenAI shape)
|
||||
// POST /v1/rerank {model, query, documents, top_n} (Jina-ish shape)
|
||||
//
|
||||
// Both paths are in llama-swap's normal model-routed tables — no /upstream
|
||||
// needed. The two surfaces are minted separately because they are two
|
||||
// DIFFERENT server instances on the host: llama-server with --embeddings
|
||||
// and --rerank enabled together returns all-zero embeddings (llama.cpp
|
||||
// #20085), so the host runs one of each and the ids differ.
|
||||
package llamaswap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/embeddings"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
||||
)
|
||||
|
||||
// EmbedModel implements embeddings.EmbedProvider. The id selects which
|
||||
// upstream llama-swap loads (a persistent CPU member on the reference host,
|
||||
// so calls are cheap and never evict GPU models).
|
||||
func (p *Provider) EmbedModel(id string, opts ...embeddings.EmbedModelOption) (embeddings.EmbedModel, error) {
|
||||
if err := p.requireBaseURL(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = embeddings.ApplyEmbedModelOptions(opts)
|
||||
return &embedModel{p: p, id: id}, nil
|
||||
}
|
||||
|
||||
type embedModel struct {
|
||||
p *Provider
|
||||
id string
|
||||
}
|
||||
|
||||
// Embed implements embeddings.EmbedModel via POST {base}/v1/embeddings.
|
||||
func (m *embedModel) Embed(ctx context.Context, req embeddings.EmbedRequest, opts ...embeddings.EmbedOption) (*embeddings.EmbedResult, error) {
|
||||
|
|
||||
req = req.Apply(opts...)
|
||||
if len(req.Inputs) == 0 {
|
||||
return nil, fmt.Errorf("%w: embedding requires at least one input", llm.ErrUnsupported)
|
||||
}
|
||||
for i, in := range req.Inputs {
|
||||
if strings.TrimSpace(in) == "" {
|
||||
return nil, fmt.Errorf("%w: embedding input %d is empty", llm.ErrUnsupported, i)
|
||||
}
|
||||
}
|
||||
wire := struct {
|
||||
Model string `json:"model"`
|
||||
Input []string `json:"input"`
|
||||
}{Model: m.id, Input: req.Inputs}
|
||||
|
||||
var resp struct {
|
||||
Data []struct {
|
||||
Index int `json:"index"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := m.p.doJSON(ctx, http.MethodPost, "/v1/embeddings", m.id, &wire, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(resp.Data) != len(req.Inputs) {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: fmt.Sprintf("embeddings response has %d vectors for %d inputs", len(resp.Data), len(req.Inputs))}
|
||||
}
|
||||
// The OpenAI shape carries an index per entry; order by it rather than
|
||||
// trusting response order.
|
||||
vectors := make([][]float32, len(req.Inputs))
|
||||
|
gitea-actions
commented
🟡 Duplicate indices in embedding response silently overwrite, causing misleading error messages error-handling · flagged by 1 model
🪰 Gadfly · advisory 🟡 **Duplicate indices in embedding response silently overwrite, causing misleading error messages**
_error-handling · flagged by 1 model_
2. **`provider/llamaswap/embed.go:72-78` — Duplicate indices in embedding response silently overwrite** If the provider returns duplicate `index` values (e.g., two entries both with `index: 0`), the second overwrites the first (`vectors[d.Index] = d.Embedding`). The subsequent nil-check loop (lines 80-85) will catch the resulting gap, but the error message ("missing vector for input X") misattributes the root cause. **Impact:** Confusing error messages if a buggy provider returns duplicate indic…
<sub>🪰 Gadfly · advisory</sub>
|
||||
for _, d := range resp.Data {
|
||||
if d.Index < 0 || d.Index >= len(vectors) || len(d.Embedding) == 0 {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: fmt.Sprintf("embeddings response entry index %d invalid or empty", d.Index)}
|
||||
}
|
||||
if vectors[d.Index] != nil {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: fmt.Sprintf("embeddings response repeats index %d", d.Index)}
|
||||
}
|
||||
vectors[d.Index] = d.Embedding
|
||||
}
|
||||
for i, v := range vectors {
|
||||
if v == nil {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: fmt.Sprintf("embeddings response missing vector for input %d", i)}
|
||||
}
|
||||
}
|
||||
return &embeddings.EmbedResult{Vectors: vectors, Raw: &resp}, nil
|
||||
}
|
||||
|
||||
// RerankModel implements embeddings.RerankProvider.
|
||||
func (p *Provider) RerankModel(id string, opts ...embeddings.RerankModelOption) (embeddings.RerankModel, error) {
|
||||
if err := p.requireBaseURL(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = embeddings.ApplyRerankModelOptions(opts)
|
||||
return &rerankModel{p: p, id: id}, nil
|
||||
}
|
||||
|
||||
type rerankModel struct {
|
||||
p *Provider
|
||||
id string
|
||||
}
|
||||
|
||||
|
gitea-actions
commented
🟠 Missing input length validation on rerank Query/Documents enables DoS security · flagged by 1 model
🪰 Gadfly · advisory 🟠 **Missing input length validation on rerank Query/Documents enables DoS**
_security · flagged by 1 model_
- **`provider/llamaswap/embed.go:106-116` — Missing input length validation for rerank (DoS vector)** Same issue as embeddings: `Query` and `Documents` are validated for emptiness but not bounded in length. **Fix**: Add per-document and query length limits. **Verified**: Read `embed.go:106-116`; confirmed only emptiness checks exist.
<sub>🪰 Gadfly · advisory</sub>
|
||||
// Rerank implements embeddings.RerankModel via POST {base}/v1/rerank. The
|
||||
// response parser reads only results[].index and results[].relevance_score —
|
||||
// llama-server documents the shape as "might change", so stay minimal.
|
||||
func (m *rerankModel) Rerank(ctx context.Context, req embeddings.RerankRequest, opts ...embeddings.RerankOption) (*embeddings.RerankResult, error) {
|
||||
req = req.Apply(opts...)
|
||||
if strings.TrimSpace(req.Query) == "" {
|
||||
return nil, fmt.Errorf("%w: rerank requires a query", llm.ErrUnsupported)
|
||||
}
|
||||
if len(req.Documents) == 0 {
|
||||
return nil, fmt.Errorf("%w: rerank requires at least one document", llm.ErrUnsupported)
|
||||
}
|
||||
if req.TopN < 0 {
|
||||
return nil, fmt.Errorf("%w: rerank top_n must be >= 0, got %d", llm.ErrUnsupported, req.TopN)
|
||||
}
|
||||
wire := struct {
|
||||
Model string `json:"model"`
|
||||
Query string `json:"query"`
|
||||
Documents []string `json:"documents"`
|
||||
TopN int `json:"top_n,omitempty"`
|
||||
}{Model: m.id, Query: req.Query, Documents: req.Documents, TopN: req.TopN}
|
||||
|
||||
var resp struct {
|
||||
Results []struct {
|
||||
Index int `json:"index"`
|
||||
RelevanceScore float64 `json:"relevance_score"`
|
||||
} `json:"results"`
|
||||
}
|
||||
if err := m.p.doJSON(ctx, http.MethodPost, "/v1/rerank", m.id, &wire, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(resp.Results) == 0 {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "rerank response contained no results"}
|
||||
}
|
||||
out := &embeddings.RerankResult{Raw: &resp}
|
||||
for _, r := range resp.Results {
|
||||
if r.Index < 0 || r.Index >= len(req.Documents) {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: fmt.Sprintf("rerank result index %d out of range", r.Index)}
|
||||
}
|
||||
out.Results = append(out.Results, embeddings.RerankItem{Index: r.Index, Score: r.RelevanceScore})
|
||||
}
|
||||
sort.SliceStable(out.Results, func(i, j int) bool { return out.Results[i].Score > out.Results[j].Score })
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,265 @@
|
||||
// music.go implements musicgen.Provider against an ACE-Step-1.5-style API
|
||||
// server reached through llama-swap's /upstream passthrough (ADR-0021):
|
||||
//
|
||||
// POST /upstream/<id>/release_task {prompt, lyrics, ...} -> {task_id}
|
||||
// POST /upstream/<id>/query_result {task_id_list: [...]} -> status+result
|
||||
// GET /upstream/<id>/<result file URL> -> audio bytes
|
||||
//
|
||||
// The backend is an async job queue; Generate wraps it into the blocking
|
||||
// one-call contract by polling, so a context deadline is the caller's
|
||||
// budget for the whole job (mort's tool timeout sits well under its
|
||||
// agent-runtime ceiling for exactly this reason).
|
||||
package llamaswap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/musicgen"
|
||||
)
|
||||
|
||||
// musicPollInterval is the delay between query_result polls. Long enough to
|
||||
// be polite to the queue, short enough that a ~10s xl-turbo song isn't
|
||||
// dominated by poll latency. A var so tests can shrink it.
|
||||
var musicPollInterval = 2 * time.Second
|
||||
|
gitea-actions
commented
🟡 Hardcoded 2s poll interval forces TestMusicGenerate to burn 2+ real seconds on every CI run performance · flagged by 1 model
🪰 Gadfly · advisory 🟡 **Hardcoded 2s poll interval forces TestMusicGenerate to burn 2+ real seconds on every CI run**
_performance · flagged by 1 model_
`provider/llamaswap/music.go:29`
<sub>🪰 Gadfly · advisory</sub>
|
||||
|
||||
// musicPollMaxConsecutiveFailures bounds how many consecutive BAD polls
|
||||
// (transport error, unparseable payload, task momentarily absent) are
|
||||
// tolerated before aborting. A multi-minute GPU job must not die to one
|
||||
// blip; a genuinely broken upstream still fails within ~5 intervals.
|
||||
const musicPollMaxConsecutiveFailures = 5
|
||||
|
||||
// MusicModel implements musicgen.Provider. The id selects which upstream
|
||||
// llama-swap loads.
|
||||
func (p *Provider) MusicModel(id string, opts ...musicgen.ModelOption) (musicgen.Model, error) {
|
||||
if err := p.requireBaseURL(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = musicgen.ApplyModelOptions(opts)
|
||||
return &musicModel{p: p, id: id}, nil
|
||||
}
|
||||
|
||||
type musicModel struct {
|
||||
|
gitea-actions
commented
🟠 Comment references non-existent docs/en/API.md file maintainability · flagged by 1 model
🪰 Gadfly · advisory 🟠 **Comment references non-existent docs/en/API.md file**
_maintainability · flagged by 1 model_
- **provider/llamaswap/music.go:47** — Comment references `docs/en/API.md` for smoke verification of the `audio_duration` param, but this file does not exist in the repository. The ADR-0021 (`docs/adr/0021-musicgen-interface.md:27`) is the actual design document that specifies this param name. **Fix:** Update the comment to reference the ADR instead, or remove the stale reference. *Verified by grep for `docs/en/API.md` — only one match found, in the comment itself at line 47.*
<sub>🪰 Gadfly · advisory</sub>
|
||||
p *Provider
|
||||
id string
|
||||
}
|
||||
|
||||
// releaseTaskRequest is the ACE-Step POST /release_task shape. audio_duration
|
||||
// is the v1 param name; verify against the upstream ACE-Step-1.5 repo's
|
||||
// docs/en/API.md at smoke time — an
|
||||
// unknown field is ignored upstream, degrading to the default clip length,
|
||||
// never an error.
|
||||
type releaseTaskRequest struct {
|
||||
Prompt string `json:"prompt"`
|
||||
Lyrics string `json:"lyrics,omitempty"`
|
||||
AudioFormat string `json:"audio_format,omitempty"`
|
||||
TaskType string `json:"task_type"`
|
||||
AudioDuration int `json:"audio_duration,omitempty"`
|
||||
InferenceSteps *int `json:"inference_steps,omitempty"`
|
||||
Seed *int64 `json:"seed,omitempty"`
|
||||
}
|
||||
|
||||
// queryItem is one task's poll state. `result` arrives as a JSON-encoded
|
||||
// STRING (the ACE-Step API double-encodes it).
|
||||
|
gitea-actions
commented
🟡 musicFormatMIME duplicates audio.go speechMIME's format→MIME table; two copies will drift maintainability · flagged by 2 models
🪰 Gadfly · advisory 🟡 **musicFormatMIME duplicates audio.go speechMIME's format→MIME table; two copies will drift**
_maintainability · flagged by 2 models_
- **`provider/llamaswap/music.go:69` + `:224` — duplicated audio-MIME resolution that will drift from `speechMIME`.** `audio.go:79` already defines `speechMIME(contentType, format string)` doing exactly the two-step the music path open-codes: try `mimeFromContentType(contentType, "audio/")`, then fall back to a `format`→MIME switch, defaulting to `audio/mpeg`. `music.go` re-implements the fallback switch as `musicFormatMIME` (69–81) and re-inlines the content-type-first check at `fetchResult` (2…
<sub>🪰 Gadfly · advisory</sub>
|
||||
type queryItem struct {
|
||||
TaskID string `json:"task_id"`
|
||||
Status int `json:"status"` // 0 queued/running, 1 succeeded, 2 failed
|
||||
Result string `json:"result"`
|
||||
}
|
||||
|
||||
// musicFormatMIME resolves the clip MIME from the response Content-Type
|
||||
// or the requested format, reusing speechMIME's table (one format->MIME
|
||||
|
gitea-actions
commented
🟡 musicFormatMIME("opus") returns audio/ogg instead of the correct audio/opus (RFC 7845 §9) correctness, error-handling · flagged by 2 models
🪰 Gadfly · advisory 🟡 **musicFormatMIME("opus") returns audio/ogg instead of the correct audio/opus (RFC 7845 §9)**
_correctness, error-handling · flagged by 2 models_
- **`provider/llamaswap/music.go:76` — Wrong MIME type for Opus in the fallback path**
<sub>🪰 Gadfly · advisory</sub>
|
||||
// map for the whole provider). wav32 is ACE-Step-specific: normalize it
|
||||
// to wav before the shared lookup.
|
||||
func musicFormatMIME(contentType, format string) string {
|
||||
format = strings.ToLower(strings.TrimSpace(format))
|
||||
if format == "wav32" {
|
||||
format = "wav"
|
||||
}
|
||||
return speechMIME(contentType, format)
|
||||
}
|
||||
|
gitea-actions
commented
🟠 Missing input length validation on Prompt/Lyrics enables DoS security · flagged by 1 model
🪰 Gadfly · advisory 🟠 **Missing input length validation on Prompt/Lyrics enables DoS**
_security · flagged by 1 model_
- **`provider/llamaswap/music.go:85-92` — Missing input length validation (DoS vector)** The `Generate` function validates that `Prompt` is non-empty and `DurationSeconds >= 0`, but does not bound the length of `Prompt` or `Lyrics`. An attacker could send extremely large strings (MBs) that get JSON-encoded and sent upstream, consuming memory and bandwidth. This is especially relevant since the PR description notes callers "budget the whole job with a ctx deadline" but doesn't mention input size…
<sub>🪰 Gadfly · advisory</sub>
|
||||
|
||||
// Generate implements musicgen.Model.
|
||||
func (m *musicModel) Generate(ctx context.Context, req musicgen.Request, opts ...musicgen.Option) (*musicgen.Result, error) {
|
||||
req = req.Apply(opts...)
|
||||
if strings.TrimSpace(req.Prompt) == "" {
|
||||
return nil, fmt.Errorf("%w: music generation requires a prompt", llm.ErrUnsupported)
|
||||
|
gitea-actions
commented
🟡 WithSteps(0) passes validation and sends inference_steps=0 to backend (invalid) error-handling · flagged by 1 model 🪰 Gadfly · advisory 🟡 **WithSteps(0) passes validation and sends inference_steps=0 to backend (invalid)**
_error-handling · flagged by 1 model_
<sub>🪰 Gadfly · advisory</sub>
|
||||
}
|
||||
if req.DurationSeconds < 0 {
|
||||
return nil, fmt.Errorf("%w: duration must be >= 0, got %d", llm.ErrUnsupported, req.DurationSeconds)
|
||||
}
|
||||
if req.Steps != nil && *req.Steps <= 0 {
|
||||
return nil, fmt.Errorf("%w: inference steps must be > 0, got %d", llm.ErrUnsupported, *req.Steps)
|
||||
}
|
||||
|
||||
taskID, err := m.releaseTask(ctx, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item, err := m.pollResult(ctx, taskID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return m.fetchResult(ctx, req.Format, item)
|
||||
}
|
||||
|
||||
// releaseTask submits the job and returns its task id.
|
||||
func (m *musicModel) releaseTask(ctx context.Context, req musicgen.Request) (string, error) {
|
||||
path, err := upstreamPath(m.id, "/release_task")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
wire := releaseTaskRequest{
|
||||
Prompt: req.Prompt,
|
||||
Lyrics: req.Lyrics,
|
||||
AudioFormat: req.Format,
|
||||
TaskType: "text2music",
|
||||
AudioDuration: req.DurationSeconds,
|
||||
InferenceSteps: req.Steps,
|
||||
Seed: req.Seed,
|
||||
}
|
||||
// Tolerant envelope: {"data": {"task_id": ...}} per the docs, with a
|
||||
// top-level fallback in case the wrapper changes.
|
||||
var resp struct {
|
||||
Data struct {
|
||||
TaskID string `json:"task_id"`
|
||||
} `json:"data"`
|
||||
TaskID string `json:"task_id"`
|
||||
}
|
||||
if err := m.p.doJSON(ctx, http.MethodPost, path, m.id, &wire, &resp); err != nil {
|
||||
return "", err
|
||||
}
|
||||
taskID := resp.Data.TaskID
|
||||
if taskID == "" {
|
||||
taskID = resp.TaskID
|
||||
}
|
||||
if taskID == "" {
|
||||
return "", &llm.APIError{Provider: m.p.name, Model: m.id, Message: "release_task returned no task_id"}
|
||||
}
|
||||
return taskID, nil
|
||||
}
|
||||
|
||||
// pollResult polls query_result until the task succeeds, fails, or ctx
|
||||
// expires. Transient trouble — a transport blip, a momentarily
|
||||
// unparseable payload, the task briefly absent from the response — is
|
||||
|
gitea-actions
commented
🔴 Ticker in pollResult causes back-to-back polls when backend is slow performance · flagged by 1 model
🪰 Gadfly · advisory 🔴 **Ticker in pollResult causes back-to-back polls when backend is slow**
_performance · flagged by 1 model_
* **`provider/llamaswap/music.go:149` — `time.NewTicker` in polling loop creates back-to-back polls when backend is slow.** `pollResult` creates a `time.NewTicker(musicPollInterval)` before the loop and then `select`s on `ticker.C` at the bottom. A Go `Ticker` has a buffer of 1. If `doJSON` + `findQueryItem` ever take longer than 2 s, the buffered tick is consumed immediately on the next `select`, so the loop fires the next request with **zero delay**. On a slow or congested upstream this turns…
<sub>🪰 Gadfly · advisory</sub>
|
||||
// tolerated up to musicPollMaxConsecutiveFailures in a row: a
|
||||
// multi-minute exclusive-GPU job must not die to one flaky poll. Only an
|
||||
// explicit status=2, a run of consecutive failures, or the ctx deadline
|
||||
// aborts.
|
||||
func (m *musicModel) pollResult(ctx context.Context, taskID string) (*queryItem, error) {
|
||||
|
gitea-actions
commented
🟠 No transient-error retry in polling loop — single HTTP blip kills the entire Generate job correctness, error-handling · flagged by 5 models
🪰 Gadfly · advisory 🟠 **No transient-error retry in polling loop — single HTTP blip kills the entire Generate job**
_correctness, error-handling · flagged by 5 models_
- **No transient-error retry in polling loop** (`provider/llamaswap/music.go:154`): If `doJSON` fails with a transient HTTP error (502, 503, connection reset) during any poll iteration, the entire `Generate` call fails immediately — even though the async job may still be running successfully upstream. The caller gets an error with no way to recover the result of an expensive GPU job. This is consistent with the rest of the codebase (no retry infrastructure exists anywhere in `provider/llamaswap`…
<sub>🪰 Gadfly · advisory</sub>
|
||||
path, err := upstreamPath(m.id, "/query_result")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body := map[string]any{"task_id_list": []string{taskID}}
|
||||
ticker := time.NewTicker(musicPollInterval)
|
||||
defer ticker.Stop()
|
||||
failures := 0
|
||||
var lastErr error
|
||||
for {
|
||||
item, pollErr := m.pollOnce(ctx, path, body, taskID)
|
||||
switch {
|
||||
case pollErr != nil:
|
||||
// Ctx expiry is never transient — bail with the deadline error.
|
||||
if ctx.Err() != nil {
|
||||
return nil, fmt.Errorf("llama-swap: music generation: %w", ctx.Err())
|
||||
}
|
||||
failures++
|
||||
lastErr = pollErr
|
||||
if failures >= musicPollMaxConsecutiveFailures {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: fmt.Sprintf("music poll failed %d times in a row: %v", failures, lastErr)}
|
||||
}
|
||||
case item.Status == 1:
|
||||
return item, nil
|
||||
case item.Status == 2:
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: "music generation failed upstream: " + truncateForError([]byte(item.Result))}
|
||||
default:
|
||||
|
gitea-actions
commented
🔴 findQueryItem treats {"data": null} as fatal parse error instead of empty response error-handling · flagged by 1 model
🪰 Gadfly · advisory 🔴 **findQueryItem treats {"data": null} as fatal parse error instead of empty response**
_error-handling · flagged by 1 model_
* **`provider/llamaswap/music.go:183` — `findQueryItem` chokes on `{"data": null}`** When the upstream returns `{"data": null}`, `env.Data` becomes `json.RawMessage("null")` whose length is 4, so the `len(env.Data) > 0` guard passes and `candidates` is set to the literal `"null"`. Both `json.Unmarshal` attempts into `[]queryItem` and `queryItem` then fail, yielding `"unrecognized query_result payload shape"`. `pollResult` wraps this as a hard `APIError`, killing the generation. A null `data` fie…
<sub>🪰 Gadfly · advisory</sub>
|
||||
failures = 0 // healthy queued/running poll
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, fmt.Errorf("llama-swap: music generation: %w", ctx.Err())
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pollOnce performs one query_result round trip and locates the task.
|
||||
func (m *musicModel) pollOnce(ctx context.Context, path string, body any, taskID string) (*queryItem, error) {
|
||||
|
gitea-actions
commented
🟠 findQueryItem accepts empty JSON object {} as valid, causing poll spin until deadline error-handling, maintainability · flagged by 5 models
🪰 Gadfly · advisory 🟠 **findQueryItem accepts empty JSON object {} as valid, causing poll spin until deadline**
_error-handling, maintainability · flagged by 5 models_
- **`findQueryItem` accepts `{}` as a valid single-item response** (`provider/llamaswap/music.go:195`): When the upstream returns `{"data": {}}` (a degenerate case of the `{"data": {...}}` envelope), `json.Unmarshal` into `[]queryItem` fails, then `json.Unmarshal` into a single `queryItem` succeeds with all zero values. The fallback at line 195 matches it (`TaskID == "" && len(items) == 1`), returning a `queryItem` with `Status: 0`. The polling loop treats status 0 as "still queued/running" and…
<sub>🪰 Gadfly · advisory</sub>
|
||||
// Tolerant envelope: items under "data" or a bare array.
|
||||
var raw json.RawMessage
|
||||
if err := m.p.doJSON(ctx, http.MethodPost, path, m.id, body, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return findQueryItem(raw, taskID)
|
||||
}
|
||||
|
||||
// findQueryItem digs the task's entry out of the query_result payload,
|
||||
// tolerating {"data": [...]}, {"data": {...}}, and bare-array envelopes.
|
||||
func findQueryItem(raw json.RawMessage, taskID string) (*queryItem, error) {
|
||||
var env struct {
|
||||
Data json.RawMessage `json:"data"`
|
||||
}
|
||||
candidates := raw
|
||||
if json.Unmarshal(raw, &env) == nil && len(env.Data) > 0 {
|
||||
candidates = env.Data
|
||||
}
|
||||
|
gitea-actions
commented
🔴 Path traversal via untrusted result.File from upstream security · flagged by 3 models
🪰 Gadfly · advisory 🔴 **Path traversal via untrusted result.File from upstream**
_security · flagged by 3 models_
- **`provider/llamaswap/music.go:213` — Path traversal / SSRF via untrusted `result.File`** The `fetchResult` function reads `result.File` from the upstream's double-encoded JSON blob and passes it directly to `upstreamPath()` without any sanitization. The `upstreamPath` function (verified at `upstream.go:20-31`) only validates the `model` argument, not the `rest` path component. A malicious or compromised upstream could return `"file": "../../../etc/passwd"` or an absolute path, resulting in a…
<sub>🪰 Gadfly · advisory</sub>
|
||||
var items []queryItem
|
||||
if err := json.Unmarshal(candidates, &items); err != nil {
|
||||
var one queryItem
|
||||
if err := json.Unmarshal(candidates, &one); err != nil {
|
||||
return nil, fmt.Errorf("unrecognized query_result payload shape")
|
||||
}
|
||||
items = []queryItem{one}
|
||||
}
|
||||
for i := range items {
|
||||
// Single-item responses without a task_id echo are assumed to be
|
||||
// ours — we only ever poll one task; a mismatch surfaces as a
|
||||
// transient miss and is retried by the caller.
|
||||
if items[i].TaskID == taskID || (items[i].TaskID == "" && len(items) == 1) {
|
||||
return &items[i], nil
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("query_result did not include task %s", taskID)
|
||||
}
|
||||
|
||||
// fetchResult downloads the finished clip named by the job's result blob.
|
||||
func (m *musicModel) fetchResult(ctx context.Context, format string, item *queryItem) (*musicgen.Result, error) {
|
||||
var result struct {
|
||||
File string `json:"file"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(item.Result), &result); err != nil || result.File == "" {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: "music result blob missing file URL: " + truncateForError([]byte(item.Result))}
|
||||
}
|
||||
// The file URL is server-relative (e.g. "/v1/audio?path=..."); route it
|
||||
// back through the same upstream. upstreamPath additionally refuses
|
||||
// dot-dot/scheme smuggling in this SERVER-SUPPLIED value.
|
||||
if !strings.HasPrefix(result.File, "/") {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
|
||||
Message: "music result file URL is not server-relative: " + truncateForError([]byte(result.File))}
|
||||
}
|
||||
path, err := upstreamPath(m.id, result.File)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw, contentType, err := m.p.doRaw(ctx, http.MethodGet, path, m.id, "", nil, maxResponseBytes)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(raw) == 0 {
|
||||
return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "music response contained no audio"}
|
||||
}
|
||||
mimeType := musicFormatMIME(contentType, format)
|
||||
return &musicgen.Result{
|
||||
Audio: musicgen.Audio{Data: raw, MIME: mimeType},
|
||||
Raw: json.RawMessage(item.Result),
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
package llamaswap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/embeddings"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
||||
"gitea.stevedudenhoeffer.com/steve/majordomo/musicgen"
|
||||
)
|
||||
|
||||
// aceStepStub emulates the ACE-Step job API: one queued poll, then success.
|
||||
func aceStepStub(t *testing.T, mp3 []byte) (*httptest.Server, *atomic.Int32) {
|
||||
t.Helper()
|
||||
var polls atomic.Int32
|
||||
var gotRelease map[string]any
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/upstream/musicgen-acestep/release_task":
|
||||
_ = json.NewDecoder(r.Body).Decode(&gotRelease)
|
||||
if gotRelease["task_type"] != "text2music" {
|
||||
t.Errorf("task_type = %v", gotRelease["task_type"])
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"data": {"task_id": "t-1", "status": "queued"}}`))
|
||||
case "/upstream/musicgen-acestep/query_result":
|
||||
n := polls.Add(1)
|
||||
if n == 1 {
|
||||
_, _ = w.Write([]byte(`{"data": [{"task_id": "t-1", "status": 0, "result": ""}]}`))
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"data": [{"task_id": "t-1", "status": 1,
|
||||
"result": "{\"file\": \"/v1/audio?path=out.mp3\", \"metas\": {\"bpm\": 120}}"}]}`))
|
||||
case "/upstream/musicgen-acestep/v1/audio":
|
||||
if got := r.URL.Query().Get("path"); got != "out.mp3" {
|
||||
t.Errorf("audio path = %q", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "audio/mpeg")
|
||||
_, _ = w.Write(mp3)
|
||||
default:
|
||||
t.Errorf("unexpected path %q", r.URL.Path)
|
||||
w.WriteHeader(404)
|
||||
}
|
||||
}))
|
||||
return srv, &polls
|
||||
}
|
||||
|
||||
func TestMusicGenerate(t *testing.T) {
|
||||
old := musicPollInterval
|
||||
musicPollInterval = 5 * time.Millisecond
|
||||
defer func() { musicPollInterval = old }()
|
||||
|
||||
mp3 := []byte("ID3fakeaudio")
|
||||
srv, polls := aceStepStub(t, mp3)
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
mm, err := p.MusicModel("musicgen-acestep")
|
||||
if err != nil {
|
||||
t.Fatalf("MusicModel: %v", err)
|
||||
}
|
||||
// Shrink the poll interval indirectly by bounding the whole call: the
|
||||
// stub succeeds on poll #2, so a generous deadline still finishes fast.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
res, err := mm.Generate(ctx,
|
||||
musicgen.Request{Prompt: "chiptune anthem about mortbux"},
|
||||
musicgen.WithLyrics("mort mort mort"), musicgen.WithDuration(30))
|
||||
if err != nil {
|
||||
t.Fatalf("Generate: %v", err)
|
||||
}
|
||||
if polls.Load() < 2 {
|
||||
t.Errorf("polls = %d, want >= 2 (queued then done)", polls.Load())
|
||||
}
|
||||
if res.Audio.MIME != "audio/mpeg" || string(res.Audio.Data) != string(mp3) {
|
||||
t.Fatalf("audio = %q/%d bytes", res.Audio.MIME, len(res.Audio.Data))
|
||||
}
|
||||
}
|
||||
|
||||
func TestMusicGenerateUpstreamFailure(t *testing.T) {
|
||||
old := musicPollInterval
|
||||
musicPollInterval = 5 * time.Millisecond
|
||||
defer func() { musicPollInterval = old }()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/upstream/musicgen-acestep/release_task":
|
||||
_, _ = w.Write([]byte(`{"data": {"task_id": "t-2"}}`))
|
||||
case "/upstream/musicgen-acestep/query_result":
|
||||
_, _ = w.Write([]byte(`{"data": [{"task_id": "t-2", "status": 2, "result": "{\"error\": \"OOM\"}"}]}`))
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
mm, _ := p.MusicModel("musicgen-acestep")
|
||||
_, err := mm.Generate(context.Background(), musicgen.Request{Prompt: "p"})
|
||||
var apiErr *llm.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %v, want APIError for failed job", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMusicGenerateRejectsBadArgs(t *testing.T) {
|
||||
p := New(WithBaseURL("http://unused"))
|
||||
mm, _ := p.MusicModel("musicgen-acestep")
|
||||
if _, err := mm.Generate(context.Background(), musicgen.Request{}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("no prompt: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmbed(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/embeddings" {
|
||||
t.Errorf("path = %q", r.URL.Path)
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&gotBody)
|
||||
// Deliberately out of order: the client must sort by index.
|
||||
_, _ = w.Write([]byte(`{"object":"list","data":[
|
||||
{"object":"embedding","index":1,"embedding":[0.3,0.4]},
|
||||
{"object":"embedding","index":0,"embedding":[0.1,0.2]}
|
||||
]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
em, err := p.EmbedModel("embed-qwen3-0.6b")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedModel: %v", err)
|
||||
}
|
||||
res, err := em.Embed(context.Background(), embeddings.EmbedRequest{Inputs: []string{"a", "b"}})
|
||||
if err != nil {
|
||||
t.Fatalf("Embed: %v", err)
|
||||
}
|
||||
if gotBody["model"] != "embed-qwen3-0.6b" {
|
||||
t.Errorf("model = %v", gotBody["model"])
|
||||
}
|
||||
if len(res.Vectors) != 2 || res.Vectors[0][0] != 0.1 || res.Vectors[1][0] != 0.3 {
|
||||
t.Fatalf("vectors = %+v (index ordering broken?)", res.Vectors)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmbedCountMismatchIsError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1]}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
em, _ := p.EmbedModel("embed-qwen3-0.6b")
|
||||
_, err := em.Embed(context.Background(), embeddings.EmbedRequest{Inputs: []string{"a", "b"}})
|
||||
var apiErr *llm.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %v, want APIError for count mismatch", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmbedRejectsEmptyInputs(t *testing.T) {
|
||||
p := New(WithBaseURL("http://unused"))
|
||||
em, _ := p.EmbedModel("embed-qwen3-0.6b")
|
||||
if _, err := em.Embed(context.Background(), embeddings.EmbedRequest{}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("no inputs: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
if _, err := em.Embed(context.Background(), embeddings.EmbedRequest{Inputs: []string{"a", " "}}); !errors.Is(err, llm.ErrUnsupported) {
|
||||
t.Errorf("blank input: err = %v, want ErrUnsupported", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRerank(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/rerank" {
|
||||
t.Errorf("path = %q", r.URL.Path)
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&gotBody)
|
||||
// Out of score order: the client must sort descending.
|
||||
_, _ = w.Write([]byte(`{"results":[
|
||||
{"index":0,"relevance_score":0.11},
|
||||
{"index":2,"relevance_score":0.93},
|
||||
{"index":1,"relevance_score":0.42}
|
||||
]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
rm, err := p.RerankModel("rerank-bge-v2-m3")
|
||||
if err != nil {
|
||||
t.Fatalf("RerankModel: %v", err)
|
||||
}
|
||||
res, err := rm.Rerank(context.Background(),
|
||||
embeddings.RerankRequest{Query: "what is a panda?", Documents: []string{"a", "b", "c"}},
|
||||
embeddings.WithTopN(3))
|
||||
if err != nil {
|
||||
t.Fatalf("Rerank: %v", err)
|
||||
}
|
||||
if gotBody["top_n"] != float64(3) || gotBody["query"] != "what is a panda?" {
|
||||
t.Errorf("top_n/query = %v/%v", gotBody["top_n"], gotBody["query"])
|
||||
}
|
||||
if len(res.Results) != 3 || res.Results[0].Index != 2 || res.Results[2].Index != 0 {
|
||||
t.Fatalf("results = %+v (descending sort broken?)", res.Results)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRerankRejectsOutOfRangeIndex(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"results":[{"index":7,"relevance_score":0.9}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
rm, _ := p.RerankModel("rerank-bge-v2-m3")
|
||||
_, err := rm.Rerank(context.Background(),
|
||||
embeddings.RerankRequest{Query: "q", Documents: []string{"a"}})
|
||||
var apiErr *llm.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("err = %v, want APIError for out-of-range index", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstructedQuery(t *testing.T) {
|
||||
got := embeddings.InstructedQuery("", "how tall is everest")
|
||||
want := "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: how tall is everest"
|
||||
if got != want {
|
||||
t.Errorf("InstructedQuery = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMusicGenerateSurvivesTransientPollFailures(t *testing.T) {
|
||||
old := musicPollInterval
|
||||
musicPollInterval = 5 * time.Millisecond
|
||||
defer func() { musicPollInterval = old }()
|
||||
|
||||
mp3 := []byte("ID3fakeaudio")
|
||||
var polls atomic.Int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/upstream/musicgen-acestep/release_task":
|
||||
_, _ = w.Write([]byte(`{"data": {"task_id": "t-3"}}`))
|
||||
case "/upstream/musicgen-acestep/query_result":
|
||||
switch polls.Add(1) {
|
||||
case 1:
|
||||
w.WriteHeader(http.StatusBadGateway) // transient transport blip
|
||||
case 2:
|
||||
_, _ = w.Write([]byte(`{"data": []}`)) // task momentarily absent
|
||||
default:
|
||||
_, _ = w.Write([]byte(`{"data": [{"task_id": "t-3", "status": 1,
|
||||
"result": "{\"file\": \"/v1/audio?path=out.mp3\"}"}]}`))
|
||||
}
|
||||
case "/upstream/musicgen-acestep/v1/audio":
|
||||
w.Header().Set("Content-Type", "audio/mpeg")
|
||||
_, _ = w.Write(mp3)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
mm, _ := p.MusicModel("musicgen-acestep")
|
||||
res, err := mm.Generate(context.Background(), musicgen.Request{Prompt: "p"})
|
||||
if err != nil {
|
||||
t.Fatalf("Generate should survive 2 transient failures: %v", err)
|
||||
}
|
||||
if len(res.Audio.Data) == 0 {
|
||||
t.Fatal("no audio")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMusicGenerateRejectsHostileFileURL(t *testing.T) {
|
||||
old := musicPollInterval
|
||||
musicPollInterval = 5 * time.Millisecond
|
||||
defer func() { musicPollInterval = old }()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/upstream/musicgen-acestep/release_task":
|
||||
_, _ = w.Write([]byte(`{"data": {"task_id": "t-4"}}`))
|
||||
case "/upstream/musicgen-acestep/query_result":
|
||||
_, _ = w.Write([]byte(`{"data": [{"task_id": "t-4", "status": 1,
|
||||
"result": "{\"file\": \"/v1/audio?path=../../api/models/unload\"}"}]}`))
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
||||
mm, _ := p.MusicModel("musicgen-acestep")
|
||||
_, err := mm.Generate(context.Background(), musicgen.Request{Prompt: "p"})
|
||||
if err == nil {
|
||||
t.Fatal("dot-dot result file URL accepted")
|
||||
}
|
||||
}
|
||||
🟠 Missing input length validation on embedding Inputs enables DoS
security · flagged by 1 model
provider/llamaswap/embed.go:42-50— Missing input length validation for embeddings (DoS vector) TheEmbedfunction rejects empty inputs and blank strings but does not limit the length of each input string or the total batch size. Large inputs could cause memory exhaustion upstream or in the HTTP client. Fix: Add per-input and total batch size limits consistent with the embedding model's actual capabilities. Verified: Readembed.go:42-50; confirmed only emptiness checks exist.🪰 Gadfly · advisory