Janady13/missionmatch / internal/platform/inference/client.go

1// Package inference is the typed Go client for the shared, stateless MEGAMIND inference plane. It is how 2// the FIPS API binary reaches embeddings, reranking, generation, and faithfulness checking WITHOUT 3// hosting any ML in process: the heavy models run out of process on the inference plane (Apple Silicon, 4// Metal), and this client speaks plain HTTP/JSON to them with timeouts, bounded retries, and a circuit 5// breaker so a sick plane degrades the product gracefully instead of hanging it. 6// 7// The client sends no client identifier to the plane on the inference paths (it sends text in, gets 8// vectors or text out), and the one teaching path is gated by the scrub boundary, so MEGAMIND learns 9// only the craft of the work and persists nothing about any tenant. 10package inference 11 12import ( 13 "bytes" 14 "context" 15 "encoding/json" 16 "errors" 17 "fmt" 18 "io" 19 "net/http" 20 "time" 21 22 "github.com/missionworkforce/missionmatch/internal/platform/scrub" 23) 24 25// Config configures the client. The model names and dimension come from server side config so a model 26// swap (for example all-minilm to bge-small-en-v1.5, both 384) is a config change, not a code change. 27type Config struct { 28 BaseURL string // MEGAMIND gateway base, e.g. http://100.94.8.94:8090 29 EmbedModel string // 384 dim embedder (all-minilm or bge-small-en-v1.5) 30 RerankModel string // cross encoder (bge-reranker-v2-m3) or cosine fallback 31 GenerateModel string // local instruct model (e.g. llama3.1:8b) 32 Dimension int // expected embedding width; a mismatch is a hard error (column is vector(384)) 33 Timeout time.Duration // per request 34 MaxRetries int // additional attempts on a transient failure 35 BreakerTrips int // consecutive failures before the breaker opens 36 BreakerCool time.Duration // how long the breaker stays open 37 TeachingEnabled bool // when false, Teach validates the boundary but does not transmit 38} 39 40// Client talks to the MEGAMIND plane. It is safe for concurrent use. 41type Client struct { 42 cfg Config 43 http *http.Client 44 cb *breaker 45} 46 47// New builds a client. BaseURL is required; sensible defaults fill the rest. 48func New(cfg Config) (*Client, error) { 49 if cfg.BaseURL == "" { 50 return nil, errors.New("inference: BaseURL is required") 51 } 52 if cfg.Dimension <= 0 { 53 cfg.Dimension = 384 54 } 55 if cfg.Timeout <= 0 { 56 cfg.Timeout = 20 * time.Second 57 } 58 if cfg.BreakerTrips <= 0 { 59 cfg.BreakerTrips = 5 60 } 61 if cfg.BreakerCool <= 0 { 62 cfg.BreakerCool = 30 * time.Second 63 } 64 return &Client{ 65 cfg: cfg, 66 http: &http.Client{Timeout: cfg.Timeout}, 67 cb: newBreaker(cfg.BreakerTrips, cfg.BreakerCool), 68 }, nil 69} 70 71// Health pings the plane. 72func (c *Client) Health(ctx context.Context) error { 73 return c.do(ctx, http.MethodGet, "/healthz", nil, nil) 74} 75 76// Embed returns one 384 dim vector per input. It hard fails if the plane returns the wrong width, since a 77// width mismatch would silently corrupt the pgvector(384) columns. 78func (c *Client) Embed(ctx context.Context, inputs []string) (*EmbedResponse, error) { 79 if len(inputs) == 0 { 80 return &EmbedResponse{Model: c.cfg.EmbedModel, Dimension: c.cfg.Dimension, Embeddings: nil}, nil 81 } 82 var out EmbedResponse 83 if err := c.do(ctx, http.MethodPost, "/v1/embed", EmbedRequest{Model: c.cfg.EmbedModel, Inputs: inputs}, &out); err != nil { 84 return nil, err 85 } 86 if len(out.Embeddings) != len(inputs) { 87 return nil, fmt.Errorf("inference: embed returned %d vectors for %d inputs", len(out.Embeddings), len(inputs)) 88 } 89 for i, v := range out.Embeddings { 90 if len(v) != c.cfg.Dimension { 91 return nil, fmt.Errorf("inference: embedding %d has width %d, expected %d (the column is vector(%d))", i, len(v), c.cfg.Dimension, c.cfg.Dimension) 92 } 93 } 94 return &out, nil 95} 96 97// Rerank scores documents against a query, highest first. 98func (c *Client) Rerank(ctx context.Context, query string, documents []string, topN int) (*RerankResponse, error) { 99 var out RerankResponse 100 if err := c.do(ctx, http.MethodPost, "/v1/rerank", RerankRequest{Model: c.cfg.RerankModel, Query: query, Documents: documents, TopN: topN}, &out); err != nil { 101 return nil, err 102 } 103 return &out, nil 104} 105 106// Generate runs one stateless generation. 107func (c *Client) Generate(ctx context.Context, req GenerateRequest) (*GenerateResponse, error) { 108 if req.Model == "" { 109 req.Model = c.cfg.GenerateModel 110 } 111 var out GenerateResponse 112 if err := c.do(ctx, http.MethodPost, "/v1/generate", req, &out); err != nil { 113 return nil, err 114 } 115 return &out, nil 116} 117 118// CheckFaithfulness asks the plane whether an answer is supported by its evidence (the self correction gate). 119func (c *Client) CheckFaithfulness(ctx context.Context, answer string, evidence []string) (*FaithfulnessResponse, error) { 120 var out FaithfulnessResponse 121 if err := c.do(ctx, http.MethodPost, "/v1/faithfulness", FaithfulnessRequest{Answer: answer, Evidence: evidence}, &out); err != nil { 122 return nil, err 123 } 124 return &out, nil 125} 126 127// Teach sends a contents scrubbed craft signal to the plane's teaching channel. The boundary is enforced 128// here (scrub.Marshal refuses any client fact) AND again on the gateway. When TeachingEnabled is false 129// the signal is validated but not transmitted, which is the default posture until multiple live clients 130// exist (the federated loop is scaffolded, not active). 131func (c *Client) Teach(ctx context.Context, signal scrub.TeachingSignal) error { 132 body, err := scrub.Marshal(signal) // refuses to emit a payload carrying a client fact 133 if err != nil { 134 return err 135 } 136 if !c.cfg.TeachingEnabled { 137 return nil // boundary checked, transmission disabled by posture 138 } 139 return c.doRaw(ctx, http.MethodPost, "/v1/teach", body, nil) 140} 141 142// do marshals req (if any), runs the request through the breaker with bounded retries, and decodes resp. 143func (c *Client) do(ctx context.Context, method, path string, req any, resp any) error { 144 var body []byte 145 if req != nil { 146 var err error 147 if body, err = json.Marshal(req); err != nil { 148 return err 149 } 150 } 151 return c.doRaw(ctx, method, path, body, resp) 152} 153 154// ErrPlaneUnavailable marks that the inference plane could not be reached after retries (a network error, 155// a timeout, or a 5xx) even though the circuit breaker has not opened yet. The caller maps it to a 503, so a 156// plane outage degrades to a typed error, never a generic 4xx and never a hang. 157var ErrPlaneUnavailable = errors.New("inference: plane unavailable") 158 159func (c *Client) doRaw(ctx context.Context, method, path string, body []byte, resp any) error { 160 if !c.cb.allow() { 161 return ErrCircuitOpen 162 } 163 var lastErr error 164 for attempt := 0; attempt <= c.cfg.MaxRetries; attempt++ { 165 if attempt > 0 { 166 // exponential backoff, cancellable 167 select { 168 case <-ctx.Done(): 169 return ctx.Err() 170 case <-time.After(time.Duration(attempt) * 150 * time.Millisecond): 171 } 172 } 173 err := c.attempt(ctx, method, path, body, resp) 174 if err == nil { 175 c.cb.record(true) 176 return nil 177 } 178 lastErr = err 179 if !transient(err) { 180 c.cb.record(false) 181 return err 182 } 183 } 184 c.cb.record(false) 185 // All retries exhausted on transient failures without tripping the breaker yet: a plane outage. Wrap the 186 // typed sentinel so the caller returns a 503, never a generic 4xx and never a hang. 187 return fmt.Errorf("%w: %s %s failed after %d attempts: %v", ErrPlaneUnavailable, method, path, c.cfg.MaxRetries+1, lastErr) 188} 189 190func (c *Client) attempt(ctx context.Context, method, path string, body []byte, resp any) error { 191 var r io.Reader 192 if body != nil { 193 r = bytes.NewReader(body) 194 } 195 httpReq, err := http.NewRequestWithContext(ctx, method, c.cfg.BaseURL+path, r) 196 if err != nil { 197 return err 198 } 199 if body != nil { 200 httpReq.Header.Set("Content-Type", "application/json") 201 } 202 res, err := c.http.Do(httpReq) 203 if err != nil { 204 return &transientError{err} 205 } 206 defer func() { _, _ = io.Copy(io.Discard, res.Body); _ = res.Body.Close() }() 207 if res.StatusCode >= 500 { 208 return &transientError{fmt.Errorf("inference: plane returned %d", res.StatusCode)} 209 } 210 if res.StatusCode >= 400 { 211 snippet, _ := io.ReadAll(io.LimitReader(res.Body, 512)) 212 return fmt.Errorf("inference: %s %s: status %d: %s", method, path, res.StatusCode, bytes.TrimSpace(snippet)) 213 } 214 if resp != nil { 215 return json.NewDecoder(res.Body).Decode(resp) 216 } 217 return nil 218} 219 220// transientError marks a failure as retryable (network or 5xx). 221type transientError struct{ err error } 222 223func (e *transientError) Error() string { return e.err.Error() } 224func (e *transientError) Unwrap() error { return e.err } 225 226func transient(err error) bool { 227 var t *transientError 228 return errors.As(err, &t) 229} 230