Building an Adaptive LLM Routing Gateway for Multi-Model Cost-Latency Tradeoffs — code snippets

snippet-1.go

package gateway

import "time"

// ModelConfig represents the static metadata and dynamic tracking of an LLM upstream.
type ModelConfig struct {
	ID             string        `json:"id"`
	Provider       string        `json:"provider"`
	InputCostPerM  float64       `json:"input_cost_per_m"`  // USD per 1M input tokens
	OutputCostPerM float64       `json:"output_cost_per_m"` // USD per 1M output tokens
	MaxTokens      int           `json:"max_tokens"`
	P50Latency     time.Duration `json:"p50_latency"`
	P99Latency     time.Duration `json:"p99_latency"`
	QualityScore   float64       `json:"quality_score"` // Internal benchmark score [0.0 - 1.0]
}

// RequestMetadata holds the analyzed attributes of the incoming client request.
type RequestMetadata struct {
	Prompt             string        `json:"prompt"`
	EstimatedTokens    int           `json:"estimated_tokens"`
	RequiredQuality    float64       `json:"required_quality"` // Threshold required for the task
	LatencySLA         time.Duration `json:"latency_sla"`       // Maximum tolerable latency
	RequiresReasoning  bool          `json:"requires_reasoning"`
	RequiresStructured bool          `json:"requires_structured"`
}

// RouteResult indicates the model chosen by the Policy Engine and its fallback sequence.
type RouteResult struct {
	PrimaryModel ModelConfig   `json:"primary_model"`
	FallbackChain []ModelConfig `json:"fallback_chain"`
	CacheHit      bool          `json:"cache_hit"`
	CacheKey      string        `json:"cache_key"`
}

snippet-2.go

package gateway

import (
	"context"
	"fmt"
	"unsafe"

	"github.com/redis/go-redis/v9"
)

type SemanticCache struct {
	redisClient *redis.Client
	threshold   float64
}

func NewSemanticCache(rdb *redis.Client, similarityThreshold float64) *SemanticCache {
	return &SemanticCache{
		redisClient: rdb,
		threshold:   similarityThreshold,
	}
}

// QueryCache searches Redis for a similar prompt embedding and returns the cached answer if found.
func (s *SemanticCache) QueryCache(ctx context.Context, embedding []float32) (string, bool, error) {
	// Convert float32 slice to bytes for Redis RediSearch vector query
	vectorBytes := make([]byte, len(embedding)*4)
	for i, f := range embedding {
		bits := *(*uint32)(unsafe.Pointer(&f))
		idx := i * 4
		vectorBytes[idx] = byte(bits)
		vectorBytes[idx+1] = byte(bits >> 8)
		vectorBytes[idx+2] = byte(bits >> 16)
		vectorBytes[idx+3] = byte(bits >> 24)
	}

	// Construct KNN query: search top 1 nearest neighbor
	// Index must be configured with HNSW vector similarity
	query := fmt.Sprintf("*=>[KNN 1 @prompt_vector $vec AS distance]")
	res, err := s.redisClient.Do(ctx, "FT.SEARCH", "idx:semantic_cache", query, 
		"PARAMS", "2", "vec", vectorBytes, 
		"DIALECT", "2", 
		"RETURN", "2", "response", "distance",
	).Result()
	if err != nil {
		return "", false, err
	}

	results, ok := res.([]interface{})
	if !ok || len(results) < 2 {
		return "", false, nil
	}

	// RediSearch return format: [total_results, key_1, [field_1, val_1, ...]]
	numResults, _ := results[0].(int64)
	if numResults == 0 {
		return "", false, nil
	}

	fields, ok := results[2].([]interface{})
	if !ok {
		return "", false, nil
	}

	var response string
	var distance float64
	for i := 0; i < len(fields); i += 2 {
		key := fields[i].(string)
		val := fields[i+1].(string)
		if key == "response" {
			response = val
		} else if key == "distance" {
			fmt.Sscanf(val, "%f", &distance)
		}
	}

	// For cosine distance, similarity = 1 - distance
	similarity := 1.0 - distance
	if similarity >= s.threshold {
		return response, true, nil
	}

	return "", false, nil
}

snippet-3.go

package gateway

import (
	"math/rand"
	"time"
)

type PolicyEngine struct {
	models []ModelConfig
	epsilon float64 // Exploration rate (e.g., 0.05)
}

func NewPolicyEngine(models []ModelConfig, epsilon float64) *PolicyEngine {
	return &PolicyEngine{
		models:  models,
		epsilon: epsilon,
	}
}

