Luminary/internal/gateway/health_check.go

package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/provider"
)

func (s *Service) healthCheckEndpoints(p *model.Provider) []provider.EndpointType {
	var out []provider.EndpointType
	if s.providerAllowsEndpoint(p, provider.EndpointListModels) {
		out = append(out, provider.EndpointListModels)
	}
	switch p.Category {
	case model.CategoryEmbedding:
		if s.providerAllowsEndpoint(p, provider.EndpointEmbeddings) {
			out = append(out, provider.EndpointEmbeddings)
		}
	case model.CategoryRerank:
		if s.providerAllowsEndpoint(p, provider.EndpointRerank) {
			out = append(out, provider.EndpointRerank)
		}
	}
	return out
}

func (s *Service) runHealthChecks(
	ctx context.Context,
	p *model.Provider,
	keys []model.ProviderKey,
	cfg provider.ProviderConfig,
) (model.ProviderStatus, int64, string) {
	endpoints := s.healthCheckEndpoints(p)
	if len(endpoints) == 0 {
		return model.ProviderUnhealthy, 0, "no health check endpoint enabled"
	}
	start := time.Now()
	var lastErr string
	for _, ep := range endpoints {
		status, _, msg := s.runHealthCheck(ctx, p, keys, cfg, ep)
		if status == model.ProviderHealthy {
			return model.ProviderHealthy, time.Since(start).Milliseconds(), ""
		}
		if msg != "" {
			lastErr = msg
		}
	}
	if lastErr == "" {
		lastErr = "health check failed"
	}
	return model.ProviderUnhealthy, time.Since(start).Milliseconds(), lastErr
}

func (s *Service) runHealthCheck(
	ctx context.Context,
	p *model.Provider,
	keys []model.ProviderKey,
	cfg provider.ProviderConfig,
	endpoint provider.EndpointType,
) (model.ProviderStatus, int64, string) {
	start := time.Now()
	latency := func() int64 { return time.Since(start).Milliseconds() }

	switch endpoint {
	case provider.EndpointListModels:
		if len(keys) == 0 && p.Type == model.ProviderOllama {
			if err := s.openai.HealthCheckWithConfig(ctx, "", cfg); err == nil {
				return model.ProviderHealthy, latency(), ""
			}
			return model.ProviderUnhealthy, latency(), "health check failed"
		}
		for _, k := range keys {
			if !k.Enabled {
				continue
			}
			apiKey, err := s.DecryptAPIKey(k.APIKeyEnc)
			if err != nil {
				continue
			}
			var checkErr error
			if p.Type == model.ProviderAnthropic {
				checkErr = s.anthropic.HealthCheckWithConfig(ctx, apiKey, cfg)
			} else {
				checkErr = s.openai.HealthCheckWithConfig(ctx, apiKey, cfg)
			}
			if checkErr == nil {
				return model.ProviderHealthy, latency(), ""
			}
		}
		return model.ProviderUnhealthy, latency(), "health check failed"

	case provider.EndpointEmbeddings, provider.EndpointRerank:
		body, err := healthProbeBody(s, p, endpoint)
		if err != nil {
			return model.ProviderUnhealthy, latency(), err.Error()
		}
		if len(keys) == 0 && p.Type == model.ProviderOllama {
			ok, msg := s.probeEndpoint(ctx, "", cfg, endpoint, body)
			if ok {
				return model.ProviderHealthy, latency(), ""
			}
			return model.ProviderUnhealthy, latency(), msg
		}
		var lastErr string
		for _, k := range keys {
			if !k.Enabled {
				continue
			}
			apiKey, err := s.DecryptAPIKey(k.APIKeyEnc)
			if err != nil {
				continue
			}
			if ok, msg := s.probeEndpoint(ctx, apiKey, cfg, endpoint, body); ok {
				return model.ProviderHealthy, latency(), ""
			} else if msg != "" {
				lastErr = msg
			}
		}
		if lastErr == "" {
			lastErr = "health check failed"
		}
		return model.ProviderUnhealthy, latency(), lastErr
	default:
		return model.ProviderUnhealthy, latency(), "unsupported health check endpoint"
	}
}

func (s *Service) probeEndpoint(
	ctx context.Context,
	apiKey string,
	cfg provider.ProviderConfig,
	endpoint provider.EndpointType,
	body []byte,
) (bool, string) {
	resp, err := s.openai.ForwardWithConfig(ctx, apiKey, cfg, endpoint, body)
	if err != nil {
		return false, err.Error()
	}
	if resp == nil {
		return false, "empty response"
	}
	if isHealthCheckHTTPStatus(resp.StatusCode) {
		return true, ""
	}
	return false, fmt.Sprintf("%s: HTTP %d", endpoint, resp.StatusCode)
}

func healthProbeBody(s *Service, p *model.Provider, endpoint provider.EndpointType) ([]byte, error) {
	modelID := probeModelID(s, p.ID, endpoint)
	switch endpoint {
	case provider.EndpointEmbeddings:
		return json.Marshal(map[string]string{"model": modelID, "input": "."})
	case provider.EndpointRerank:
		return json.Marshal(map[string]any{
			"model":     modelID,
			"query":     "ping",
			"documents": []string{"ping"},
		})
	default:
		return nil, fmt.Errorf("unsupported probe endpoint %s", endpoint)
	}
}

func probeModelID(s *Service, providerID uint, endpoint provider.EndpointType) string {
	for _, m := range s.cache.GetModels(providerID) {
		if m.Enabled && m.ModelID != "" {
			return m.ModelID
		}
	}
	switch endpoint {
	case provider.EndpointEmbeddings:
		return "text-embedding-3-small"
	case provider.EndpointRerank:
		return "rerank-multilingual-v3.0"
	default:
		return "health-check"
	}
}

func isHealthCheckHTTPStatus(code int) bool {
	if code >= 200 && code < 300 {
		return true
	}
	// Upstream reachable but rejected payload/model — still counts as healthy.
	return code == 400 || code == 422
}