diff --git a/docs/adr/0021-musicgen-interface.md b/docs/adr/0021-musicgen-interface.md new file mode 100644 index 0000000..dc82d83 --- /dev/null +++ b/docs/adr/0021-musicgen-interface.md @@ -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//` + 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. diff --git a/docs/adr/0022-embeddings-rerank-interface.md b/docs/adr/0022-embeddings-rerank-interface.md new file mode 100644 index 0000000..c5e7178 --- /dev/null +++ b/docs/adr/0022-embeddings-rerank-interface.md @@ -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). diff --git a/embeddings/embeddings.go b/embeddings/embeddings.go new file mode 100644 index 0000000..b3672d8 --- /dev/null +++ b/embeddings/embeddings.go @@ -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) +} diff --git a/musicgen/musicgen.go b/musicgen/musicgen.go new file mode 100644 index 0000000..e009f3c --- /dev/null +++ b/musicgen/musicgen.go @@ -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) +} diff --git a/provider/llamaswap/embed.go b/provider/llamaswap/embed.go new file mode 100644 index 0000000..2705955 --- /dev/null +++ b/provider/llamaswap/embed.go @@ -0,0 +1,146 @@ +// 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)) + 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)} + } + 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 +} + +// 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 +} diff --git a/provider/llamaswap/music.go b/provider/llamaswap/music.go new file mode 100644 index 0000000..a3631cc --- /dev/null +++ b/provider/llamaswap/music.go @@ -0,0 +1,232 @@ +// 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//release_task {prompt, lyrics, ...} -> {task_id} +// POST /upstream//query_result {task_id_list: [...]} -> status+result +// GET /upstream// -> 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. +const musicPollInterval = 2 * time.Second + +// 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 { + p *Provider + id string +} + +// releaseTaskRequest is the ACE-Step POST /release_task shape. audio_duration +// is the v1 param name; verify against 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). +type queryItem struct { + TaskID string `json:"task_id"` + Status int `json:"status"` // 0 queued/running, 1 succeeded, 2 failed + Result string `json:"result"` +} + +// musicFormatMIME maps ACE-Step's audio_format values to MIME types. +func musicFormatMIME(format string) string { + switch strings.ToLower(strings.TrimSpace(format)) { + case "wav", "wav32": + return "audio/wav" + case "flac": + return "audio/flac" + case "opus": + return "audio/ogg" + case "aac": + return "audio/aac" + default: // "", "mp3" + return "audio/mpeg" + } +} + +// 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) + } + if req.DurationSeconds < 0 { + return nil, fmt.Errorf("%w: duration must be >= 0, got %d", llm.ErrUnsupported, req.DurationSeconds) + } + + 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. +func (m *musicModel) pollResult(ctx context.Context, taskID string) (*queryItem, error) { + 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() + for { + // 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 + } + item, err := findQueryItem(raw, taskID) + if err != nil { + return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: err.Error()} + } + switch item.Status { + case 1: + return item, nil + case 2: + return nil, &llm.APIError{Provider: m.p.name, Model: m.id, + Message: "music generation failed upstream: " + truncateForError([]byte(item.Result))} + } + select { + case <-ctx.Done(): + return nil, fmt.Errorf("llama-swap: music generation: %w", ctx.Err()) + case <-ticker.C: + } + } +} + +// 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 + } + 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 { + 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. + 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 := mimeFromContentType(contentType, "audio/") + if mimeType == "" { + mimeType = musicFormatMIME(format) + } + return &musicgen.Result{ + Audio: musicgen.Audio{Data: raw, MIME: mimeType}, + Raw: json.RawMessage(item.Result), + }, nil +} diff --git a/provider/llamaswap/music_embed_test.go b/provider/llamaswap/music_embed_test.go new file mode 100644 index 0000000..ea3f7a5 --- /dev/null +++ b/provider/llamaswap/music_embed_test.go @@ -0,0 +1,224 @@ +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) { + 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) { + 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) + } +}