majordomo/provider/llamaswap/embed.go

// 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)}
		}
		if vectors[d.Index] != nil {
			return nil, &llm.APIError{Provider: m.p.name, Model: m.id,
				Message: fmt.Sprintf("embeddings response repeats index %d", 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
}