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 <noreply@anthropic.com>

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

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

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