fix(proxy): wire LLM latency metrics into the streaming response path (#63)
CI / CI (push) Successful in 34m13s

`NewLLMMetricsCapture`/`ResponseWriterWrapper` had **zero call sites**, so `llm_ttft_seconds`, `llm_itl_seconds` and `llm_tokens_total` were exported as HELP/TYPE stubs with no samples (verified against the live `/metrics`).

These are the only per-request LLM latency series covering the Ollama models (`ornith:35b`, `qwen2.5:3b-instruct`), which expose no `/metrics` of their own — so those models had zero latency visibility.

## Change

Wire the wrapper around the response writer in `ServeHTTP` for model-routed requests, and delete the unused `io.WriteCloser` duplicate that made the gap hard to spot (two near-identical implementations, neither constructed).

## Correctness details

- **Only SSE responses are measured.** A non-streaming completion is written in one `Write`, which would otherwise be recorded as a TTFT sample equal to the whole request latency, with no ITL samples — plausible-looking but meaningless numbers.
- **`requestStart` is captured at the top of `ServeHTTP`** so TTFT measures from request acceptance, not from when the upstream response is copied.
- **Token count skips the `data: [DONE]` terminator.**
- **`Unwrap()`** keeps `http.ResponseController` (Flush/Hijack/deadlines) working through the wrapper, so SSE unbuffering (#33) is unaffected.

## Tests

| Test | Guards |
|---|---|
| `TestResponseWriterWrapperRecordsSSEMetrics` | TTFT=1, ITL>0, tokens=2 (terminator excluded) |
| `TestResponseWriterWrapperIgnoresNonStreaming` | JSON response records nothing |
| `TestResponseWriterWrapperPreservesStreaming` | bytes + Flush pass through |
| `TestHandlerRecordsLLMMetricsForStreaming` | end-to-end through `ServeHTTP` |

The last one fails if the wrapper is ever wired nowhere again — verified by disabling the wiring locally and watching it fail.

## Verification needed post-merge

Confirm samples appear:
```
curl -s https://api.riotpiao.com/metrics | grep -E "llm_(ttft|itl)_seconds_count|llm_tokens_total"
```
This commit was merged in pull request #63.
This commit is contained in:
2026-09-29 15:46:39 +00:00
parent 28753d11ed
commit aac68386d8
3 changed files with 251 additions and 97 deletions
+68 -96
View File
@@ -3,7 +3,6 @@ package proxy
import (
"bufio"
"fmt"
"io"
"net"
"net/http"
"strings"
@@ -12,87 +11,37 @@ import (
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
)
// LLMMetricsCapture wraps a response writer to capture TTFT and ITL metrics
type LLMMetricsCapture struct {
writer io.WriteCloser
model string
metrics *observability.Metrics
firstTokenTime time.Time
lastTokenTime time.Time
requestStartTime time.Time
ttftRecorded bool
tokenCount int64
responseStartTime time.Time
}
// NewLLMMetricsCapture creates a new metrics capture wrapper
func NewLLMMetricsCapture(writer io.WriteCloser, model string, metrics *observability.Metrics, startTime time.Time) *LLMMetricsCapture {
return &LLMMetricsCapture{
writer: writer,
model: model,
metrics: metrics,
requestStartTime: startTime,
responseStartTime: time.Now(),
}
}
// Write intercepts writes to detect tokens and record metrics
func (c *LLMMetricsCapture) Write(p []byte) (int, error) {
// Record first token time
if !c.ttftRecorded && len(p) > 0 {
now := time.Now()
ttft := now.Sub(c.requestStartTime).Milliseconds()
c.metrics.RecordTTFT(c.model, ttft)
c.ttftRecorded = true
c.firstTokenTime = now
c.lastTokenTime = now
}
// Count tokens in SSE stream (simple: count "data: " lines)
if c.ttftRecorded {
tokenCount := strings.Count(string(p), "data: ")
if tokenCount > 0 {
now := time.Now()
if !c.firstTokenTime.IsZero() && c.lastTokenTime != now {
itl := now.Sub(c.lastTokenTime).Milliseconds()
c.metrics.RecordITL(c.model, itl)
}
c.lastTokenTime = now
c.tokenCount += int64(tokenCount)
}
}
return c.writer.Write(p)
}
// Close records final metrics and closes writer
func (c *LLMMetricsCapture) Close() error {
if c.tokenCount > 0 {
c.metrics.RecordTokenCount(c.model, c.tokenCount)
}
return c.writer.Close()
}
// ResponseWriterWrapper wraps http.ResponseWriter to capture metrics
// 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
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
// 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,
writer: w,
model: model,
metrics: metrics,
startTime: startTime,
statusCode: 200,
}
}
@@ -101,29 +50,48 @@ 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) {
// Record TTFT on first write
if !w.ttftRecorded && len(b) > 0 {
if len(b) > 0 && w.isStreaming() {
now := time.Now()
ttft := now.Sub(w.startTime).Milliseconds()
w.metrics.RecordTTFT(w.model, ttft)
w.ttftRecorded = true
w.firstByteTime = now
w.lastWriteTime = now
}
// Record ITL for subsequent writes (for streaming)
if w.ttftRecorded && len(b) > 0 {
now := time.Now()
if !w.firstByteTime.IsZero() && w.lastWriteTime != now {
itl := now.Sub(w.lastWriteTime).Milliseconds()
// Only record if ITL > 0 (avoid recording same millisecond twice)
if itl > 0 {
w.metrics.RecordITL(w.model, itl)
}
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)
@@ -150,4 +118,8 @@ func (w *ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
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
}