// SelectRoute chooses the best model configuration based on the request constraints and policy weights.
func (p *PolicyEngine) SelectRoute(req RequestMetadata) RouteResult {
	// Filter models that can support the request's token requirements
	var candidates []ModelConfig
	for _, m := range p.models {
		if m.MaxTokens >= req.EstimatedTokens && m.QualityScore >= req.RequiredQuality {
			candidates = append(candidates, m)
		}
	}

	if len(candidates) == 0 {
		// Fallback to absolute default if no model meets criteria
		return RouteResult{PrimaryModel: p.models[0]}
	}

	// Epsilon-Greedy Exploration: Route to a random candidate 5% of the time to collect fresh metrics
	if rand.Float64() < p.epsilon {
		chosen := candidates[rand.Intn(len(candidates))]
		return RouteResult{
			PrimaryModel:  chosen,
			FallbackChain: p.buildFallbackChain(chosen, candidates),
		}
	}

	// Exploitation: Select the model maximizing the utility score
	var bestModel ModelConfig
	bestScore := -999999.0
	
	// Weights for routing prioritization
	const (
		wQuality = 10.0
		wCost    = 5.0
		wLatency = 4.0
	)

	for _, m := range candidates {
		// Compute normalized cost (estimate total cost for 1K tokens)
		estCost := (float64(req.EstimatedTokens)/1000.0) * (m.InputCostPerM / 1000.0)
		
		// Latency SLA penalty
		latencyPenalty := 0.0
		if m.P99Latency > req.LatencySLA {
			// Severe penalty if p99 latency breaks the requested SLA
			latencyPenalty = float64(m.P99Latency - req.LatencySLA) / float64(time.Second) * 15.0
		}

		// Utility function: maximize quality, minimize cost, minimize latency penalty
		score := (wQuality * m.QualityScore) - (wCost * estCost) - (wLatency * latencyPenalty)

		if score > bestScore {
			bestScore = score
			bestModel = m
		}
	}

	return RouteResult{
		PrimaryModel:  bestModel,
		FallbackChain: p.buildFallbackChain(bestModel, candidates),
	}
}

func (p *PolicyEngine) buildFallbackChain(primary ModelConfig, candidates []ModelConfig) []ModelConfig {
	var chain []ModelConfig
	for _, m := range candidates {
		if m.ID != primary.ID {
			chain = append(chain, m)
		}
	}
	// Sort chain by quality score descending
	return chain
}

snippet-4.go

package gateway

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net/http"
	"sync"
	"time"
)

type CircuitState int

const (
	StateClosed CircuitState = iota
	StateHalfOpen
	StateOpen
)

type CircuitBreaker struct {
	mu           sync.RWMutex
	state        CircuitState
	failures     int
	successes    int
	totalRequests int
	threshold    float64 // Error percentage threshold (e.g., 0.15)
	openUntil    time.Time
	cooldown     time.Duration
}

func NewCircuitBreaker(threshold float64, cooldown time.Duration) *CircuitBreaker {
	return &CircuitBreaker{
		state:     StateClosed,
		threshold: threshold,
		cooldown:  cooldown,
	}
}

func (cb *CircuitBreaker) AllowRequest() bool {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	if cb.state == StateOpen {
		if time.Now().After(cb.openUntil) {
			cb.state = StateHalfOpen
			return true
		}
		return false
	}
	return true
}

func (cb *CircuitBreaker) RecordResult(err error) {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	cb.totalRequests++
	if err != nil {
		cb.failures++
	} else {
		cb.successes++
	}

	if cb.totalRequests >= 20 { // Minimum sample size
		errRate := float64(cb.failures) / float64(cb.totalRequests)
		if cb.state == StateClosed && errRate >= cb.threshold {
			cb.state = StateOpen
			cb.openUntil = time.Now().Add(cb.cooldown)
			cb.failures = 0
			cb.successes = 0
			cb.totalRequests = 0
		} else if cb.state == StateHalfOpen {
			if err == nil && cb.successes >= 5 {
				cb.state = StateClosed
				cb.failures = 0
				cb.successes = 0
				cb.totalRequests = 0
			} else if err != nil {
				cb.state = StateOpen
				cb.openUntil = time.Now().Add(cb.cooldown)
			}
		}
	}
}

type FailoverProxyClient struct {
	client   *http.Client
	breakers map[string]*CircuitBreaker
	mu       sync.Mutex
}

func NewFailoverProxyClient(client *http.Client) *FailoverProxyClient {
	return &FailoverProxyClient{
		client:   client,
		breakers: make(map[string]*CircuitBreaker),
	}
}

