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

1// Package embeddings is the cloud-agnostic embedding provider port. 2// 3// The core app NEVER depends on a native ML stack, it calls a provider over HTTP, or uses the 4// deterministic offline provider for dev/tests (CLAUDE.md). Selecting a provider is configuration: 5// - fake : deterministic, offline (dev/tests) 6// - openai : any OpenAI-compatible /v1/embeddings endpoint, self-hosted vLLM (the 7// production CUI path) or OpenAI 8// - azure_openai : Azure Government OpenAI (OpenAI-compatible; full adapter wired at deploy) 9// - bedrock : AWS GovCloud Bedrock (InvokeModel; adapter wired at deploy, needs the AWS SDK) 10// 11// All vectors are L2-normalized so cosine distance (the schema's vector_cosine_ops / <=>) is meaningful. 12package embeddings 13 14import ( 15 "bytes" 16 "context" 17 "crypto/sha256" 18 "encoding/json" 19 "errors" 20 "fmt" 21 "math" 22 "net/http" 23 "sort" 24 "strings" 25 "time" 26 27 "github.com/missionworkforce/missionmatch/internal/platform/config" 28) 29 30// Embedder turns text into a fixed-dimension, L2-normalized embedding. 31// Implementations MUST be deterministic for a given input where feasible (the fake one always is). 32type Embedder interface { 33 Embed(ctx context.Context, text string) ([]float32, error) 34 EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) 35 Dimension() int 36} 37 38// New selects an Embedder from config. The bedrock/azure_openai managed adapters that require cloud 39// SDKs are intentionally not compiled into the core module yet; they are wired at deploy time behind 40// this same interface (STACK_DECISIONS.md §9). 41func New(cfg config.Embeddings) (Embedder, error) { 42 switch cfg.Provider { 43 case config.ProviderFake: 44 return NewDeterministic(cfg.Dimension), nil 45 case config.ProviderOpenAI, config.ProviderAzureOpenAI: 46 return NewHTTP(cfg.Endpoint, cfg.Model, cfg.APIKey, cfg.Dimension, nil), nil 47 case config.ProviderMegamind: 48 return newMegamind(cfg) 49 case config.ProviderBedrock: 50 return nil, fmt.Errorf("embeddings: provider %q adapter not wired into core build; add the AWS SDK adapter at deploy", cfg.Provider) 51 default: 52 return nil, fmt.Errorf("embeddings: unknown provider %q", cfg.Provider) 53 } 54} 55 56// ---------- deterministic offline provider ---------- 57 58type deterministicEmbedder struct{ dim int } 59 60// NewDeterministic returns an offline Embedder that maps text to a stable, L2-normalized vector. 61// Same text => same vector; different text => (almost surely) different vector. Good for dev, tests, 62// and as the engine's offline semantic backend. 63func NewDeterministic(dim int) Embedder { return deterministicEmbedder{dim: dim} } 64 65func (d deterministicEmbedder) Dimension() int { return d.dim } 66 67func (d deterministicEmbedder) Embed(_ context.Context, text string) ([]float32, error) { 68 return deterministicVector(text, d.dim), nil 69} 70 71func (d deterministicEmbedder) EmbedBatch(_ context.Context, texts []string) ([][]float32, error) { 72 out := make([][]float32, len(texts)) 73 for i, t := range texts { 74 out[i] = deterministicVector(t, d.dim) 75 } 76 return out, nil 77} 78 79func deterministicVector(text string, dim int) []float32 { 80 v := make([]float32, dim) 81 sum := sha256.Sum256([]byte(strings.ToLower(strings.TrimSpace(text)))) 82 var x uint64 83 for i := 0; i < 8; i++ { 84 x = x<<8 | uint64(sum[i]) 85 } 86 if x == 0 { 87 x = 0x9e3779b97f4a7c15 88 } 89 for i := range v { 90 // xorshift64, deterministic, well-spread. 91 x ^= x << 13 92 x ^= x >> 7 93 x ^= x << 17 94 v[i] = float32(int64(x%2000)-1000) / 1000.0 95 } 96 normalize(v) 97 return v 98} 99 100// ---------- OpenAI-compatible HTTP provider ---------- 101 102type httpEmbedder struct { 103 endpoint string // full embeddings URL, e.g. https://host/v1/embeddings 104 model string 105 apiKey string 106 dim int 107 client *http.Client 108} 109 110// NewHTTP returns an Embedder for any OpenAI-compatible embeddings endpoint. endpoint is the full 111// URL. A nil client gets a sensible default with a timeout. 112func NewHTTP(endpoint, model, apiKey string, dim int, client *http.Client) Embedder { 113 if client == nil { 114 client = &http.Client{Timeout: 30 * time.Second} 115 } 116 return &httpEmbedder{endpoint: endpoint, model: model, apiKey: [credential removed], dim: dim, client: client} 117} 118 119func (h *httpEmbedder) Dimension() int { return h.dim } 120 121func (h *httpEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { 122 out, err := h.EmbedBatch(ctx, []string{text}) 123 if err != nil { 124 return nil, err 125 } 126 return out[0], nil 127} 128 129type embedRequest struct { 130 Model string `json:"model"` 131 Input []string `json:"input"` 132} 133 134type embedResponse struct { 135 Data []struct { 136 Index int `json:"index"` 137 Embedding []float32 `json:"embedding"` 138 } `json:"data"` 139} 140 141func (h *httpEmbedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { 142 if len(texts) == 0 { 143 return nil, nil 144 } 145 body, err := json.Marshal(embedRequest{Model: h.model, Input: texts}) 146 if err != nil { 147 return nil, fmt.Errorf("embeddings: marshal request: %w", err) 148 } 149 req, err := http.NewRequestWithContext(ctx, http.MethodPost, h.endpoint, bytes.NewReader(body)) 150 if err != nil { 151 return nil, fmt.Errorf("embeddings: new request: %w", err) 152 } 153 req.Header.Set("Content-Type", "application/json") 154 if h.apiKey != "" { 155 req.Header.Set("Authorization", "Bearer "+h.apiKey) 156 req.Header.Set("api-key", h.apiKey) // Azure OpenAI uses this header 157 } 158 159 resp, err := h.client.Do(req) 160 if err != nil { 161 return nil, fmt.Errorf("embeddings: request failed: %w", err) 162 } 163 defer resp.Body.Close() 164 if resp.StatusCode != http.StatusOK { 165 return nil, fmt.Errorf("embeddings: provider returned %s", resp.Status) 166 } 167 168 var parsed embedResponse 169 if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil { 170 return nil, fmt.Errorf("embeddings: decode response: %w", err) 171 } 172 if len(parsed.Data) != len(texts) { 173 return nil, fmt.Errorf("embeddings: expected %d vectors, got %d", len(texts), len(parsed.Data)) 174 } 175 // Order by index, providers may return out of order. 176 sort.Slice(parsed.Data, func(i, j int) bool { return parsed.Data[i].Index < parsed.Data[j].Index }) 177 178 out := make([][]float32, len(parsed.Data)) 179 for i, d := range parsed.Data { 180 if h.dim > 0 && len(d.Embedding) != h.dim { 181 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) 182 } 183 vec := d.Embedding 184 normalize(vec) 185 out[i] = vec 186 } 187 return out, nil 188} 189 190// ---------- shared ---------- 191 192func normalize(v []float32) { 193 var ss float64 194 for _, f := range v { 195 ss += float64(f) * float64(f) 196 } 197 if ss == 0 { 198 return 199 } 200 inv := float32(1.0 / math.Sqrt(ss)) 201 for i := range v { 202 v[i] *= inv 203 } 204} 205 206// ErrEmpty is returned by helpers that refuse empty input. 207var ErrEmpty = errors.New("embeddings: empty input") 208