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>
This commit is contained in:
@@ -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/<model>/`
|
||||||
|
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.
|
||||||
@@ -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).
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user