func (f *FailoverProxyClient) ExecuteWithFailover(ctx context.Context, route RouteResult, payload []byte) ([]byte, error) {
	modelsToTry := append([]ModelConfig{route.PrimaryModel}, route.FallbackChain...)

	for _, model := range modelsToTry {
		f.mu.Lock()
		cb, exists := f.breakers[model.ID]
		if !exists {
			cb = NewCircuitBreaker(0.15, 10*time.Second)
			f.breakers[model.ID] = cb
		}
		f.mu.Unlock()

		if !cb.AllowRequest() {
			continue // Skip tripped model and try fallback
		}

		req, err := http.NewRequestWithContext(ctx, "POST", GetProviderURL(model), bytes.NewReader(payload))
		if err != nil {
			cb.RecordResult(err)
			continue
		}
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", GetAuthHeader(model))

		resp, err := f.client.Do(req)
		if err != nil {
			cb.RecordResult(err)
			continue
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			body, _ := io.ReadAll(resp.Body)
			cb.RecordResult(errors.New(string(body)))
			continue
		}

		respBody, err := io.ReadAll(resp.Body)
		if err != nil {
			cb.RecordResult(err)
			continue
		}

		cb.RecordResult(nil)
		return respBody, nil
	}

	return nil, errors.New("all downstream routing paths failed or were circuit-broken")
}

func GetProviderURL(m ModelConfig) string {
	if m.Provider == "openai" {
		return "https://api.openai.com/v1/chat/completions"
	}
	return "https://api.anthropic.com/v1/messages"
}

func GetAuthHeader(m ModelConfig) string {
	return "Bearer token_placeholder"
}

snippet-5.go

package gateway

import (
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

var (
	tokenCounter = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "llm_gateway_tokens_total",
			Help: "Total number of LLM tokens processed by model and type",
		},
		[]string{"model", "token_type"},
	)

	costCounter = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "llm_gateway_cost_usd_total",
			Help: "Calculated API expenditure in USD",
		},
		[]string{"model"},
	)

	latencyHistogram = promauto.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "llm_gateway_latency_seconds",
			Help:    "Response latency of upstream models in seconds",
			Buckets: []float64{0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.0},
		},
		[]string{"model"},
	)
)

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
}

type TelemetryPayload struct {
	ModelID      string
	Duration     time.Duration
	InputTokens  int
	OutputTokens int
	Error        bool
}

type TelemetryQueue struct {
	queueChan chan TelemetryPayload
}

func NewTelemetryQueue(bufferSize int) *TelemetryQueue {
	tq := &TelemetryQueue{
		queueChan: make(chan TelemetryPayload, bufferSize),
	}
	go tq.worker()
	return tq
}

func (tq *TelemetryQueue) Publish(payload TelemetryPayload) {
	select {
	case tq.queueChan <- payload:
	default:
		// Queue full, drop telemetry to prevent blocking the request path
	}
}

func (tq *TelemetryQueue) worker() {
	for payload := range tq.queueChan {
		if payload.Error {
			continue
		}

		// Log token counts to Prometheus
		tokenCounter.WithLabelValues(payload.ModelID, "input").Add(float64(payload.InputTokens))
		tokenCounter.WithLabelValues(payload.ModelID, "output").Add(float64(payload.OutputTokens))

		// Log latency
		latencyHistogram.WithLabelValues(payload.ModelID).Observe(payload.Duration.Seconds())

		// Calculate and log USD costs
		cost := calculateCost(payload.ModelID, payload.InputTokens, payload.OutputTokens)
		costCounter.WithLabelValues(payload.ModelID).Add(cost)

		// Persist performance to update Policy Engine state
		updateModelPerformanceState(payload.ModelID, payload.Duration, cost)
	}
}

func calculateCost(modelID string, input, output int) float64 {
	switch modelID {
	case "claude-3-5-sonnet":
		return (float64(input)/1_000_000.0)*3.0 + (float64(output)/1_000_000.0)*15.0
	case "gemini-1-5-flash":
		return (float64(input)/1_000_000.0)*0.075 + (float64(output)/1_000_000.0)*0.30
	default:
		return 0.0
	}
}

func updateModelPerformanceState(modelID string, latency time.Duration, cost float64) {
	// Write to Redis sorted set or hash to compute rolling averages
}

snippet-6.go

package gateway

import (
	"encoding/json"
	"fmt"
)

