ce7cb702218693d08ea6d83464d2457c9250027c
This commit is contained in:
commit
34c5f8963c
10 changed files with 145336 additions and 0 deletions
3
services/agentgateway-model-catalog-sync/go.mod
Normal file
3
services/agentgateway-model-catalog-sync/go.mod
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
module github.com/quorate/ms2k-internal/services/agentgateway-model-catalog-sync
|
||||
|
||||
go 1.26.4
|
||||
13
services/agentgateway-model-catalog-sync/k8s/configmap.yaml
Normal file
13
services/agentgateway-model-catalog-sync/k8s/configmap.yaml
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
apiVersion: v1
|
||||
kind: ConfigMap
|
||||
metadata:
|
||||
name: agentgateway-model-costs
|
||||
namespace: agentgateway-system
|
||||
annotations:
|
||||
kustomize.toolkit.fluxcd.io/ssa: IfNotPresent
|
||||
labels:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
app.kubernetes.io/part-of: agentgateway
|
||||
data:
|
||||
catalog.json: '{"providers":{}}'
|
||||
100
services/agentgateway-model-catalog-sync/k8s/deployment.yaml
Normal file
100
services/agentgateway-model-catalog-sync/k8s/deployment.yaml
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: model-catalog
|
||||
namespace: agentgateway-system
|
||||
labels:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
app.kubernetes.io/part-of: agentgateway
|
||||
app.kubernetes.io/version: "0.1.1"
|
||||
spec:
|
||||
replicas: 2
|
||||
strategy:
|
||||
type: RollingUpdate
|
||||
rollingUpdate:
|
||||
maxSurge: 0
|
||||
maxUnavailable: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
app.kubernetes.io/part-of: agentgateway
|
||||
app.kubernetes.io/version: "0.1.1"
|
||||
spec:
|
||||
automountServiceAccountToken: true
|
||||
serviceAccountName: model-catalog
|
||||
securityContext:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65532
|
||||
runAsGroup: 65532
|
||||
seccompProfile:
|
||||
type: RuntimeDefault
|
||||
topologySpreadConstraints:
|
||||
- maxSkew: 1
|
||||
topologyKey: kubernetes.io/hostname
|
||||
whenUnsatisfiable: DoNotSchedule
|
||||
labelSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
containers:
|
||||
- name: catalog
|
||||
image: registry.quorate.internal/quorate/agentgateway-catalog:0.1.1@sha256:919e63ef02b9d6f40efa613570a5cb04af40764e9f377da4db4f58b4c3f23389
|
||||
imagePullPolicy: IfNotPresent
|
||||
args:
|
||||
- run
|
||||
env:
|
||||
- name: POD_NAMESPACE
|
||||
valueFrom:
|
||||
fieldRef:
|
||||
fieldPath: metadata.namespace
|
||||
- name: ENABLED_PROVIDERS
|
||||
value: google-vertex,novita
|
||||
- name: NOVITA_API_KEY
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: agentgateway-novita
|
||||
key: api_key
|
||||
- name: SYNC_INTERVAL
|
||||
value: 6h
|
||||
- name: SYNC_RETRY_INTERVAL
|
||||
value: 5m
|
||||
ports:
|
||||
- name: metrics
|
||||
containerPort: 8080
|
||||
protocol: TCP
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /readyz
|
||||
port: metrics
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /healthz
|
||||
port: metrics
|
||||
resources:
|
||||
requests:
|
||||
cpu: 10m
|
||||
memory: 16Mi
|
||||
limits:
|
||||
memory: 64Mi
|
||||
securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop:
|
||||
- ALL
|
||||
volumeMounts:
|
||||
- name: catalog
|
||||
mountPath: /catalog
|
||||
readOnly: true
|
||||
volumes:
|
||||
- name: catalog
|
||||
configMap:
|
||||
name: agentgateway-model-costs
|
||||
optional: true
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
apiVersion: kustomize.config.k8s.io/v1beta1
|
||||
kind: Kustomization
|
||||
namespace: agentgateway-system
|
||||
resources:
|
||||
- service-account.yaml
|
||||
- configmap.yaml
|
||||
- deployment.yaml
|
||||
- service.yaml
|
||||
- poddisruptionbudget.yaml
|
||||
- network-policy.yaml
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
apiVersion: cilium.io/v2
|
||||
kind: CiliumNetworkPolicy
|
||||
metadata:
|
||||
name: model-catalog
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
endpointSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
ingress:
|
||||
- fromEntities:
|
||||
- host
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "8080"
|
||||
protocol: TCP
|
||||
- fromEndpoints:
|
||||
- matchLabels:
|
||||
k8s:io.kubernetes.pod.namespace: victoriametrics
|
||||
app.kubernetes.io/name: vmagent
|
||||
app.kubernetes.io/instance: main
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "8080"
|
||||
protocol: TCP
|
||||
- fromEndpoints:
|
||||
- matchLabels:
|
||||
k8s:io.kubernetes.pod.namespace: agentgateway-system
|
||||
gateway.networking.k8s.io/gateway-name: agentgateway-proxy
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "8080"
|
||||
protocol: TCP
|
||||
egress:
|
||||
- toEntities:
|
||||
- kube-apiserver
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "6443"
|
||||
protocol: TCP
|
||||
- toEndpoints:
|
||||
- matchLabels:
|
||||
k8s:io.kubernetes.pod.namespace: kube-system
|
||||
k8s-app: kube-dns
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "53"
|
||||
protocol: UDP
|
||||
- port: "53"
|
||||
protocol: TCP
|
||||
- toEntities:
|
||||
- world
|
||||
toPorts:
|
||||
- serverNames:
|
||||
- models.dev
|
||||
- api.novita.ai
|
||||
ports:
|
||||
- port: "443"
|
||||
protocol: TCP
|
||||
---
|
||||
apiVersion: cilium.io/v2
|
||||
kind: CiliumClusterwideNetworkPolicy
|
||||
metadata:
|
||||
name: host-egress-agentgateway-model-catalog
|
||||
spec:
|
||||
nodeSelector: {}
|
||||
egress:
|
||||
- toEndpoints:
|
||||
- matchLabels:
|
||||
k8s:io.kubernetes.pod.namespace: agentgateway-system
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
toPorts:
|
||||
- ports:
|
||||
- port: "8080"
|
||||
protocol: TCP
|
||||
|
|
@ -0,0 +1,11 @@
|
|||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: model-catalog
|
||||
namespace: agentgateway-system
|
||||
spec:
|
||||
minAvailable: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
apiVersion: v1
|
||||
kind: ServiceAccount
|
||||
metadata:
|
||||
name: model-catalog
|
||||
namespace: agentgateway-system
|
||||
automountServiceAccountToken: true
|
||||
imagePullSecrets:
|
||||
- name: quay-cluster-pull
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: Role
|
||||
metadata:
|
||||
name: model-catalog
|
||||
namespace: agentgateway-system
|
||||
rules:
|
||||
- apiGroups:
|
||||
- ""
|
||||
resources:
|
||||
- configmaps
|
||||
verbs:
|
||||
- create
|
||||
- apiGroups:
|
||||
- ""
|
||||
resources:
|
||||
- configmaps
|
||||
resourceNames:
|
||||
- agentgateway-model-costs
|
||||
verbs:
|
||||
- patch
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: RoleBinding
|
||||
metadata:
|
||||
name: model-catalog
|
||||
namespace: agentgateway-system
|
||||
roleRef:
|
||||
apiGroup: rbac.authorization.k8s.io
|
||||
kind: Role
|
||||
name: model-catalog
|
||||
subjects:
|
||||
- kind: ServiceAccount
|
||||
name: model-catalog
|
||||
namespace: agentgateway-system
|
||||
18
services/agentgateway-model-catalog-sync/k8s/service.yaml
Normal file
18
services/agentgateway-model-catalog-sync/k8s/service.yaml
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: model-catalog
|
||||
namespace: agentgateway-system
|
||||
labels:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
app.kubernetes.io/component: metrics
|
||||
spec:
|
||||
selector:
|
||||
app.kubernetes.io/name: model-catalog
|
||||
app.kubernetes.io/instance: agentgateway
|
||||
ports:
|
||||
- name: metrics
|
||||
port: 8080
|
||||
targetPort: metrics
|
||||
protocol: TCP
|
||||
729
services/agentgateway-model-catalog-sync/main.go
Normal file
729
services/agentgateway-model-catalog-sync/main.go
Normal file
|
|
@ -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
|
||||
}
|
||||
630
services/agentgateway-model-catalog-sync/main_test.go
Normal file
630
services/agentgateway-model-catalog-sync/main_test.go
Normal file
|
|
@ -0,0 +1,630 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestDurationFromEnvironment(t *testing.T) {
|
||||
t.Setenv("TEST_INTERVAL", "")
|
||||
got, err := durationFromEnvironment("TEST_INTERVAL", "6h")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != 6*time.Hour {
|
||||
t.Fatalf("durationFromEnvironment() = %s, want 6h", got)
|
||||
}
|
||||
|
||||
t.Setenv("TEST_INTERVAL", "45m")
|
||||
got, err = durationFromEnvironment("TEST_INTERVAL", "6h")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != 45*time.Minute {
|
||||
t.Fatalf("durationFromEnvironment() = %s, want 45m", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDurationFromEnvironmentRejectsNonPositive(t *testing.T) {
|
||||
t.Setenv("TEST_INTERVAL", "0s")
|
||||
if _, err := durationFromEnvironment("TEST_INTERVAL", "6h"); err == nil {
|
||||
t.Fatal("durationFromEnvironment() accepted a non-positive duration")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDurationFromEnvironmentRejectsInvalidDuration(t *testing.T) {
|
||||
t.Setenv("TEST_INTERVAL", "eventually")
|
||||
if _, err := durationFromEnvironment("TEST_INTERVAL", "6h"); err == nil {
|
||||
t.Fatal("durationFromEnvironment() accepted an invalid duration")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvidersFromEnvironment(t *testing.T) {
|
||||
t.Setenv("ENABLED_PROVIDERS", "")
|
||||
providers, err := providersFromEnvironment()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := mapsKeys(providers)
|
||||
slices.Sort(got)
|
||||
if want := []string{"google-vertex", "novita"}; !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("default providers = %v, want %v", got, want)
|
||||
}
|
||||
|
||||
t.Setenv("ENABLED_PROVIDERS", "openai, google-vertex")
|
||||
providers, err = providersFromEnvironment()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got = mapsKeys(providers)
|
||||
slices.Sort(got)
|
||||
if want := []string{"google-vertex", "openai"}; !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("configured providers = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProvidersFromEnvironmentRejectsInvalidSelections(t *testing.T) {
|
||||
for _, value := range []string{"unknown", "novita,novita", ","} {
|
||||
t.Run(value, func(t *testing.T) {
|
||||
t.Setenv("ENABLED_PROVIDERS", value)
|
||||
if _, err := providersFromEnvironment(); err == nil {
|
||||
t.Fatalf("providersFromEnvironment() accepted %q", value)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func mapsKeys[T any](input map[string]T) []string {
|
||||
keys := make([]string, 0, len(input))
|
||||
for key := range input {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
func selectProviders(names ...string) providerSelection {
|
||||
definitions := buildProviderDefinitions()
|
||||
providers := providerSelection{}
|
||||
for _, name := range names {
|
||||
providers[name] = definitions[name]
|
||||
}
|
||||
return providers
|
||||
}
|
||||
|
||||
func number(value string) *json.Number {
|
||||
number := json.Number(value)
|
||||
return &number
|
||||
}
|
||||
|
||||
func TestConvertModelsDev(t *testing.T) {
|
||||
provider := modelsDevProvider{Models: map[string]modelsDevModel{
|
||||
"model-a": {Cost: modelsDevCost{
|
||||
Input: number("1.5"),
|
||||
Output: number("7.5"),
|
||||
CacheRead: number("0.15"),
|
||||
Tiers: []modelsDevTier{{
|
||||
Input: number("3"),
|
||||
Output: number("15"),
|
||||
Tier: modelsDevTierSelector{Type: "context", Size: 200000},
|
||||
}},
|
||||
}},
|
||||
"unpriced": {},
|
||||
}}
|
||||
want := map[string]catalogModel{
|
||||
"model-a": {
|
||||
Rates: map[string]string{"input": "1.5", "output": "7.5", "cacheRead": "0.15"},
|
||||
Tiers: []catalogTier{{ContextOver: 200000, Rates: map[string]string{"input": "3", "output": "15"}}},
|
||||
},
|
||||
}
|
||||
if got := convertModelsDev(provider); !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("convertModelsDev() = %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertNovita(t *testing.T) {
|
||||
models, err := convertNovita(novitaCatalog{Data: []novitaModel{{
|
||||
ID: "model-a",
|
||||
InputPrice: json.Number("3900"),
|
||||
OutputPrice: json.Number("12500"),
|
||||
}}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := map[string]catalogModel{
|
||||
"model-a": {Rates: map[string]string{"input": "0.39", "output": "1.25"}},
|
||||
}
|
||||
if !reflect.DeepEqual(models, want) {
|
||||
t.Fatalf("convertNovita() = %#v, want %#v", models, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertNovitaSkipsIncompleteModels(t *testing.T) {
|
||||
models, err := convertNovita(novitaCatalog{Data: []novitaModel{
|
||||
{ID: "", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000")},
|
||||
{ID: "missing-input", OutputPrice: json.Number("20000")},
|
||||
{ID: "missing-output", InputPrice: json.Number("10000")},
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(models) != 0 {
|
||||
t.Fatalf("convertNovita() = %#v, want no models", models)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildCatalogUsesAgentgatewayProviderIDs(t *testing.T) {
|
||||
input := modelsDevCatalog{
|
||||
"openai": {Models: map[string]modelsDevModel{"openai-model": {Cost: modelsDevCost{Input: number("1")}}}},
|
||||
"google-vertex": {Models: map[string]modelsDevModel{
|
||||
"vertex-model": {Cost: modelsDevCost{Input: number("2")}},
|
||||
}},
|
||||
}
|
||||
output, err := buildCatalog(input, novitaCatalog{Data: []novitaModel{{
|
||||
ID: "novita-model", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000"),
|
||||
}}}, selectProviders("openai", "google-vertex", "novita"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := []string{"custom", "gcp.vertex_ai", "openai"}
|
||||
got := make([]string, 0, len(output.Providers))
|
||||
for provider := range output.Providers {
|
||||
got = append(got, provider)
|
||||
}
|
||||
slices.Sort(got)
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("provider IDs = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildCatalogRejectsMissingOrInvalidSources(t *testing.T) {
|
||||
valid := modelsDevCatalog{
|
||||
"google-vertex": {Models: map[string]modelsDevModel{
|
||||
"vertex-model": {Cost: modelsDevCost{Input: number("2")}},
|
||||
}},
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
input modelsDevCatalog
|
||||
novita novitaCatalog
|
||||
}{
|
||||
{
|
||||
name: "missing models.dev provider",
|
||||
input: modelsDevCatalog{},
|
||||
novita: novitaCatalog{Data: []novitaModel{{
|
||||
ID: "novita-model", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000"),
|
||||
}}},
|
||||
},
|
||||
{
|
||||
name: "missing Novita models",
|
||||
input: valid,
|
||||
},
|
||||
{
|
||||
name: "invalid Novita price",
|
||||
input: valid,
|
||||
novita: novitaCatalog{Data: []novitaModel{{
|
||||
ID: "novita-model", InputPrice: json.Number("-1"), OutputPrice: json.Number("20000"),
|
||||
}}},
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if _, err := buildCatalog(
|
||||
test.input,
|
||||
test.novita,
|
||||
selectProviders("google-vertex", "novita"),
|
||||
); err == nil {
|
||||
t.Fatal("buildCatalog() unexpectedly succeeded")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTenThousandths(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{input: "0", want: "0"},
|
||||
{input: "1", want: "0.0001"},
|
||||
{input: "3900", want: "0.39"},
|
||||
{input: "10000", want: "1"},
|
||||
{input: "12500", want: "1.25"},
|
||||
{input: "00012500", want: "1.25"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.input, func(t *testing.T) {
|
||||
got, err := tenThousandths(json.Number(test.input))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != test.want {
|
||||
t.Fatalf("tenThousandths(%q) = %q, want %q", test.input, got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTenThousandthsRejectsInvalidNumbers(t *testing.T) {
|
||||
for _, input := range []string{"0.39", "-1", "1e4"} {
|
||||
t.Run(input, func(t *testing.T) {
|
||||
if _, err := tenThousandths(json.Number(input)); err == nil {
|
||||
t.Fatalf("tenThousandths(%q) unexpectedly succeeded", input)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderMetrics(t *testing.T) {
|
||||
input := catalog{Providers: map[string]catalogProvider{
|
||||
"custom": {Models: map[string]catalogModel{
|
||||
"vendor/model": {Rates: map[string]string{"output": "1.25", "cacheRead": "0.1"}},
|
||||
}},
|
||||
}}
|
||||
want := "# HELP agentgateway_model_token_price_usd_per_million Model token price in USD per million tokens.\n" +
|
||||
"# TYPE agentgateway_model_token_price_usd_per_million gauge\n" +
|
||||
"agentgateway_model_token_price_usd_per_million{gen_ai_system=\"custom\",gen_ai_request_model=\"vendor/model\",gen_ai_token_type=\"cache_read\",source_provider=\"novita\"} 0.1\n" +
|
||||
"agentgateway_model_token_price_usd_per_million{gen_ai_system=\"custom\",gen_ai_request_model=\"vendor/model\",gen_ai_token_type=\"output\",source_provider=\"novita\"} 1.25\n"
|
||||
if got := string(renderMetrics(input)); got != want {
|
||||
t.Fatalf("renderMetrics() = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderModelList(t *testing.T) {
|
||||
input := catalog{Providers: map[string]catalogProvider{
|
||||
"gcp.vertex_ai": {Models: map[string]catalogModel{"gemini-z": {}, "gemini-a": {}}},
|
||||
"custom": {Models: map[string]catalogModel{"vendor/model": {}}},
|
||||
}}
|
||||
want := modelList{Object: "list", Data: []modelListItem{
|
||||
{ID: "novita/vendor/model", Object: "model", Created: 0, OwnedBy: "novita"},
|
||||
{ID: "vertex_ai/gemini-a", Object: "model", Created: 0, OwnedBy: "google-vertex"},
|
||||
{ID: "vertex_ai/gemini-z", Object: "model", Created: 0, OwnedBy: "google-vertex"},
|
||||
}}
|
||||
if got := renderModelList(input); !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("renderModelList() = %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadCatalog(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "catalog.json")
|
||||
want := catalog{Providers: map[string]catalogProvider{
|
||||
"gcp.vertex_ai": {Models: map[string]catalogModel{"gemini": {Rates: map[string]string{"input": "1"}}}},
|
||||
}}
|
||||
contents, err := json.Marshal(want)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, contents, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := readCatalog(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("readCatalog() = %#v, want %#v", got, want)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(path, []byte("{"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := readCatalog(path); err == nil {
|
||||
t.Fatal("readCatalog() accepted malformed JSON")
|
||||
}
|
||||
if _, err := readCatalog(filepath.Join(t.TempDir(), "missing")); err == nil {
|
||||
t.Fatal("readCatalog() accepted a missing file")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCatalogHandler(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "catalog.json")
|
||||
contents := `{"providers":{
|
||||
"gcp.vertex_ai":{"models":{"gemini":{"rates":{"input":"2"}}}},
|
||||
"custom":{"models":{"novita/model":{"rates":{"input":"3"}}}}
|
||||
}}`
|
||||
if err := os.WriteFile(path, []byte(contents), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
handler := catalogHandler(path, selectProviders("google-vertex", "novita"))
|
||||
tests := []struct {
|
||||
name string
|
||||
method string
|
||||
path string
|
||||
status int
|
||||
contentType string
|
||||
body string
|
||||
}{
|
||||
{name: "health", method: http.MethodGet, path: "/healthz", status: http.StatusNoContent},
|
||||
{name: "ready", method: http.MethodGet, path: "/readyz", status: http.StatusNoContent},
|
||||
{
|
||||
name: "metrics",
|
||||
method: http.MethodGet,
|
||||
path: "/metrics",
|
||||
status: http.StatusOK,
|
||||
contentType: "text/plain; version=0.0.4",
|
||||
body: `gen_ai_request_model="gemini"`,
|
||||
},
|
||||
{
|
||||
name: "models",
|
||||
method: http.MethodGet,
|
||||
path: "/v1/models",
|
||||
status: http.StatusOK,
|
||||
contentType: "application/json",
|
||||
body: `"id":"vertex_ai/gemini"`,
|
||||
},
|
||||
{name: "models method", method: http.MethodPost, path: "/v1/models", status: http.StatusMethodNotAllowed},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
request := httptest.NewRequest(test.method, test.path, nil)
|
||||
response := httptest.NewRecorder()
|
||||
handler.ServeHTTP(response, request)
|
||||
if response.Code != test.status {
|
||||
t.Fatalf("status = %d, want %d", response.Code, test.status)
|
||||
}
|
||||
if test.contentType != "" && response.Header().Get("Content-Type") != test.contentType {
|
||||
t.Fatalf("Content-Type = %q, want %q", response.Header().Get("Content-Type"), test.contentType)
|
||||
}
|
||||
if test.body != "" && !strings.Contains(response.Body.String(), test.body) {
|
||||
t.Fatalf("body = %q, want substring %q", response.Body.String(), test.body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCatalogHandlerRejectsInvalidCatalog(t *testing.T) {
|
||||
missing := filepath.Join(t.TempDir(), "missing")
|
||||
providers := selectProviders("google-vertex", "novita")
|
||||
handler := catalogHandler(missing, providers)
|
||||
for _, path := range []string{"/readyz", "/metrics", "/v1/models"} {
|
||||
t.Run(path, func(t *testing.T) {
|
||||
request := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
response := httptest.NewRecorder()
|
||||
handler.ServeHTTP(response, request)
|
||||
if response.Code != http.StatusServiceUnavailable {
|
||||
t.Fatalf("status = %d, want %d", response.Code, http.StatusServiceUnavailable)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
empty := filepath.Join(t.TempDir(), "catalog.json")
|
||||
if err := os.WriteFile(empty, []byte(`{"providers":{}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := httptest.NewRequest(http.MethodGet, "/readyz", nil)
|
||||
response := httptest.NewRecorder()
|
||||
catalogHandler(empty, providers).ServeHTTP(response, request)
|
||||
if response.Code != http.StatusServiceUnavailable {
|
||||
t.Fatalf("empty catalog status = %d, want %d", response.Code, http.StatusServiceUnavailable)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchJSON(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||||
if request.Header.Get("Accept") != "application/json" {
|
||||
t.Errorf("Accept = %q", request.Header.Get("Accept"))
|
||||
}
|
||||
if request.Header.Get("Authorization") != "Bearer test" {
|
||||
t.Errorf("Authorization = %q", request.Header.Get("Authorization"))
|
||||
}
|
||||
return response(http.StatusOK, `{"value":12345678901234567890}`), nil
|
||||
})}
|
||||
|
||||
var target struct {
|
||||
Value json.Number `json:"value"`
|
||||
}
|
||||
if err := fetchJSON(context.Background(), client, "https://models.invalid", "Bearer test", &target); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if target.Value.String() != "12345678901234567890" {
|
||||
t.Fatalf("value = %q", target.Value)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchJSONRejectsHTTPAndJSONErrors(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
status int
|
||||
body string
|
||||
}{
|
||||
{name: "status", status: http.StatusBadGateway, body: `{"error":"upstream"}`},
|
||||
{name: "JSON", status: http.StatusOK, body: "{"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return response(test.status, test.body), nil
|
||||
})}
|
||||
if err := fetchJSON(context.Background(), client, "https://models.invalid", "", &map[string]any{}); err == nil {
|
||||
t.Fatal("fetchJSON() unexpectedly succeeded")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteConfigMapPatchesExistingConfigMap(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||||
if request.Method != http.MethodPatch {
|
||||
t.Errorf("method = %q", request.Method)
|
||||
}
|
||||
if request.URL.Path != "/configmaps/agentgateway-model-costs" {
|
||||
t.Errorf("path = %q", request.URL.Path)
|
||||
}
|
||||
if request.Header.Get("Authorization") != "Bearer token" {
|
||||
t.Errorf("Authorization = %q", request.Header.Get("Authorization"))
|
||||
}
|
||||
if request.Header.Get("Content-Type") != "application/merge-patch+json" {
|
||||
t.Errorf("Content-Type = %q", request.Header.Get("Content-Type"))
|
||||
}
|
||||
var patch configMapPatch
|
||||
if err := json.NewDecoder(request.Body).Decode(&patch); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if patch.Data["catalog.json"] != `{"providers":{}}` {
|
||||
t.Errorf("catalog.json = %q", patch.Data["catalog.json"])
|
||||
}
|
||||
return response(http.StatusOK, `{}`), nil
|
||||
})}
|
||||
|
||||
err := writeConfigMap(
|
||||
context.Background(),
|
||||
client,
|
||||
"https://kubernetes.invalid/configmaps",
|
||||
[]byte("token"),
|
||||
[]byte(`{"providers":{}}`),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteConfigMapCreatesMissingConfigMap(t *testing.T) {
|
||||
requests := 0
|
||||
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||||
requests++
|
||||
if requests == 1 {
|
||||
if request.Method != http.MethodPatch {
|
||||
t.Errorf("first method = %q", request.Method)
|
||||
}
|
||||
return response(http.StatusNotFound, ""), nil
|
||||
}
|
||||
if request.Method != http.MethodPost {
|
||||
t.Errorf("second method = %q", request.Method)
|
||||
}
|
||||
if request.URL.Path != "/configmaps" {
|
||||
t.Errorf("path = %q", request.URL.Path)
|
||||
}
|
||||
var resource configMapResource
|
||||
if err := json.NewDecoder(request.Body).Decode(&resource); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if resource.APIVersion != "v1" || resource.Kind != "ConfigMap" {
|
||||
t.Errorf("resource = %#v", resource)
|
||||
}
|
||||
if resource.Metadata.Name != configMapName {
|
||||
t.Errorf("name = %q", resource.Metadata.Name)
|
||||
}
|
||||
if resource.Metadata.Annotations["kustomize.toolkit.fluxcd.io/prune"] != "disabled" {
|
||||
t.Errorf("annotations = %#v", resource.Metadata.Annotations)
|
||||
}
|
||||
if resource.Metadata.Labels["app.kubernetes.io/instance"] != "agentgateway" ||
|
||||
resource.Metadata.Labels["app.kubernetes.io/name"] != "model-catalog" ||
|
||||
resource.Metadata.Labels["app.kubernetes.io/part-of"] != "agentgateway" {
|
||||
t.Errorf("labels = %#v", resource.Metadata.Labels)
|
||||
}
|
||||
if resource.Data["catalog.json"] != `{"providers":{}}` {
|
||||
t.Errorf("catalog.json = %q", resource.Data["catalog.json"])
|
||||
}
|
||||
return response(http.StatusCreated, `{}`), nil
|
||||
})}
|
||||
|
||||
if err := writeConfigMap(
|
||||
context.Background(),
|
||||
client,
|
||||
"https://kubernetes.invalid/configmaps",
|
||||
[]byte("token"),
|
||||
[]byte(`{"providers":{}}`),
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if requests != 2 {
|
||||
t.Fatalf("requests = %d, want 2", requests)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteConfigMapRetriesPatchAfterCreateConflict(t *testing.T) {
|
||||
requests := 0
|
||||
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||||
requests++
|
||||
switch requests {
|
||||
case 1:
|
||||
return response(http.StatusNotFound, ""), nil
|
||||
case 2:
|
||||
return response(http.StatusConflict, ""), nil
|
||||
default:
|
||||
if request.Method != http.MethodPatch {
|
||||
t.Errorf("third method = %q", request.Method)
|
||||
}
|
||||
return response(http.StatusOK, `{}`), nil
|
||||
}
|
||||
})}
|
||||
|
||||
if err := writeConfigMap(
|
||||
context.Background(),
|
||||
client,
|
||||
"https://kubernetes.invalid/configmaps",
|
||||
[]byte("token"),
|
||||
[]byte(`{"providers":{}}`),
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if requests != 3 {
|
||||
t.Fatalf("requests = %d, want 3", requests)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteConfigMapRejectsHTTPAndTransportErrors(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return response(http.StatusForbidden, ""), nil
|
||||
})}
|
||||
if err := writeConfigMap(context.Background(), client, "https://kubernetes.invalid", []byte("token"), []byte("{}")); err == nil {
|
||||
t.Fatal("writeConfigMap() accepted an HTTP error")
|
||||
}
|
||||
|
||||
client = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return nil, errors.New("transport failed")
|
||||
})}
|
||||
if err := writeConfigMap(context.Background(), client, "https://kubernetes.invalid", []byte("token"), []byte("{}")); err == nil {
|
||||
t.Fatal("writeConfigMap() accepted a transport error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKubernetesRequestDecodesResponse(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||||
if request.Header.Get("Content-Type") != "" {
|
||||
t.Errorf("Content-Type = %q, want empty", request.Header.Get("Content-Type"))
|
||||
}
|
||||
return response(http.StatusOK, `{"value":"decoded"}`), nil
|
||||
})}
|
||||
target := map[string]string{}
|
||||
status, err := kubernetesRequest(
|
||||
context.Background(),
|
||||
client,
|
||||
http.MethodGet,
|
||||
"https://kubernetes.invalid",
|
||||
[]byte("token"),
|
||||
nil,
|
||||
&target,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if status != http.StatusOK || target["value"] != "decoded" {
|
||||
t.Fatalf("status = %d, target = %#v", status, target)
|
||||
}
|
||||
}
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
return function(request)
|
||||
}
|
||||
|
||||
func response(status int, body string) *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Status: http.StatusText(status),
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue