ce7cb702218693d08ea6d83464d2457c9250027c

services/agentgateway-model-catalog-sync/main.go

@@ -0,0 +1,729 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+	"time"
+)
+
+const (
+	modelsDevURL             = "https://models.dev/api.json"
+	novitaModelsURL          = "https://api.novita.ai/openai/v1/models"
+	configMapName            = "agentgateway-model-costs"
+	serviceAccountDir        = "/var/run/secrets/kubernetes.io/serviceaccount"
+	defaultSyncInterval      = "6h"
+	defaultSyncRetryInterval = "5m"
+	defaultEnabledProviders  = "google-vertex,novita"
+)
+
+type modelsDevCatalog map[string]modelsDevProvider
+
+type modelsDevProvider struct {
+	Models map[string]modelsDevModel `json:"models"`
+}
+
+type modelsDevModel struct {
+	Cost modelsDevCost `json:"cost"`
+}
+
+type modelsDevCost struct {
+	Input       *json.Number    `json:"input"`
+	Output      *json.Number    `json:"output"`
+	CacheRead   *json.Number    `json:"cache_read"`
+	CacheWrite  *json.Number    `json:"cache_write"`
+	Reasoning   *json.Number    `json:"reasoning"`
+	InputAudio  *json.Number    `json:"input_audio"`
+	OutputAudio *json.Number    `json:"output_audio"`
+	Tiers       []modelsDevTier `json:"tiers"`
+}
+
+type modelsDevTier struct {
+	Input       *json.Number          `json:"input"`
+	Output      *json.Number          `json:"output"`
+	CacheRead   *json.Number          `json:"cache_read"`
+	CacheWrite  *json.Number          `json:"cache_write"`
+	Reasoning   *json.Number          `json:"reasoning"`
+	InputAudio  *json.Number          `json:"input_audio"`
+	OutputAudio *json.Number          `json:"output_audio"`
+	Tier        modelsDevTierSelector `json:"tier"`
+}
+
+type modelsDevTierSelector struct {
+	Type string `json:"type"`
+	Size int    `json:"size"`
+}
+
+type novitaCatalog struct {
+	Data []novitaModel `json:"data"`
+}
+
+type novitaModel struct {
+	ID          string      `json:"id"`
+	InputPrice  json.Number `json:"input_token_price_per_m"`
+	OutputPrice json.Number `json:"output_token_price_per_m"`
+}
+
+type catalog struct {
+	Providers map[string]catalogProvider `json:"providers"`
+}
+
+type catalogProvider struct {
+	Models map[string]catalogModel `json:"models"`
+}
+
+type catalogModel struct {
+	Rates map[string]string `json:"rates"`
+	Tiers []catalogTier     `json:"tiers,omitempty"`
+}
+
+type catalogTier struct {
+	ContextOver int               `json:"contextOver"`
+	Rates       map[string]string `json:"rates"`
+}
+
+type configMapPatch struct {
+	Data map[string]string `json:"data"`
+}
+
+type configMapResource struct {
+	APIVersion string            `json:"apiVersion"`
+	Kind       string            `json:"kind"`
+	Metadata   configMapMetadata `json:"metadata"`
+	Data       map[string]string `json:"data"`
+}
+
+type configMapMetadata struct {
+	Name        string            `json:"name"`
+	Annotations map[string]string `json:"annotations"`
+	Labels      map[string]string `json:"labels"`
+}
+
+type modelList struct {
+	Object string          `json:"object"`
+	Data   []modelListItem `json:"data"`
+}
+
+type modelListItem struct {
+	ID      string `json:"id"`
+	Object  string `json:"object"`
+	Created int64  `json:"created"`
+	OwnedBy string `json:"owned_by"`
+}
+
+type providerDefinition struct {
+	CatalogID      string
+	ModelsDevID    string
+	ModelPrefix    string
+	SourceProvider string
+	Novita         bool
+}
+
+type providerSelection map[string]providerDefinition
+
+var modelsDevProviderIDs = map[string]string{
+	"amazon-bedrock": "aws.bedrock",
+	"anthropic":      "anthropic",
+	"azure":          "azure",
+	"baseten":        "baseten",
+	"cerebras":       "cerebras",
+	"cohere":         "cohere",
+	"deepinfra":      "deepinfra",
+	"deepseek":       "deepseek",
+	"fireworks-ai":   "fireworks",
+	"github-copilot": "copilot",
+	"google":         "gcp.gemini",
+	"google-vertex":  "gcp.vertex_ai",
+	"groq":           "groq",
+	"huggingface":    "huggingface",
+	"mistral":        "mistral",
+	"openai":         "openai",
+	"openrouter":     "openrouter",
+	"togetherai":     "togetherai",
+	"xai":            "xai",
+}
+
+func main() {
+	var err error
+	if len(os.Args) != 2 {
+		err = errors.New("expected sync or run")
+	} else {
+		var providers providerSelection
+		providers, err = providersFromEnvironment()
+		if err == nil {
+			switch os.Args[1] {
+			case "sync":
+				err = syncCatalog(context.Background(), providers)
+			case "run":
+				err = run(context.Background(), providers)
+			default:
+				err = fmt.Errorf("unknown command %q", os.Args[1])
+			}
+		}
+	}
+	if err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+}
+
+func providersFromEnvironment() (providerSelection, error) {
+	value := os.Getenv("ENABLED_PROVIDERS")
+	if value == "" {
+		value = defaultEnabledProviders
+	}
+	definitions := buildProviderDefinitions()
+	providers := providerSelection{}
+	for _, name := range strings.Split(value, ",") {
+		name = strings.TrimSpace(name)
+		definition, found := definitions[name]
+		if !found {
+			return nil, fmt.Errorf("unsupported provider %q", name)
+		}
+		if _, found := providers[name]; found {
+			return nil, fmt.Errorf("duplicate provider %q", name)
+		}
+		providers[name] = definition
+	}
+	if len(providers) == 0 {
+		return nil, errors.New("ENABLED_PROVIDERS must select at least one provider")
+	}
+	return providers, nil
+}
+
+func buildProviderDefinitions() map[string]providerDefinition {
+	definitions := make(map[string]providerDefinition, len(modelsDevProviderIDs)+1)
+	for sourceID, catalogID := range modelsDevProviderIDs {
+		prefix := catalogID
+		if separator := strings.LastIndexByte(prefix, '.'); separator >= 0 {
+			prefix = prefix[separator+1:]
+		}
+		definitions[sourceID] = providerDefinition{
+			CatalogID:      catalogID,
+			ModelsDevID:    sourceID,
+			ModelPrefix:    prefix + "/",
+			SourceProvider: sourceID,
+		}
+	}
+	definitions["novita"] = providerDefinition{
+		CatalogID:      "custom",
+		ModelPrefix:    "novita/",
+		SourceProvider: "novita",
+		Novita:         true,
+	}
+	return definitions
+}
+
+func syncCatalog(ctx context.Context, providers providerSelection) error {
+	namespace := os.Getenv("POD_NAMESPACE")
+	if namespace == "" {
+		return errors.New("POD_NAMESPACE is required")
+	}
+
+	client := &http.Client{Timeout: 30 * time.Second}
+	var upstream modelsDevCatalog
+	for _, provider := range providers {
+		if provider.ModelsDevID != "" {
+			if err := fetchJSON(ctx, client, modelsDevURL, "", &upstream); err != nil {
+				return fmt.Errorf("fetch models.dev: %w", err)
+			}
+			break
+		}
+	}
+	var novita novitaCatalog
+	if _, enabled := providers["novita"]; enabled {
+		novitaAPIKey := os.Getenv("NOVITA_API_KEY")
+		if novitaAPIKey == "" {
+			return errors.New("NOVITA_API_KEY is required when Novita is enabled")
+		}
+		if err := fetchJSON(ctx, client, novitaModelsURL, "Bearer "+novitaAPIKey, &novita); err != nil {
+			return fmt.Errorf("fetch Novita: %w", err)
+		}
+	}
+
+	output, err := buildCatalog(upstream, novita, providers)
+	if err != nil {
+		return err
+	}
+	encoded, err := json.Marshal(output)
+	if err != nil {
+		return fmt.Errorf("encode catalog: %w", err)
+	}
+	return updateConfigMap(ctx, namespace, encoded)
+}
+
+func run(ctx context.Context, providers providerSelection) error {
+	syncInterval, err := durationFromEnvironment("SYNC_INTERVAL", defaultSyncInterval)
+	if err != nil {
+		return err
+	}
+	syncRetryInterval, err := durationFromEnvironment("SYNC_RETRY_INTERVAL", defaultSyncRetryInterval)
+	if err != nil {
+		return err
+	}
+	go syncOnSchedule(ctx, syncInterval, syncRetryInterval, providers)
+	return serveCatalog(providers)
+}
+
+func durationFromEnvironment(name, defaultValue string) (time.Duration, error) {
+	value := os.Getenv(name)
+	if value == "" {
+		value = defaultValue
+	}
+	duration, err := time.ParseDuration(value)
+	if err != nil {
+		return 0, fmt.Errorf("parse %s: %w", name, err)
+	}
+	if duration <= 0 {
+		return 0, fmt.Errorf("%s must be positive", name)
+	}
+	return duration, nil
+}
+
+func syncOnSchedule(
+	ctx context.Context,
+	syncInterval time.Duration,
+	syncRetryInterval time.Duration,
+	providers providerSelection,
+) {
+	delay := time.Duration(0)
+	for {
+		timer := time.NewTimer(delay)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			return
+		case <-timer.C:
+		}
+		if err := syncCatalog(ctx, providers); err != nil {
+			fmt.Fprintln(os.Stderr, err)
+			delay = syncRetryInterval
+		} else {
+			delay = syncInterval
+		}
+	}
+}
+
+func serveCatalog(providers providerSelection) error {
+	server := &http.Server{
+		Addr:              ":8080",
+		Handler:           catalogHandler("/catalog/catalog.json", providers),
+		ReadHeaderTimeout: 5 * time.Second,
+		IdleTimeout:       60 * time.Second,
+	}
+	return server.ListenAndServe()
+}
+
+func catalogHandler(path string, providers providerSelection) http.Handler {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/healthz", func(response http.ResponseWriter, _ *http.Request) {
+		response.WriteHeader(http.StatusNoContent)
+	})
+	mux.HandleFunc("/readyz", func(response http.ResponseWriter, _ *http.Request) {
+		input, err := readCatalog(path)
+		if err != nil || !catalogReady(input, providers) {
+			http.Error(response, "catalog invalid", http.StatusServiceUnavailable)
+			return
+		}
+		response.WriteHeader(http.StatusNoContent)
+	})
+	mux.HandleFunc("/metrics", func(response http.ResponseWriter, _ *http.Request) {
+		input, err := readCatalog(path)
+		if err != nil {
+			http.Error(response, "catalog invalid", http.StatusServiceUnavailable)
+			return
+		}
+		response.Header().Set("Content-Type", "text/plain; version=0.0.4")
+		_, _ = response.Write(renderMetrics(input))
+	})
+	mux.HandleFunc("/v1/models", func(response http.ResponseWriter, request *http.Request) {
+		if request.Method != http.MethodGet {
+			response.WriteHeader(http.StatusMethodNotAllowed)
+			return
+		}
+		input, err := readCatalog(path)
+		if err != nil {
+			http.Error(response, "catalog invalid", http.StatusServiceUnavailable)
+			return
+		}
+		response.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(response).Encode(renderModelList(input))
+	})
+	return mux
+}
+
+func catalogReady(input catalog, providers providerSelection) bool {
+	for _, provider := range providers {
+		if len(input.Providers[provider.CatalogID].Models) == 0 {
+			return false
+		}
+	}
+	return true
+}
+
+func readCatalog(path string) (catalog, error) {
+	contents, err := os.ReadFile(path)
+	if err != nil {
+		return catalog{}, err
+	}
+	input := catalog{}
+	if err := json.Unmarshal(contents, &input); err != nil {
+		return catalog{}, err
+	}
+	return input, nil
+}
+
+func renderMetrics(input catalog) []byte {
+	var output strings.Builder
+	output.WriteString("# HELP agentgateway_model_token_price_usd_per_million Model token price in USD per million tokens.\n")
+	output.WriteString("# TYPE agentgateway_model_token_price_usd_per_million gauge\n")
+	providers := sortedKeys(input.Providers)
+	for _, providerID := range providers {
+		provider := input.Providers[providerID]
+		for _, modelName := range sortedKeys(provider.Models) {
+			model := provider.Models[modelName]
+			for _, tokenType := range sortedKeys(model.Rates) {
+				fmt.Fprintf(
+					&output,
+					"agentgateway_model_token_price_usd_per_million{gen_ai_system=\"%s\",gen_ai_request_model=\"%s\",gen_ai_token_type=\"%s\",source_provider=\"%s\"} %s\n",
+					escapeLabel(providerID),
+					escapeLabel(modelName),
+					escapeLabel(metricTokenType(tokenType)),
+					escapeLabel(sourceProvider(providerID)),
+					model.Rates[tokenType],
+				)
+			}
+		}
+	}
+	return []byte(output.String())
+}
+
+func renderModelList(input catalog) modelList {
+	output := modelList{Object: "list"}
+	for _, providerID := range sortedKeys(input.Providers) {
+		provider := input.Providers[providerID]
+		for _, modelName := range sortedKeys(provider.Models) {
+			output.Data = append(output.Data, modelListItem{
+				ID:      modelPrefix(providerID) + modelName,
+				Object:  "model",
+				Created: 0,
+				OwnedBy: sourceProvider(providerID),
+			})
+		}
+	}
+	return output
+}
+
+func sortedKeys[Value any](values map[string]Value) []string {
+	keys := make([]string, 0, len(values))
+	for key := range values {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	return keys
+}
+
+func escapeLabel(value string) string {
+	return strings.NewReplacer("\\", "\\\\", "\n", "\\n", "\"", "\\\"").Replace(value)
+}
+
+func metricTokenType(value string) string {
+	replacements := map[string]string{
+		"cacheRead":   "cache_read",
+		"cacheWrite":  "cache_write",
+		"inputAudio":  "input_audio",
+		"outputAudio": "output_audio",
+	}
+	if replacement, found := replacements[value]; found {
+		return replacement
+	}
+	return value
+}
+
+func sourceProvider(providerID string) string {
+	for _, provider := range buildProviderDefinitions() {
+		if provider.CatalogID == providerID {
+			return provider.SourceProvider
+		}
+	}
+	return providerID
+}
+
+func modelPrefix(providerID string) string {
+	for _, provider := range buildProviderDefinitions() {
+		if provider.CatalogID == providerID {
+			return provider.ModelPrefix
+		}
+	}
+	return ""
+}
+
+func fetchJSON(ctx context.Context, client *http.Client, url, authorization string, target any) error {
+	request, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return err
+	}
+	request.Header.Set("Accept", "application/json")
+	if authorization != "" {
+		request.Header.Set("Authorization", authorization)
+	}
+	response, err := client.Do(request)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		_ = response.Body.Close()
+	}()
+	if response.StatusCode != http.StatusOK {
+		_, _ = io.Copy(io.Discard, response.Body)
+		return fmt.Errorf("unexpected HTTP status %s", response.Status)
+	}
+	decoder := json.NewDecoder(io.LimitReader(response.Body, 16<<20))
+	decoder.UseNumber()
+	return decoder.Decode(target)
+}
+
+func buildCatalog(
+	upstream modelsDevCatalog,
+	novita novitaCatalog,
+	enabledProviders providerSelection,
+) (catalog, error) {
+	providers := make(map[string]catalogProvider, len(enabledProviders))
+	for name, definition := range enabledProviders {
+		var models map[string]catalogModel
+		if definition.Novita {
+			var err error
+			models, err = convertNovita(novita)
+			if err != nil {
+				return catalog{}, fmt.Errorf("convert %s catalog: %w", name, err)
+			}
+		} else {
+			models = convertModelsDev(upstream[definition.ModelsDevID])
+		}
+		if len(models) == 0 {
+			return catalog{}, fmt.Errorf("provider %s returned no priced models", name)
+		}
+		providers[definition.CatalogID] = catalogProvider{Models: models}
+	}
+	return catalog{Providers: providers}, nil
+}
+
+func convertModelsDev(provider modelsDevProvider) map[string]catalogModel {
+	models := make(map[string]catalogModel)
+	for name, model := range provider.Models {
+		rates := ratesFromModelsDevCost(model.Cost)
+		if len(rates) == 0 {
+			continue
+		}
+		converted := catalogModel{Rates: rates}
+		for _, tier := range model.Cost.Tiers {
+			if tier.Tier.Type != "context" || tier.Tier.Size <= 0 {
+				continue
+			}
+			tierRates := ratesFromModelsDevTier(tier)
+			if len(tierRates) > 0 {
+				converted.Tiers = append(converted.Tiers, catalogTier{ContextOver: tier.Tier.Size, Rates: tierRates})
+			}
+		}
+		sort.Slice(converted.Tiers, func(left, right int) bool {
+			return converted.Tiers[left].ContextOver < converted.Tiers[right].ContextOver
+		})
+		models[name] = converted
+	}
+	return models
+}
+
+func ratesFromModelsDevCost(cost modelsDevCost) map[string]string {
+	return ratesFromNumbers(cost.Input, cost.Output, cost.CacheRead, cost.CacheWrite, cost.Reasoning, cost.InputAudio, cost.OutputAudio)
+}
+
+func ratesFromModelsDevTier(tier modelsDevTier) map[string]string {
+	return ratesFromNumbers(tier.Input, tier.Output, tier.CacheRead, tier.CacheWrite, tier.Reasoning, tier.InputAudio, tier.OutputAudio)
+}
+
+func ratesFromNumbers(input, output, cacheRead, cacheWrite, reasoning, inputAudio, outputAudio *json.Number) map[string]string {
+	values := []struct {
+		name  string
+		value *json.Number
+	}{
+		{"input", input},
+		{"output", output},
+		{"cacheRead", cacheRead},
+		{"cacheWrite", cacheWrite},
+		{"reasoning", reasoning},
+		{"inputAudio", inputAudio},
+		{"outputAudio", outputAudio},
+	}
+	rates := make(map[string]string)
+	for _, value := range values {
+		if value.value != nil {
+			rates[value.name] = value.value.String()
+		}
+	}
+	return rates
+}
+
+func convertNovita(upstream novitaCatalog) (map[string]catalogModel, error) {
+	models := make(map[string]catalogModel)
+	for _, model := range upstream.Data {
+		if model.ID == "" || model.InputPrice == "" || model.OutputPrice == "" {
+			continue
+		}
+		input, err := tenThousandths(model.InputPrice)
+		if err != nil {
+			return nil, fmt.Errorf("model %q input price: %w", model.ID, err)
+		}
+		output, err := tenThousandths(model.OutputPrice)
+		if err != nil {
+			return nil, fmt.Errorf("model %q output price: %w", model.ID, err)
+		}
+		models[model.ID] = catalogModel{Rates: map[string]string{"input": input, "output": output}}
+	}
+	return models, nil
+}
+
+func tenThousandths(value json.Number) (string, error) {
+	raw := value.String()
+	if strings.ContainsAny(raw, ".eE-") {
+		return "", fmt.Errorf("expected a non-negative integer, got %q", raw)
+	}
+	raw = strings.TrimLeft(raw, "0")
+	if raw == "" {
+		return "0", nil
+	}
+	if len(raw) <= 4 {
+		raw = strings.Repeat("0", 4-len(raw)+1) + raw
+	}
+	integer := raw[:len(raw)-4]
+	fraction := strings.TrimRight(raw[len(raw)-4:], "0")
+	if fraction == "" {
+		return integer, nil
+	}
+	return integer + "." + fraction, nil
+}
+
+func updateConfigMap(ctx context.Context, namespace string, catalogJSON []byte) error {
+	token, err := os.ReadFile(filepath.Join(serviceAccountDir, "token"))
+	if err != nil {
+		return fmt.Errorf("read service account token: %w", err)
+	}
+	certificate, err := os.ReadFile(filepath.Join(serviceAccountDir, "ca.crt"))
+	if err != nil {
+		return fmt.Errorf("read Kubernetes CA: %w", err)
+	}
+	pool := x509.NewCertPool()
+	if !pool.AppendCertsFromPEM(certificate) {
+		return errors.New("parse Kubernetes CA")
+	}
+	client := &http.Client{
+		Timeout: 30 * time.Second,
+		Transport: &http.Transport{TLSClientConfig: &tls.Config{
+			MinVersion: tls.VersionTLS12,
+			RootCAs:    pool,
+		}},
+	}
+	url := "https://kubernetes.default.svc/api/v1/namespaces/" + namespace + "/configmaps"
+	return writeConfigMap(ctx, client, url, token, catalogJSON)
+}
+
+func writeConfigMap(ctx context.Context, client *http.Client, collectionURL string, token, catalogJSON []byte) error {
+	patch := configMapPatch{
+		Data: map[string]string{"catalog.json": string(catalogJSON)},
+	}
+	itemURL := collectionURL + "/" + configMapName
+	status, err := kubernetesRequest(ctx, client, http.MethodPatch, itemURL, token, patch, nil)
+	if err != nil {
+		return fmt.Errorf("patch ConfigMap: %w", err)
+	}
+	if status == http.StatusOK {
+		return nil
+	}
+	if status != http.StatusNotFound {
+		return fmt.Errorf("patch ConfigMap: unexpected HTTP status %d", status)
+	}
+
+	resource := configMapResource{
+		APIVersion: "v1",
+		Kind:       "ConfigMap",
+		Metadata: configMapMetadata{
+			Name: configMapName,
+			Annotations: map[string]string{
+				"kustomize.toolkit.fluxcd.io/prune": "disabled",
+			},
+			Labels: map[string]string{
+				"app.kubernetes.io/instance": "agentgateway",
+				"app.kubernetes.io/name":     "model-catalog",
+				"app.kubernetes.io/part-of":  "agentgateway",
+			},
+		},
+		Data: patch.Data,
+	}
+	status, err = kubernetesRequest(ctx, client, http.MethodPost, collectionURL, token, resource, nil)
+	if err != nil {
+		return fmt.Errorf("create ConfigMap: %w", err)
+	}
+	if status == http.StatusCreated {
+		return nil
+	}
+	if status != http.StatusConflict {
+		return fmt.Errorf("create ConfigMap: unexpected HTTP status %d", status)
+	}
+
+	status, err = kubernetesRequest(ctx, client, http.MethodPatch, itemURL, token, patch, nil)
+	if err != nil {
+		return fmt.Errorf("patch ConfigMap after create conflict: %w", err)
+	}
+	if status != http.StatusOK {
+		return fmt.Errorf("patch ConfigMap after create conflict: unexpected HTTP status %d", status)
+	}
+	return nil
+}
+
+func kubernetesRequest(ctx context.Context, client *http.Client, method, url string, token []byte, body any, target any) (int, error) {
+	var encoded io.Reader
+	if body != nil {
+		contents, err := json.Marshal(body)
+		if err != nil {
+			return 0, err
+		}
+		encoded = bytes.NewReader(contents)
+	}
+	request, err := http.NewRequestWithContext(ctx, method, url, encoded)
+	if err != nil {
+		return 0, err
+	}
+	request.Header.Set("Authorization", "Bearer "+string(token))
+	request.Header.Set("Accept", "application/json")
+	if body != nil {
+		contentType := "application/json"
+		if method == http.MethodPatch {
+			contentType = "application/merge-patch+json"
+		}
+		request.Header.Set("Content-Type", contentType)
+	}
+	response, err := client.Do(request)
+	if err != nil {
+		return 0, err
+	}
+	defer func() {
+		_ = response.Body.Close()
+	}()
+	if target != nil && response.StatusCode >= 200 && response.StatusCode < 300 {
+		if err := json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(target); err != nil {
+			return response.StatusCode, err
+		}
+	} else {
+		_, _ = io.Copy(io.Discard, response.Body)
+	}
+	return response.StatusCode, nil
+}