// UnifiedMessage is the system-agnostic chat message schema.
type UnifiedMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type UnifiedRequest struct {
	Messages    []UnifiedMessage `json:"messages"`
	Temperature float64          `json:"temperature"`
}

// Adapts unified schema to OpenAI Chat Completion structure
type OpenAIRequest struct {
	Model       string           `json:"model"`
	Messages    []UnifiedMessage `json:"messages"`
	Temperature float64          `json:"temperature"`
}

// Adapts unified schema to Anthropic Messages structure
type AnthropicRequest struct {
	Model       string           `json:"model"`
	System      string           `json:"system,omitempty"`
	Messages    []UnifiedMessage `json:"messages"`
	MaxTokens   int              `json:"max_tokens"`
	Temperature float64          `json:"temperature"`
}

func AdaptPayload(unified UnifiedRequest, modelConfig ModelConfig) ([]byte, error) {
	if modelConfig.Provider == "openai" {
		payload := OpenAIRequest{
			Model:       modelConfig.ID,
			Messages:    unified.Messages,
			Temperature: unified.Temperature,
		}
		return json.Marshal(payload)
	}

	if modelConfig.Provider == "anthropic" {
		var systemPrompt string
		var chatMessages []UnifiedMessage
		
		for _, msg := range unified.Messages {
			if msg.Role == "system" {
				systemPrompt = msg.Content
			} else {
				chatMessages = append(chatMessages, msg)
			}
		}

		payload := AnthropicRequest{
			Model:       modelConfig.ID,
			System:      systemPrompt,
			Messages:    chatMessages,
			MaxTokens:   4096,
			Temperature: unified.Temperature,
		}
		return json.Marshal(payload)
	}

	return nil, fmt.Errorf("unsupported provider: %s", modelConfig.Provider)
}

snippet-7.go

package gateway

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"time"
)

type GatewayHandler struct {
	policyEngine *PolicyEngine
	proxyClient  *FailoverProxyClient
	telemetry    *TelemetryQueue
}

func NewGatewayHandler(pe *PolicyEngine, pc *FailoverProxyClient, tq *TelemetryQueue) *GatewayHandler {
	return &GatewayHandler{
		policyEngine: pe,
		proxyClient:  pc,
		telemetry:    tq,
	}
}

func (gh *GatewayHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "Failed to read request body", http.StatusBadRequest)
		return
	}

	var unifiedReq UnifiedRequest
	if err := json.Unmarshal(body, &unifiedReq); err != nil {
		http.Error(w, "Invalid request JSON", http.StatusBadRequest)
		return
	}

	// Extract SLA from headers, default to 1500ms
	slaHeader := r.Header.Get("X-Latency-SLA-Ms")
	slaDuration := 1500 * time.Millisecond
	if slaHeader != "" {
		if d, err := time.ParseDuration(slaHeader + "ms"); err == nil {
			slaDuration = d
		}
	}

	// Apply heuristics to evaluate request metadata
	promptLength := len(unifiedReq.Messages[len(unifiedReq.Messages)-1].Content)
	reqMeta := RequestMetadata{
		Prompt:          unifiedReq.Messages[len(unifiedReq.Messages)-1].Content,
		EstimatedTokens: promptLength / 4,
		RequiredQuality: 0.70, // Baseline minimum quality threshold
		LatencySLA:      slaDuration,
	}

	// Run policy selection
	routeResult := gh.policyEngine.SelectRoute(reqMeta)

	// Adapt system agnostic request payload to model target
	adaptedPayload, err := AdaptPayload(unifiedReq, routeResult.PrimaryModel)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	startTime := time.Now()
	ctx, cancel := context.WithTimeout(r.Context(), slaDuration)
	defer cancel()

	// Execute requests with circuit breaking and fallback routes
	respBytes, err := gh.proxyClient.ExecuteWithFailover(ctx, routeResult, adaptedPayload)
	duration := time.Since(startTime)

	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		gh.telemetry.Publish(TelemetryPayload{
			ModelID: routeResult.PrimaryModel.ID,
			Error:   true,
		})
		return
	}

	// Parse usage asynchronously and publish telemetry
	gh.telemetry.Publish(TelemetryPayload{
		ModelID:      routeResult.PrimaryModel.ID,
		Duration:     duration,
		InputTokens:  reqMeta.EstimatedTokens,
		OutputTokens: 250, // Approximation; parse actual token counts in production
		Error:        false,
	})

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("X-Routed-Model", routeResult.PrimaryModel.ID)
	w.WriteHeader(http.StatusOK)
	w.Write(respBytes)
}