majordomo/provider/llamaswap/music.go

// 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
}