llmgw/internal/router/cache.go

package router

import (
	"bytes"
	"container/list"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
	storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Reusing answers, for a project that asks: the same request to the same
// model, from the same project and within its TTL, is answered with what
// the model said the first time rather than asking it again.
//
// Only an answer that would come out the same is reused, unless the project
// says otherwise: a chat or a completion asked at temperature 0 or with a
// seed, and an embedding, a rerank or a score, which have no sampling at
// all. A request sampled above temperature 0 is asked for a new answer, and
// handing it an old one would quietly turn variety into repetition.
//
// The answers are held in this process's memory and nowhere else. A shared
// cache in the store would write every answer of every project that turned
// this on to the database — the very thing the request log is opt-in for —
// and add a round trip to each request to save one; with two control planes
// behind a balancer each fills its own, which costs some hits and no
// correctness.

// cacheHeader says what the cache did for a request, hit, miss or bypass,
// and is how a client asks it to stay out of the way: "no".
const cacheHeader = "x-llmgw-cache"

// Defaults for router.cache.
const (
	defaultCacheBytes      = 128 << 20
	defaultCacheEntryBytes = 1 << 20
)

// cacheEndpoints are those whose answers can be reused. The chat and
// completion endpoints only when the answer is reproducible; the others
// always, as they sample nothing. Audio is a form, not JSON, and is not
// among them.
var cacheEndpoints = map[string]bool{
	"/v1/chat/completions": true,
	"/v1/completions":      true,
	"/v1/embeddings":       true,
	"/v1/rerank":           true,
	"/v1/score":            true,
}

// sampledEndpoints are those whose answer depends on sampling.
var sampledEndpoints = map[string]bool{
	"/v1/chat/completions": true,
	"/v1/completions":      true,
}

// cacheState is where a request's answer is kept, for one whose answer
// may be reused and was not found.
type cacheState struct {
	// key is the answer's place in the cache; empty when it is not kept.
	key string
	// project is whose answer it is, and ttl how long it is reused.
	project string
	ttl     time.Duration
}

// cachedAnswer is one reusable answer.
type cachedAnswer struct {
	key, project string
	// model is the model that answered, by its catalog name.
	model string
	path  string
	// body is the answer as one JSON object — a chat.completion, a
	// text_completion, a list of embeddings — even when it was streamed.
	body []byte
	at   time.Time
	// expires is when the TTL in force when it was kept runs out. The TTL
	// in force when it is asked for is checked too, so shortening a
	// project's TTL shortens the life of what it already has.
	expires time.Time
	size    int64
}

// answerCache is a least-recently-used cache of answers, bounded in bytes.
type answerCache struct {
	maxBytes, maxEntry int64

	mu    sync.Mutex
	lru   *list.List // of *cachedAnswer, most recently used at the front
	by    map[string]*list.Element
	bytes int64
	// swept is when expired answers were last cleared out.
	swept time.Time
	// size reports bytes and answers held, after each change.
	size func(bytes int64, n int)
}

func newAnswerCache(maxBytes, maxEntry int64, size func(int64, int)) *answerCache {
	if maxBytes <= 0 {
		maxBytes = defaultCacheBytes
	}
	if maxEntry <= 0 {
		maxEntry = defaultCacheEntryBytes
	}
	maxEntry = min(maxEntry, maxBytes)
	if size == nil {
		size = func(int64, int) {}
	}
	return &answerCache{maxBytes: maxBytes, maxEntry: maxEntry, lru: list.New(),
		by: map[string]*list.Element{}, size: size}
}

// get is the answer kept under key, when it is younger than ttl and has
// not expired.
func (c *answerCache) get(key string, ttl time.Duration, now time.Time) *cachedAnswer {
	c.mu.Lock()
	defer c.mu.Unlock()
	el, ok := c.by[key]
	if !ok {
		return nil
	}
	a := el.Value.(*cachedAnswer)
	if now.After(a.expires) || now.Sub(a.at) >= ttl {
		c.remove(el)
		c.size(c.bytes, len(c.by))
		return nil
	}
	c.lru.MoveToFront(el)
	return a
}

// put keeps an answer, making room for it by dropping those used longest
// ago. One larger than an entry may be is not kept.
func (c *answerCache) put(a *cachedAnswer, now time.Time) bool {
	// What the entry costs besides its body: the key twice, the names, and
	// the bookkeeping, roughly.
	a.size = int64(len(a.body) + 2*len(a.key) + len(a.project) + len(a.model) + len(a.path) + 160)
	if a.size > c.maxEntry {
		return false
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if el, ok := c.by[a.key]; ok {
		c.remove(el)
	}
	// Now and then, those past their time, so a quiet project's answers
	// do not sit in memory until pushed out by others.
	if now.Sub(c.swept) > time.Minute {
		c.swept = now
		for el := c.lru.Back(); el != nil; {
			prev := el.Prev()
			if now.After(el.Value.(*cachedAnswer).expires) {
				c.remove(el)
			}
			el = prev
		}
	}
	for c.bytes+a.size > c.maxBytes && c.lru.Len() > 0 {
		c.remove(c.lru.Back())
	}
	c.by[a.key] = c.lru.PushFront(a)
	c.bytes += a.size
	c.size(c.bytes, len(c.by))
	return true
}

func (c *answerCache) remove(el *list.Element) {
	a := el.Value.(*cachedAnswer)
	c.lru.Remove(el)
	delete(c.by, a.key)
	c.bytes -= a.size
}

// answerCacheOf is a project's setting, read at most every requestLogTTL
// with its other settings. Unreadable, it reuses nothing: asking the model
// again is the safe way to be wrong.
func (r *Router) answerCacheOf(ctx context.Context, project string) storepkg.AnswerCache {
	if r.store == nil || project == "" {
		return storepkg.AnswerCache{}
	}
	r.reqLogs.mu.Lock()
	c, ok := r.reqLogs.by[project]
	r.reqLogs.mu.Unlock()
	if ok && time.Since(c.at) < requestLogTTL {
		return c.cache
	}
	c, err := r.readProjectPolicy(ctx, project)
	if err != nil {
		r.log.Warn("could not read whether a project reuses answers", "project", project, "err", err)
	}
	return c.cache
}

// cacheLookup decides whether a request's answer may come from its
// project's cache, says what it decided in x-llmgw-cache, and returns the
// answer when there is one. Run after the guardrails have judged the
// request, so a cached answer reaches only a request that would have been
// sent to the model, and the request it is looked up by is the one the
// model would have been sent.
func (r *Router) cacheLookup(ctx context.Context, w http.ResponseWriter, req *parsedRequest) *cachedAnswer {
	if r.answers == nil || req.internal || req.ProjectID == "" {
		return nil
	}
	set := r.answerCacheOf(ctx, req.ProjectID)
	if set.TTL <= 0 {
		// Says nothing: there is no cache here to have bypassed.
		return nil
	}
	result := "bypass"
	defer func() {
		w.Header().Set(cacheHeader, result)
		r.m.cacheRequests.WithLabelValues(r.metricModel(req.Model), result).Inc()
	}()
	if !reusable(req, set.Any) {
		return nil
	}
	model, _ := r.resolve(req.Model)
	key, ok := answerKey(req.ProjectID, req.Path, model, req.UpstreamBody)
	if !ok {
		return nil
	}
	now := time.Now()
	if a := r.answers.get(key, set.TTL, now); a != nil && a.project == req.ProjectID {
		result = "hit"
		// Age is HTTP's own way of saying how old a reused answer is.
		w.Header().Set("Age", strconv.Itoa(int(now.Sub(a.at).Seconds())))
		return a
	}
	result = "miss"
	req.cache = cacheState{key: key, project: req.ProjectID, ttl: set.TTL}
	return nil
}

// reusable says whether a request's answer may be reused at all: not one
// the client asked to be answered afresh, not one continuing a stored
// conversation (each turn is new, and is saved), not one under a policy
// that keeps nothing — a cached answer is a kept one — and, for chat and
// completions, only one asked at temperature 0 or with a seed unless the
// project reuses any answer.
func reusable(req *parsedRequest, anyAnswer bool) bool {
	if req.noStore || req.ThreadID != "" || req.ContentType != "" || !cacheEndpoints[req.Path] {
		return false
	}
	// An answer from knowledge bases is the caller's own: two keys of one
	// project can read different collections.
	if req.Header.Get(knowledgeHeader) != "" {
		return false
	}
	if v := strings.ToLower(strings.TrimSpace(req.Header.Get(cacheHeader))); v == "no" || v == "off" || v == "bypass" {
		return false
	}
	for _, v := range req.Header.Values("Cache-Control") {
		v = strings.ToLower(v)
		if strings.Contains(v, "no-cache") || strings.Contains(v, "no-store") {
			return false
		}
	}
	if !sampledEndpoints[req.Path] || anyAnswer {
		return true
	}
	var b struct {
		Temperature *float64        `json:"temperature"`
		Seed        json.RawMessage `json:"seed"`
	}
	if json.Unmarshal(req.UpstreamBody, &b) != nil {
		return false
	}
	seeded := len(b.Seed) > 0 && string(b.Seed) != "null"
	return seeded || b.Temperature != nil && *b.Temperature == 0
}

// answerFields are taken out of a request before it is looked up: they
// change how an answer is delivered or recorded, never what it says.
var answerFields = []string{"model", "stream", "stream_options", "user", "metadata", "store", "service_tier"}

// answerKey is where a request's answer is kept: the project, the endpoint,
// the model by its catalog name, and the request with what does not change
// the answer taken out and its keys sorted. The project is spelled out
// rather than only hashed, so no answer is ever looked up across projects.
func answerKey(project, path, model string, body []byte) (string, bool) {
	d := json.NewDecoder(bytes.NewReader(body))
	d.UseNumber()
	var m map[string]any
	if err := d.Decode(&m); err != nil || m == nil {
		return "", false
	}
	for _, f := range answerFields {
		delete(m, f)
	}
	// encoding/json writes map keys sorted, at every depth.
	norm, err := json.Marshal(canonicalNumbers(m))
	if err != nil {
		return "", false
	}
	h := sha256.New()
	for _, part := range [][]byte{[]byte(path), []byte(model), norm} {
		h.Write(part)
		h.Write([]byte{0})
	}
	return project + "\x00" + hex.EncodeToString(h.Sum(nil)), true
}

// canonicalNumbers writes each number one way, so 0, 0.0 and 0e0 are the
// same temperature. Whole numbers are left as written: a seed above 2^53
// would lose its last digits as a float, and two seeds become one.
func canonicalNumbers(v any) any {
	switch x := v.(type) {
	case map[string]any:
		for k, e := range x {
			x[k] = canonicalNumbers(e)
		}
	case []any:
		for i, e := range x {
			x[i] = canonicalNumbers(e)
		}
	case json.Number:
		s := string(x)
		if _, err := strconv.ParseInt(s, 10, 64); err == nil {
			return x
		}
		if f, err := strconv.ParseFloat(s, 64); err == nil {
			if f == float64(int64(f)) && f > -1<<53 && f < 1<<53 {
				return json.Number(strconv.FormatInt(int64(f), 10))
			}
			return json.Number(strconv.FormatFloat(f, 'g', -1, 64))
		}
	}
	return v
}

// cacheWriter keeps a copy of what is written to the client, up to a
// bound, so an answer can be kept once it is known to be a good one.
type cacheWriter struct {
	http.ResponseWriter
	status int
	buf    bytes.Buffer
	limit  int
	over   bool
}

func (c *cacheWriter) WriteHeader(code int) {
	if c.status == 0 {
		c.status = code
	}
	c.ResponseWriter.WriteHeader(code)
}

func (c *cacheWriter) Write(p []byte) (int, error) {
	if c.status == 0 {
		c.status = http.StatusOK
	}
	if !c.over {
		if c.buf.Len()+len(p) > c.limit {
			c.over = true
			c.buf = bytes.Buffer{}
		} else {
			c.buf.Write(p)
		}
	}
	return c.ResponseWriter.Write(p)
}

// Flush passes through, which the proxy needs to stream.
func (c *cacheWriter) Flush() {
	if f, ok := c.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func (c *cacheWriter) Unwrap() http.ResponseWriter { return c.ResponseWriter }

// cacheWriterFor wraps w when a request's answer may be kept.
func (r *Router) cacheWriterFor(w http.ResponseWriter, req *parsedRequest) (http.ResponseWriter, *cacheWriter) {
	if req.cache.key == "" {
		return w, nil
	}
	cw := &cacheWriter{ResponseWriter: w, limit: int(r.answers.maxEntry)}
	return cw, cw
}

// cacheKeep keeps an answer that may be reused: a whole 200 from the model
// asked for — not a fallback, which is another model, nor a canary, which
// is weights on trial — that no guard changed or stopped.
func (r *Router) cacheKeep(req *parsedRequest, res dispatchResult, status int, cw *cacheWriter) {
	if cw == nil || cw.over || status != http.StatusOK || cw.status != http.StatusOK || res.err != nil ||
		res.outcome.Status != http.StatusOK || res.outcome.Stopped || res.fallback || res.onCanary {
		return
	}
	if res.judged && res.guard.Action != guard.Allow && !res.guard.Shadow {
		return
	}
	var body []byte
	captured := bytes.TrimSpace(cw.buf.Bytes())
	if strings.Contains(cw.Header().Get("Content-Type"), "text/event-stream") ||
		bytes.HasPrefix(captured, []byte("data:")) {
		b, ok := assembleStream(req.Path, captured, res.outcome.Usage)
		if !ok {
			return
		}
		body = b
	} else {
		if len(captured) == 0 || captured[0] != '{' || !json.Valid(captured) {
			return
		}
		body = bytes.Clone(captured)
	}
	model := res.model
	if model == "" {
		model, _ = r.resolve(req.Model)
	}
	now := time.Now()
	r.answers.put(&cachedAnswer{key: req.cache.key, project: req.cache.project, model: model,
		path: req.Path, body: body, at: now, expires: now.Add(req.cache.ttl)}, now)
}

// serveCached answers a request from the cache, as the model would have:
// as one JSON object, or as a stream of chunks to a client that asked for
// one. The guardrails' output policy judges it as it judges any answer —
// the policy in force now, not the one when it was first given — and the
// usage is recorded as a request that generated nothing: no tokens, no
// cost, marked as a cache hit.
func (r *Router) serveCached(ctx context.Context, w http.ResponseWriter, req *parsedRequest, a *cachedAnswer, start time.Time) {
	body := a.body
	res := dispatchResult{servedModel: servedModel{model: a.model}, cached: true}
	if r.guard != nil {
		if f := r.guard.Output(ctx, r.subject(req, a.model)); f != nil {
			// A copy: the kept answer is shared by every request reading it,
			// and a filter is free to rewrite what it is given.
			body = f.Buffered(bytes.Clone(body))
			res.guard, res.judged = f.Verdict(), true
			if res.guard.Stricter(req.guardIn) {
				setGuardHeaders(w.Header(), res.guard)
			}
		}
	}
	var werr error
	if req.Stream && sampledEndpoints[req.Path] {
		werr = writeCachedStream(w, req, body)
	} else {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, werr = w.Write(body)
	}
	status := http.StatusOK
	if werr != nil {
		status = 499
	}
	res.outcome = proxy.Outcome{Status: http.StatusOK, Wrote: true, Usage: &proxy.Usage{},
		Reply: replyText(body)}
	took := time.Since(start)
	r.m.observeRequest(r.metricModel(req.Model), "cache", status, took, nil, 0)
	r.account(req, res, status, took)
	r.accessLog(req, res, status, took)
	r.observeOutput(req, res, res.judged)
	if req.logRetain > 0 {
		go r.keepRequest(context.WithoutCancel(ctx), req, res, status)
	}
}

// writeCachedStream sends a kept answer as the chunks a model streams, with
// the usage chunk only to a client that asked for one.
func writeCachedStream(w http.ResponseWriter, req *parsedRequest, body []byte) error {
	events, err := streamFromAnswer(req.Path, body, clientWantsUsage(req))
	if err != nil {
		// Not an answer that can be streamed — a guard's replacement, say.
		// Sent whole, which a client reads as a model that does not stream.
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, err = w.Write(body)
		return err
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.WriteHeader(http.StatusOK)
	for _, e := range append(events, "[DONE]") {
		if _, err := io.WriteString(w, "data: "+e+"\n\n"); err != nil {
			return err
		}
	}
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
	return nil
}

// clientWantsUsage says whether the client asked for the usage chunk
// itself, rather than the gateway asking for it on its behalf.
func clientWantsUsage(req *parsedRequest) bool {
	if req.SuppressUsageChunk {
		return false
	}
	var b struct {
		StreamOptions *struct {
			IncludeUsage bool `json:"include_usage"`
		} `json:"stream_options"`
	}
	return json.Unmarshal(req.UpstreamBody, &b) == nil && b.StreamOptions != nil && b.StreamOptions.IncludeUsage
}

// replyText is an answer's first choice as text, for the request log.
func replyText(body []byte) string {
	var c struct {
		Choices []struct {
			Message struct {
				Content *string `json:"content"`
			} `json:"message"`
			Text *string `json:"text"`
		} `json:"choices"`
	}
	if json.Unmarshal(body, &c) != nil || len(c.Choices) == 0 {
		return ""
	}
	if s := c.Choices[0].Message.Content; s != nil {
		return *s
	}
	if s := c.Choices[0].Text; s != nil {
		return *s
	}
	return ""
}