homelab-frontend/internal/proxy/llm_metrics.go

package proxy

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"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")
}

// Import http package
import "net/http"