Janady13/missionmatch / internal/platform/embeddings/embeddings.go

Source identity

Field	Recorded value
Source type	GitHub source
Original record	Janady13/missionmatch / internal/platform/embeddings/embeddings.go
Date	Not specified in source
Repository	Janady13/missionmatch
Revision	720d47fae476b0d38477ae4401697d372b1ea2bc
Branch	feature/opportunity-to-outcome
Path or archive member	internal/platform/embeddings/embeddings.go
Original SHA-256	c9f896c0996e731cf81aa3b1adf7d64b74153cf5268cc51b20b0da2237ca24e9
Git blob SHA	d2a04756e1bb1d7ff1064027a7284f593ff55631
Reading-copy SHA-256	1c5142d466cec965348eabbc8623e4f14cbfa49206fb529322aa8f32f5cdaed6
Original source ↗
Recovered source transcript

Download complete source text ↓268 lines · 8.9 KB
Credential-bearing and unrelated personal details were removed from this reading copy. The source identity and original checksum remain recorded.
// Package embeddings is the cloud-agnostic embedding provider port.
//
// The core app NEVER depends on a native ML stack, it calls a provider over HTTP, or uses the
// deterministic offline provider for dev/tests (CLAUDE.md). Selecting a provider is configuration:
//   - fake          : deterministic, offline (dev/tests)
//   - openai        : any OpenAI-compatible /v1/embeddings endpoint, self-hosted vLLM (the
//     production CUI path) or OpenAI
//   - azure_openai  : Azure Government OpenAI (OpenAI-compatible; full adapter wired at deploy)
//   - bedrock       : AWS GovCloud Bedrock (InvokeModel; adapter wired at deploy, needs the AWS SDK)
//
// All vectors are L2-normalized so cosine distance (the schema's vector_cosine_ops / <=>) is meaningful.
package embeddings

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"net/http"
	"sort"
	"strings"
	"time"

	"github.com/missionworkforce/missionmatch/internal/platform/config"
)

// SkillEmbedText is the canonical text to embed for one skill: the name, with the category in
// parentheses when present. The backfill (the writer of skills.embedding) and the boot vector model
// check (the verifier that re embeds a stored skill to confirm the model) MUST embed the SAME text, or
// the check would see a false mismatch on every categorized skill. This function is that single source.
func SkillEmbedText(name, category string) string {
	if category != "" {
		return name + " (" + category + ")"
	}
	return name
}

// OpportunityEmbedText is the canonical text to embed for one federal opportunity record. Like
// SkillEmbedText it is the SINGLE source both the backfill and any future verifier use, so a re embed of a
// stored row reproduces the same vector rather than reporting a false model mismatch.
//
// The composition is deliberate. Title carries the work itself and leads. The buying organization, the
// codes, and the place of performance are what a capture question actually asks about ("who buys data
// engineering in San Antonio", "what 8(a) work is coming in 541512"), so each is included as a labeled
// clause rather than a bare token, which keeps the sentence meaningful to a sentence embedder instead of
// reading as a code dump. Empty fields are skipped so a sparse forecast row does not embed a string of
// empty labels that would pull unrelated sparse rows together in the vector space.
func OpportunityEmbedText(title, agency, subAgency, naics, psc, setAside, city, state, incumbent, kind string) string {
	parts := make([]string, 0, 8)
	if title != "" {
		parts = append(parts, title)
	}
	if agency != "" {
		org := agency
		if subAgency != "" && subAgency != agency {
			org += ", " + subAgency
		}
		parts = append(parts, "Agency: "+org)
	}
	if naics != "" {
		parts = append(parts, "NAICS "+naics)
	}
	if psc != "" {
		parts = append(parts, "PSC "+psc)
	}
	if setAside != "" {
		parts = append(parts, "Set aside: "+setAside)
	}
	if city != "" || state != "" {
		place := city
		if place != "" && state != "" {
			place += ", " + state
		} else if place == "" {
			place = state
		}
		parts = append(parts, "Place of performance: "+place)
	}
	if incumbent != "" {
		parts = append(parts, "Incumbent: "+incumbent)
	}
	if kind != "" {
		parts = append(parts, "Record type: "+strings.ReplaceAll(kind, "_", " "))
	}
	return strings.Join(parts, ". ")
}

// Embedder turns text into a fixed-dimension, L2-normalized embedding.
// Implementations MUST be deterministic for a given input where feasible (the fake one always is).
type Embedder interface {
	Embed(ctx context.Context, text string) ([]float32, error)
	EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)
	Dimension() int
}

// New selects an Embedder from config. The bedrock/azure_openai managed adapters that require cloud
// SDKs are intentionally not compiled into the core module yet; they are wired at deploy time behind
// this same interface (STACK_DECISIONS.md §9).
func New(cfg config.Embeddings) (Embedder, error) {
	switch cfg.Provider {
	case config.ProviderFake:
		return NewDeterministic(cfg.Dimension), nil
	case config.ProviderOpenAI, config.ProviderAzureOpenAI:
		return NewHTTP(cfg.Endpoint, cfg.Model, cfg.APIKey, cfg.Dimension, nil), nil
	case config.ProviderMegamind:
		return newMegamind(cfg)
	case config.ProviderBedrock:
		return nil, fmt.Errorf("embeddings: provider %q adapter not wired into core build; add the AWS SDK adapter at deploy", cfg.Provider)
	default:
		return nil, fmt.Errorf("embeddings: unknown provider %q", cfg.Provider)
	}
}

