llmgw/internal/router/requestlog.go

package router

import (
	"bytes"
	"context"
	"fmt"
	"hash/fnv"
	"net/http"
	"sync"
	"time"

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

// Keeping requests and answers, for a project that asks: what was sent to
// the model and what came back, for following one request up and for
// sending it again to another model.

// requestLogTTL is how long a project's setting is believed before it is
// read again.
const requestLogTTL = 30 * time.Second

type requestLogs struct {
	mu sync.Mutex
	by map[string]cachedLog
}

type cachedLog struct {
	l  storepkg.RequestLog
	at time.Time
}

// requestLogOf is a project's setting, read at most every requestLogTTL.
func (r *Router) requestLogOf(ctx context.Context, project string) storepkg.RequestLog {
	if r.store == nil || project == "" {
		return storepkg.RequestLog{}
	}
	r.reqLogs.mu.Lock()
	c, ok := r.reqLogs.by[project]
	r.reqLogs.mu.Unlock()
	if ok && time.Since(c.at) < requestLogTTL {
		return c.l
	}
	l, err := r.store.RequestLogOf(ctx, project)
	if err != nil {
		r.log.Warn("could not read whether a project keeps its requests", "project", project, "err", err)
		return c.l
	}
	r.reqLogs.mu.Lock()
	if r.reqLogs.by == nil {
		r.reqLogs.by = map[string]cachedLog{}
	}
	r.reqLogs.by[project] = cachedLog{l: l, at: time.Now()}
	r.reqLogs.mu.Unlock()
	return l
}

// sampled says whether a request falls in a share of them, by its id, so
// the same id always says the same.
func sampled(id string, share float64) bool {
	if share >= 1 {
		return true
	}
	h := fnv.New32a()
	h.Write([]byte(id))
	return float64(h.Sum32()%10000) < share*10000
}

// decideRequestLog marks a request to be kept, once it is known who sent
// it and that no policy forbids keeping it.
func (r *Router) decideRequestLog(ctx context.Context, req *parsedRequest) {
	if req.noStore || req.internal {
		return
	}
	l := r.requestLogOf(ctx, req.ProjectID)
	if l.Retain > 0 && sampled(req.RequestID, l.Sample) {
		req.logRetain = l.Retain
	}
}

// keepRequest writes a finished request and its answer, for one marked.
func (r *Router) keepRequest(ctx context.Context, req *parsedRequest, res dispatchResult, status int) {
	if req.logRetain <= 0 {
		return
	}
	served := res.model
	if served == "" {
		served = req.Model
	}
	now := time.Now()
	b := storepkg.RequestBody{RequestID: req.RequestID, ProjectID: req.ProjectID, At: now,
		ExpiresAt: now.Add(req.logRetain), Path: req.Path, Model: req.Model, ServedModel: served,
		Status: status, Request: string(req.UpstreamBody), Response: res.outcome.Reply,
		Truncated: res.outcome.ReplyTruncated}
	if err := r.store.PutRequestBody(ctx, b); err != nil {
		r.log.Warn("could not keep a request", "request_id", req.RequestID, "err", err)
	}
}

// pruneRequestBodies removes kept requests past their retention, hourly.
func (r *Router) pruneRequestBodies(ctx context.Context) {
	t := time.NewTicker(time.Hour)
	defer t.Stop()
	for {
		if n, err := r.store.PruneRequestBodies(ctx, time.Now()); err != nil {
			if ctx.Err() == nil {
				r.log.Warn("could not prune kept requests", "err", err)
			}
		} else if n > 0 {
			r.log.Info("removed kept requests past their retention", "requests", n)
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

// Replay sends a kept request again to a model, as the gateway itself: past
// admission, accounting and the guardrails, which judged it the first time.
// For comparing models on real traffic. It answers the status, the body,
// and how long it took.
func (r *Router) Replay(ctx context.Context, path, model string, body []byte) (int, []byte, time.Duration, error) {
	if !batchEndpoints[path] {
		return 0, nil, 0, fmt.Errorf("a request to %s cannot be replayed", path)
	}
	body = withoutField(withModel(body, model), "stream")
	hreq, err := http.NewRequestWithContext(ctx, http.MethodPost, path, bytes.NewReader(body))
	if err != nil {
		return 0, nil, 0, err
	}
	req, err := r.parse(hreq, path)
	if err != nil {
		return 0, nil, 0, err
	}
	req.internal, req.Stream = true, false
	req.RequestID = requestID("")
	req.Header = http.Header{"X-Request-Id": []string{req.RequestID}}
	w := &bufferWriter{header: http.Header{}}
	start := time.Now()
	res := r.dispatch(ctx, w, req)
	took := time.Since(start)
	if res.err != nil {
		return 0, nil, took, res.err
	}
	status := w.status
	if status == 0 {
		status = http.StatusOK
	}
	return status, w.buf.Bytes(), took, nil
}