Add TTFT & ITL Metrics for LLM Inference (#28)

Time-to-First-Token (TTFT) and Inter-Token Latency (ITL) metrics for LLM inference observability Metrics: llm_ttft_seconds, llm_itl_seconds, llm_tokens_total Closes #31 #32 #33 --------- Co-authored-by: poimen <[email protected]> Reviewed-on: #28
2026-09-15 08:53:58 +00:00
parent 30e0a83a50
commit d439536ca9
22 changed files with 2224 additions and 199 deletions
@@ -381,3 +381,66 @@ func TestBodySizeCappedDispatch(t *testing.T) {
 		t.Errorf("expected 200 for reasonable body, got %d", resp.StatusCode)
 	}
 }
+
+// TestUpstreamModelRewrite verifies that the model field is rewritten when upstreamModel is set.
+func TestUpstreamModelRewrite(t *testing.T) {
+	var receivedModel string
+
+	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, _ := io.ReadAll(r.Body)
+		var payload map[string]interface{}
+		json.Unmarshal(body, &payload)
+		receivedModel = payload["model"].(string)
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprint(w, `{}`)
+	}))
+	defer upstreamServer.Close()
+
+	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
+
+	cfg := &config.Config{
+		Models: map[string]*config.ModelUpstream{
+			// Client sends "ornith:35b", upstream expects "qwen2.5:72b-instruct"
+			"ornith:35b": {
+				Name:          "ornith:35b",
+				Address:       upstreamAddr,
+				UpstreamModel: "qwen2.5:72b-instruct",
+			},
+			// No rewrite - upstream model same as client model
+			"reasoning": {
+				Name:    "reasoning",
+				Address: upstreamAddr,
+			},
+		},
+		Routes: make(map[string]*config.Route),
+	}
+
+	handler := New(cfg)
+	defer handler.Close()
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	// Test 1: Model should be rewritten
+	requestBody := `{"model":"ornith:35b","messages":[{"role":"user","content":"hi"}]}`
+	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
+	if err != nil {
+		t.Fatalf("request failed: %v", err)
+	}
+	resp.Body.Close()
+
+	if receivedModel != "qwen2.5:72b-instruct" {
+		t.Errorf("expected upstream to receive model 'qwen2.5:72b-instruct', got '%s'", receivedModel)
+	}
+
+	// Test 2: No rewrite when upstreamModel is empty
+	receivedModel = ""
+	requestBody = `{"model":"reasoning","messages":[{"role":"user","content":"hi"}]}`
+	resp, _ = http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
+	resp.Body.Close()
+
+	if receivedModel != "reasoning" {
+		t.Errorf("expected upstream to receive model 'reasoning', got '%s'", receivedModel)
+	}
+}
@@ -0,0 +1,153 @@
+package proxy
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"strings"
+	"time"
+
+	"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
+type ResponseWriterWrapper struct {
+	writer          http.ResponseWriter
+	statusCode      int
+	metrics         *observability.Metrics
+	model           string
+	startTime       time.Time
+	firstByteTime   time.Time
+	lastWriteTime   time.Time
+	ttftRecorded    bool
+}
+
+// NewResponseWriterWrapper creates a wrapper for response writer
+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()
+}
+
+// Write implements http.ResponseWriter
+func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
+	// Record TTFT on first write
+	if !w.ttftRecorded && len(b) > 0 {
+		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)
+			}
+		}
+		w.lastWriteTime = now
+	}
+
+	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")
+}
+
+
@@ -122,6 +122,20 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
 		}
 	}

+	// If upstream expects a different model name, rewrite the body
+	if modelUpstream.UpstreamModel != "" && modelUpstream.UpstreamModel != modelName {
+		payload["model"] = modelUpstream.UpstreamModel
+		newBody, err := json.Marshal(payload)
+		if err != nil {
+			return nil, &modelValidationError{
+				Kind:    "invalid_request",
+				Message: fmt.Sprintf("failed to rewrite model name: %v", err),
+			}
+		}
+		r.Body = io.NopCloser(bytes.NewReader(newBody))
+		r.ContentLength = int64(len(newBody))
+	}
+
 	// Determine the upstream path based on the request path
 	upstreamPath := path
 	if path == "/v1/rerank" {