// ---------- deterministic offline provider ----------

type deterministicEmbedder struct{ dim int }

// NewDeterministic returns an offline Embedder that maps text to a stable, L2-normalized vector.
// Same text => same vector; different text => (almost surely) different vector. Good for dev, tests,
// and as the engine's offline semantic backend.
func NewDeterministic(dim int) Embedder { return deterministicEmbedder{dim: dim} }

func (d deterministicEmbedder) Dimension() int { return d.dim }

func (d deterministicEmbedder) Embed(_ context.Context, text string) ([]float32, error) {
	return deterministicVector(text, d.dim), nil
}

func (d deterministicEmbedder) EmbedBatch(_ context.Context, texts []string) ([][]float32, error) {
	out := make([][]float32, len(texts))
	for i, t := range texts {
		out[i] = deterministicVector(t, d.dim)
	}
	return out, nil
}

func deterministicVector(text string, dim int) []float32 {
	v := make([]float32, dim)
	sum := sha256.Sum256([]byte(strings.ToLower(strings.TrimSpace(text))))
	var x uint64
	for i := 0; i < 8; i++ {
		x = x<<8 | uint64(sum[i])
	}
	if x == 0 {
		x = 0x9e3779b97f4a7c15
	}
	for i := range v {
		// xorshift64, deterministic, well-spread.
		x ^= x << 13
		x ^= x >> 7
		x ^= x << 17
		v[i] = float32(int64(x%2000)-1000) / 1000.0
	}
	normalize(v)
	return v
}

// ---------- OpenAI-compatible HTTP provider ----------

type httpEmbedder struct {
	endpoint string // full embeddings URL, e.g. https://host/v1/embeddings
	model    string
	apiKey   string
	dim      int
	client   *http.Client
}

// NewHTTP returns an Embedder for any OpenAI-compatible embeddings endpoint. endpoint is the full
// URL. A nil client gets a sensible default with a timeout.
func NewHTTP(endpoint, model, apiKey string, dim int, client *http.Client) Embedder {
	if client == nil {
		client = &http.Client{Timeout: 30 * time.Second}
	}
	return &httpEmbedder{endpoint: endpoint, model: model, apiKey: [credential removed], dim: dim, client: client}
}

func (h *httpEmbedder) Dimension() int { return h.dim }

func (h *httpEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
	out, err := h.EmbedBatch(ctx, []string{text})
	if err != nil {
		return nil, err
	}
	return out[0], nil
}

type embedRequest struct {
	Model string   `json:"model"`
	Input []string `json:"input"`
}

type embedResponse struct {
	Data []struct {
		Index     int       `json:"index"`
		Embedding []float32 `json:"embedding"`
	} `json:"data"`
}

func (h *httpEmbedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) {
	if len(texts) == 0 {
		return nil, nil
	}
	body, err := json.Marshal(embedRequest{Model: h.model, Input: texts})
	if err != nil {
		return nil, fmt.Errorf("embeddings: marshal request: %w", err)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, h.endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("embeddings: new request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	if h.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+h.apiKey)
		req.Header.Set("api-key", h.apiKey) // Azure OpenAI uses this header
	}

	resp, err := h.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("embeddings: request failed: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("embeddings: provider returned %s", resp.Status)
	}

	var parsed embedResponse
	if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
		return nil, fmt.Errorf("embeddings: decode response: %w", err)
	}
	if len(parsed.Data) != len(texts) {
		return nil, fmt.Errorf("embeddings: expected %d vectors, got %d", len(texts), len(parsed.Data))
	}
	// Order by index, providers may return out of order.
	sort.Slice(parsed.Data, func(i, j int) bool { return parsed.Data[i].Index < parsed.Data[j].Index })

	out := make([][]float32, len(parsed.Data))
	for i, d := range parsed.Data {
		if h.dim > 0 && len(d.Embedding) != h.dim {
			return nil, fmt.Errorf("embeddings: vector %d has dim %d, want %d (check MM_EMBED_DIMENSION vs the model)", i, len(d.Embedding), h.dim)
		}
		vec := d.Embedding
		normalize(vec)
		out[i] = vec
	}
	return out, nil
}

// ---------- shared ----------

func normalize(v []float32) {
	var ss float64
	for _, f := range v {
		ss += float64(f) * float64(f)
	}
	if ss == 0 {
		return
	}
	inv := float32(1.0 / math.Sqrt(ss))
	for i := range v {
		v[i] *= inv
	}
}

// ErrEmpty is returned by helpers that refuse empty input.
var ErrEmpty = errors.New("embeddings: empty input")