feat: musicgen + embeddings/rerank surfaces (ADR-0021, ADR-0022)
- NEW musicgen leaf package: blocking Generate over ACE-Step's async job queue (release_task -> poll query_result -> fetch file, all via /upstream); tolerant envelope parsing, double-encoded result handled - NEW embeddings leaf package: EmbedModel + RerankModel as separate mints (two server instances on the host, llama.cpp #20085); InstructedQuery helper for Qwen3-style query/document asymmetry - provider/llamaswap: /v1/embeddings + /v1/rerank clients with strict validation (index-ordered vectors, count mismatch and out-of-range index are hard errors; rerank sorted descending, minimal parser) Co-Authored-By: Claude Fable 5 <[email protected]>
This commit is contained in:
Claude Fable 5
@@ -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/<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.
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user