fix(proxy): wire LLM latency metrics into the streaming response path
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. 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. 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. - 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 (issue #33) is unaffected. Tests: unit coverage for SSE vs non-streaming, wrapper transparency, and an end-to-end handler test that fails if the wrapper is ever wired nowhere again.
This commit is contained in:
Admin Bot
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user