[beta] Add an optional deployment for AI Gateway

.github/workflows/test-helm-chart.yml

@@ -21,7 +21,7 @@ jobs:
 
       - name: Install helm-unittest
         run: |
-          helm plugin install https://github.com/quintush/helm-unittest --verify=false
+          helm plugin install https://github.com/helm-unittest/helm-unittest.git --verify=false
 
       - name: Run unit tests
         run: |

braintrust/templates/_helpers.tpl

@@ -25,6 +25,20 @@ Static fast reader query sources used by API.
 {{- join "," $sources -}}
 {{- end -}}
 
+{{/*
+Internal cluster URL for the API service.
+*/}}
+{{- define "braintrust.apiInternalUrl" -}}
+http://{{ .Values.api.service.name | default .Values.api.name }}.{{ include "braintrust.namespace" . }}:{{ .Values.api.service.port }}
+{{- end -}}
+
+{{/*
+Internal cluster URL for the AI Gateway service.
+*/}}
+{{- define "braintrust.aiGatewayInternalUrl" -}}
+http://{{ .Values.aiGateway.service.name | default .Values.aiGateway.name }}.{{ include "braintrust.namespace" . }}:{{ .Values.aiGateway.service.port }}
+{{- end -}}
+
 {{/*
 Render Brainstore container resources with provider-specific ephemeral storage.
 

braintrust/templates/ai-gateway-configmap.yaml

@@ -0,0 +1,27 @@
+{{- if .Values.aiGateway.enabled }}
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: {{ .Values.aiGateway.name }}
+  namespace: {{ include "braintrust.namespace" . }}
+  {{- with (merge (deepCopy .Values.aiGateway.labels) .Values.global.labels) }}
+  labels:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+  {{- with .Values.aiGateway.annotations.configmap }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+data:
+  GATEWAY_ENV: "production"
+  {{- if .Values.aiGateway.region }}
+  GATEWAY_REGION: {{ .Values.aiGateway.region | quote }}
+  {{- end }}
+  BRAINTRUST_APP_URL: {{ .Values.aiGateway.braintrustAppUrl | quote }}
+  BRAINTRUST_API_URL: {{ default (include "braintrust.apiInternalUrl" .) .Values.aiGateway.braintrustApiUrl | quote }}
+  GATEWAY_JSON_LOGS: "true"
+  GATEWAY_TELEMETRY: {{ .Values.global.controlPlaneTelemetry | quote }}
+  {{- $appUrl := .Values.aiGateway.braintrustAppUrl | trimSuffix "/" }}
+  OTLP_HTTP_ENDPOINT: {{ default (printf "%s/api/pulse/otel" $appUrl) .Values.aiGateway.otlpHttpEndpoint | quote }}
+{{- end }}

braintrust/templates/ai-gateway-deployment.yaml

@@ -0,0 +1,98 @@
+{{- if .Values.aiGateway.enabled }}
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: {{ .Values.aiGateway.name }}
+  namespace: {{ include "braintrust.namespace" . }}
+  {{- with (merge (deepCopy .Values.aiGateway.labels) .Values.global.labels) }}
+  labels:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+  {{- with .Values.aiGateway.annotations.deployment }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+spec:
+  replicas: {{ .Values.aiGateway.replicas }}
+  strategy:
+    type: RollingUpdate
+    rollingUpdate:
+      maxSurge: 100%
+      maxUnavailable: 0
+  selector:
+    matchLabels:
+      app: {{ .Values.aiGateway.name }}
+  template:
+    metadata:
+      labels:
+        app: {{ .Values.aiGateway.name }}
+        {{- with (merge (deepCopy .Values.aiGateway.podLabels) .Values.aiGateway.labels .Values.global.labels) }}
+        {{- toYaml . | nindent 8 }}
+        {{- end }}
+      annotations:
+        checksum/config: {{ include (print $.Template.BasePath "/ai-gateway-configmap.yaml") . | sha256sum }}
+        {{- with .Values.aiGateway.annotations.pod }}
+        {{- toYaml . | nindent 8 }}
+        {{- end }}
+    spec:
+      {{- with .Values.aiGateway.nodeSelector }}
+      nodeSelector:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.aiGateway.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.aiGateway.affinity }}
+      affinity:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      containers:
+        - name: ai-gateway
+          image: "{{ .Values.aiGateway.image.repository }}:{{ .Values.aiGateway.image.tag }}"
+          imagePullPolicy: {{ .Values.aiGateway.image.pullPolicy }}
+          ports:
+            - containerPort: {{ .Values.aiGateway.service.port }}
+          resources:
+            {{- toYaml .Values.aiGateway.resources | nindent 12 }}
+          {{- with .Values.aiGateway.livenessProbe }}
+          livenessProbe:
+            {{- toYaml . | nindent 12 }}
+          {{- end }}
+          {{- with .Values.aiGateway.readinessProbe }}
+          readinessProbe:
+            {{- toYaml . | nindent 12 }}
+          {{- end }}
+          envFrom:
+            - configMapRef:
+                name: {{ .Values.aiGateway.name }}
+          env:
+            - name: COMPLETIONS_CACHE_REDIS_URL
+              valueFrom:
+                secretKeyRef:
+                  name: braintrust-secrets
+                  key: REDIS_URL
+            - name: AUTH_CACHE_REDIS_URL
+              valueFrom:
+                secretKeyRef:
+                  name: braintrust-secrets
+                  key: REDIS_URL
+            - name: BRAINSTORE_LICENSE_KEY
+              valueFrom:
+                secretKeyRef:
+                  name: braintrust-secrets
+                  key: BRAINSTORE_LICENSE_KEY
+            {{- if .Values.aiGateway.extraEnvVars }}
+            {{- toYaml .Values.aiGateway.extraEnvVars | nindent 12 }}
+            {{- end }}
+      {{- with .Values.aiGateway.extraContainers }}
+      {{- toYaml . | nindent 8 }}
+      {{- end }}
+      volumes:
+        {{- with .Values.aiGateway.extraVolumes }}
+        {{- toYaml . | nindent 8 }}
+        {{- else }}
+        []
+        {{- end }}
+{{- end }}

braintrust/templates/ai-gateway-service.yaml

@@ -0,0 +1,25 @@
+{{- if .Values.aiGateway.enabled }}
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: {{ .Values.aiGateway.service.name | default .Values.aiGateway.name }}
+  namespace: {{ include "braintrust.namespace" . }}
+  {{- with (merge (deepCopy .Values.aiGateway.labels) .Values.global.labels) }}
+  labels:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+  {{- with .Values.aiGateway.annotations.service }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+spec:
+  selector:
+    app: {{ .Values.aiGateway.name }}
+  ports:
+    - name: {{ .Values.aiGateway.service.portName }}
+      protocol: TCP
+      port: {{ .Values.aiGateway.service.port }}
+      targetPort: {{ .Values.aiGateway.service.port }}
+  type: {{ .Values.aiGateway.service.type }}
+{{- end }}

braintrust/templates/api-configmap.yaml

@@ -64,3 +64,6 @@ data:
   {{- if or .Values.brainstoreWalFooterVersion .Values.skipPgForBrainstoreObjects }}
   BRAINSTORE_WAL_USE_EFFICIENT_FORMAT: "true"
   {{- end }}
+  {{- if .Values.aiGateway.enabled }}
+  GATEWAY_URL: {{ include "braintrust.aiGatewayInternalUrl" . | quote }}
+  {{- end }}

braintrust/tests/ai-gateway-configmap_test.yaml

@@ -0,0 +1,127 @@
+suite: test AI gateway configmap template
+templates:
+  - ai-gateway-configmap.yaml
+tests:
+  - it: should not render when ai-gateway is disabled
+    values:
+      - __fixtures__/base-values.yaml
+    release:
+      namespace: "braintrust"
+    asserts:
+      - hasDocuments:
+          count: 0
+
+  - it: should render ai-gateway configmap when enabled
+    values:
+      - __fixtures__/base-values.yaml
+    set:
+      aiGateway.enabled: true
+    release:
+      namespace: "braintrust"
+    asserts:
+      - isKind:
+          of: ConfigMap
+      - equal:
+          path: metadata.name
+          value: braintrust-ai-gateway
+      - equal:
+          path: data.GATEWAY_ENV
+          value: "production"
+      - isNull:
+          path: data.GATEWAY_REGION
+      - equal:
+          path: data.BRAINTRUST_APP_URL
+          value: "https://www.braintrust.dev"
+      - equal:
+          path: data.BRAINTRUST_API_URL
+          value: "http://braintrust-api.braintrust:8000"
+      - equal:
+          path: data.GATEWAY_JSON_LOGS
+          value: "true"
+      - equal:
+          path: data.GATEWAY_TELEMETRY
+          value: "status,metrics"
+      - equal:
+          path: data.OTLP_HTTP_ENDPOINT
+          value: "https://www.braintrust.dev/api/pulse/otel"
+
+  - it: should set GATEWAY_REGION when region is provided
+    values:
+      - __fixtures__/base-values.yaml
+    set:
+      aiGateway.enabled: true
+      aiGateway.region: "us-east-1"
+    release:
+      namespace: "braintrust"
+    asserts:
+      - equal:
+          path: data.GATEWAY_REGION
+          value: "us-east-1"
+
+  - it: should use global controlPlaneTelemetry for GATEWAY_TELEMETRY
+    values:
+      - __fixtures__/base-values.yaml
+    set:
+      aiGateway.enabled: true
+      aiGateway.region: "us-east-1"
+      global.controlPlaneTelemetry: "status,metrics,usage,logs"
+    release:
+      namespace: "braintrust"
+    asserts:
+      - equal:
+          path: data.GATEWAY_TELEMETRY
+          value: "status,metrics,usage,logs"
+
+  - it: should use custom braintrustApiUrl when provided
+    values:
+      - __fixtures__/base-values.yaml
+    set:
+      aiGateway.enabled: true
+      aiGateway.region: "us-east-1"
+      aiGateway.braintrustApiUrl: "https://api.example.com"
+    release:
+      namespace: "braintrust"
+    asserts:
+      - equal:
+          path: data.BRAINTRUST_API_URL
+          value: "https://api.example.com"
+
+  - it: should derive OTLP_HTTP_ENDPOINT from braintrustAppUrl
+    values:
+      - __fixtures__/base-values.yaml
+    set:
+      aiGateway.enabled: true
+      aiGateway.braintrustAppUrl: "https://control.example.com"
+    release:
+      namespace: "braintrust"
+    asserts:
+      - equal:
+          path: data.OTLP_HTTP_ENDPOINT
+          value: "https://control.example.com/api/pulse/otel"
+
+  - it: should derive OTLP_HTTP_ENDPOINT from braintrustAppUrl with trailing slash
+    values:
+      - __fixtures__/base-values.yaml
+    set:
+      aiGateway.enabled: true
+      aiGateway.braintrustAppUrl: "https://control.example.com/"
+    release:
+      namespace: "braintrust"
+    asserts:
+      - equal:
+          path: data.OTLP_HTTP_ENDPOINT
+          value: "https://control.example.com/api/pulse/otel"
+
+  - it: should use custom otlpHttpEndpoint when provided
+    values:
+      - __fixtures__/base-values.yaml
+    set:
+      aiGateway.enabled: true
+      aiGateway.braintrustAppUrl: "https://control.example.com"
+      aiGateway.otlpHttpEndpoint: "https://otel.example.com/v1"
+    release:
+      namespace: "braintrust"
+    asserts:
+      - equal:
+          path: data.OTLP_HTTP_ENDPOINT
+          value: "https://otel.example.com/v1"