majordomo/provider/llamaswap/music_embed_test.go

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")
	}
}