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,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
|
||||
}
|
||||
Reference in New Issue
Block a user