majordomo/embeddings/embeddings.go

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