homelab-frontend/internal/proxy/llm_metrics.go

package proxy

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"

	"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
)

// ResponseWriterWrapper wraps an http.ResponseWriter to capture per-request LLM
// latency metrics (TTFT, ITL, generated tokens) while a streaming response is
// copied from the upstream by httputil.ReverseProxy.
//
// Only server-sent-event responses are measured. A non-streaming completion is
// written in a single Write, so recording it would emit one "TTFT" sample equal
// to the whole request latency and no ITL samples at all — numbers that look
// plausible but mean nothing.
type ResponseWriterWrapper struct {
	writer        http.ResponseWriter
	statusCode    int
	metrics       *observability.Metrics
	model         string
	startTime     time.Time
	firstByteTime time.Time
	lastWriteTime time.Time
	ttftRecorded  bool
	streamChecked bool
	streaming     bool
}

// NewResponseWriterWrapper creates a wrapper for response writer.
// startTime must be the time the request entered the gateway, not the time this
// wrapper was constructed, otherwise TTFT measures the wrong interval.
func NewResponseWriterWrapper(w http.ResponseWriter, model string, metrics *observability.Metrics, startTime time.Time) *ResponseWriterWrapper {
	return &ResponseWriterWrapper{
		writer:     w,
		model:      model,
		metrics:    metrics,
		startTime:  startTime,
		statusCode: 200,
	}
}

// Header implements http.ResponseWriter
func (w *ResponseWriterWrapper) Header() http.Header {
	return w.writer.Header()
}

// isStreaming reports whether this is an SSE response. Resolved once from the
// Content-Type header, which ReverseProxy has already set before the first write.
func (w *ResponseWriterWrapper) isStreaming() bool {
	if !w.streamChecked {
		w.streamChecked = true
		w.streaming = strings.Contains(w.writer.Header().Get("Content-Type"), "text/event-stream")
	}
	return w.streaming
}

// countSSETokens counts generated tokens in a chunk of an SSE stream. These
// upstreams emit one `data:` event per token, so the event count is the token
// count; the `[DONE]` terminator is not a token.
func countSSETokens(b []byte) int64 {
	var n int64
	for _, line := range strings.Split(string(b), "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "data:") && !strings.HasPrefix(line, "data: [DONE]") {
			n++
		}
	}
	return n
}

// Write implements http.ResponseWriter
func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
	if len(b) > 0 && w.isStreaming() {
		now := time.Now()
		if !w.ttftRecorded {
			w.metrics.RecordTTFT(w.model, now.Sub(w.startTime).Milliseconds())
			w.ttftRecorded = true
			w.firstByteTime = now
		} else if itl := now.Sub(w.lastWriteTime).Milliseconds(); itl > 0 {
			// Only record positive ITL: a zero means two chunks landed in the
			// same millisecond, which carries no timing information.
			w.metrics.RecordITL(w.model, itl)
		}
		w.lastWriteTime = now

		if n := countSSETokens(b); n > 0 {
			w.metrics.RecordTokenCount(w.model, n)
		}
	}

	return w.writer.Write(b)
}

// WriteHeader implements http.ResponseWriter
func (w *ResponseWriterWrapper) WriteHeader(statusCode int) {
	w.statusCode = statusCode
	w.writer.WriteHeader(statusCode)
}

// Flush implements http.Flusher
func (w *ResponseWriterWrapper) Flush() {
	if flusher, ok := w.writer.(http.Flusher); ok {
		flusher.Flush()
	}
}

// Hijack implements http.Hijacker for streaming
func (w *ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := w.writer.(http.Hijacker); ok {
		return hijacker.Hijack()
	}
	return nil, nil, fmt.Errorf("response writer does not implement Hijacker")
}

// Unwrap exposes the underlying writer to http.ResponseController so optional
// capabilities (Flush, Hijack, deadlines) keep working through the wrapper.
func (w *ResponseWriterWrapper) Unwrap() http.ResponseWriter {
	return w.writer
}