homelab-frontend/internal/observability/prometheus.go

package observability

import (
	"fmt"
	"sort"
	"strings"
)

// PrometheusExporter exports metrics in Prometheus text format
type PrometheusExporter struct {
	metrics *Metrics
}

// NewPrometheusExporter creates a new Prometheus exporter
func NewPrometheusExporter(m *Metrics) *PrometheusExporter {
	return &PrometheusExporter{metrics: m}
}

// Export returns metrics in Prometheus text format
func (p *PrometheusExporter) Export() string {
	var lines []string

	lines = append(lines, "# HELP llm_ttft_seconds Time to first token for LLM inference (seconds)")
	lines = append(lines, "# TYPE llm_ttft_seconds histogram")
	p.exportTTFT(&lines)

	lines = append(lines, "# HELP llm_itl_seconds Inter-token latency for LLM inference (seconds)")
	lines = append(lines, "# TYPE llm_itl_seconds histogram")
	p.exportITL(&lines)

	lines = append(lines, "# HELP llm_tokens_total Total tokens generated")
	lines = append(lines, "# TYPE llm_tokens_total counter")
	p.exportTokens(&lines)

	lines = append(lines, "# HELP request_duration_seconds Request latency")
	lines = append(lines, "# TYPE request_duration_seconds histogram")
	p.exportRequestDuration(&lines)

	return strings.Join(lines, "\n") + "\n"
}

func (p *PrometheusExporter) exportTTFT(lines *[]string) {
	p.metrics.mu.RLock()
	defer p.metrics.mu.RUnlock()

	// Calculate statistics for each model
	for key, samples := range p.metrics.ttftMs {
		if len(samples) == 0 {
			continue
		}

		model := extractModel(key)
		sum := sumInt64(samples)
		avg := sum / int64(len(samples))

		// Export histogram buckets (in seconds)
		buckets := []float64{0.001, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0}
		for _, bucket := range buckets {
			count := countLessOrEqual(samples, int64(bucket*1000))
			*lines = append(*lines, fmt.Sprintf(
				`llm_ttft_seconds_bucket{model="%s",le="%.3f"} %d`,
				model, bucket, count,
			))
		}

		*lines = append(*lines, fmt.Sprintf(
			`llm_ttft_seconds_bucket{model="%s",le="+Inf"} %d`,
			model, len(samples),
		))
		*lines = append(*lines, fmt.Sprintf(
			`llm_ttft_seconds_sum{model="%s"} %.3f`,
			model, float64(sum)/1000,
		))
		*lines = append(*lines, fmt.Sprintf(
			`llm_ttft_seconds_count{model="%s"} %d`,
			model, len(samples),
		))
	}
}

func (p *PrometheusExporter) exportITL(lines *[]string) {
	p.metrics.mu.RLock()
	defer p.metrics.mu.RUnlock()

	for key, samples := range p.metrics.itlMs {
		if len(samples) == 0 {
			continue
		}

		model := extractModel(key)
		sum := sumInt64(samples)

		// Export histogram buckets (in seconds)
		buckets := []float64{0.001, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0}
		for _, bucket := range buckets {
			count := countLessOrEqual(samples, int64(bucket*1000))
			*lines = append(*lines, fmt.Sprintf(
				`llm_itl_seconds_bucket{model="%s",le="%.3f"} %d`,
				model, bucket, count,
			))
		}

		*lines = append(*lines, fmt.Sprintf(
			`llm_itl_seconds_bucket{model="%s",le="+Inf"} %d`,
			model, len(samples),
		))
		*lines = append(*lines, fmt.Sprintf(
			`llm_itl_seconds_sum{model="%s"} %.3f`,
			model, float64(sum)/1000,
		))
		*lines = append(*lines, fmt.Sprintf(
			`llm_itl_seconds_count{model="%s"} %d`,
			model, len(samples),
		))
	}
}

func (p *PrometheusExporter) exportTokens(lines *[]string) {
	p.metrics.mu.RLock()
	defer p.metrics.mu.RUnlock()

	// Sort keys for consistent output
	var keys []string
	for k := range p.metrics.tokenCount {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for _, key := range keys {
		model := extractModel(key)
		count := p.metrics.tokenCount[key]
		*lines = append(*lines, fmt.Sprintf(
			`llm_tokens_total{model="%s"} %d`,
			model, count,
		))
	}
}

func (p *PrometheusExporter) exportRequestDuration(lines *[]string) {
	p.metrics.mu.RLock()
	defer p.metrics.mu.RUnlock()

	// Sort keys for consistent output
	var keys []string
	for k := range p.metrics.requestDuration {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for _, key := range keys {
		route, upstream := parseKey(key)
		totalMs := p.metrics.requestDuration[key]
		count := int64(1) // We'd need to track count separately in real impl

		if buckets, ok := p.metrics.requestDurationBuckets[key]; ok {
			for bucket := range buckets {
				*lines = append(*lines, fmt.Sprintf(
					`request_duration_seconds_bucket{route="%s",upstream="%s",le="%.1f"} %d`,
					route, upstream, bucket, buckets[bucket],
				))
			}
		}

		*lines = append(*lines, fmt.Sprintf(
			`request_duration_seconds_sum{route="%s",upstream="%s"} %.3f`,
			route, upstream, float64(totalMs)/1000,
		))
		*lines = append(*lines, fmt.Sprintf(
			`request_duration_seconds_count{route="%s",upstream="%s"} %d`,
			route, upstream, count,
		))
	}
}

func extractModel(key string) string {
	parts := strings.Split(key, ":")
	if len(parts) >= 3 {
		return parts[2]
	}
	return key
}

func parseKey(key string) (string, string) {
	parts := strings.Split(key, ":")
	if len(parts) >= 2 {
		return parts[0], parts[1]
	}
	return key, ""
}

func countLessOrEqual(samples []int64, threshold int64) int {
	count := 0
	for _, s := range samples {
		if s <= threshold {
			count++
		}
	}
	return count
}