Repository: envoyproxy/ai-gateway
Pull request #530 (2025-03-27 18:24:01)

examples: adds prometheus scraping setup
**Commit Message**

This adds an example prometheus setup to scrape LLM metrics from extproc deployments and the case in tests/e2e to verify it's working.

**Related Issues/PRs (if applicable)**

Contributes to #245 


Patch:
diff --git a/examples/monitoring/README.md b/examples/monitoring/README.md
new file mode 100644
index 0000000..622aae5
--- /dev/null
+++ b/examples/monitoring/README.md
@@ -0,0 +1,2 @@
+This `monitoring.yaml` file is a Kubernetes manifest file that deploys a Prometheus server that scrapes metrics from pods with the standard annotations.
+More specifically, Envoy AI Gateway produced metrics can be collected by `prometheus.io/scrape: 'true'` annotation. so there's nothing Envoy AI Gateway specific in this file.
diff --git a/examples/monitoring/monitoring.yaml b/examples/monitoring/monitoring.yaml
new file mode 100644
index 0000000..924d598
--- /dev/null
+++ b/examples/monitoring/monitoring.yaml
@@ -0,0 +1,131 @@
+# Copyright Envoy AI Gateway Authors
+# SPDX-License-Identifier: Apache-2.0
+# The full text of the Apache license is available in the LICENSE file at
+# the root of the repo.
+
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: monitoring
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: prometheus-config
+  namespace: monitoring
+data:
+  prometheus.yml: |
+    global:
+      scrape_interval: 3s
+    scrape_configs:
+      - job_name: 'kubernetes-pods'
+        kubernetes_sd_configs:
+          - role: pod
+        relabel_configs:
+          - source_labels: [__meta_kubernetes_pod_annotation_prometheus_io_scrape]
+            action: keep
+            regex: true
+          - source_labels: [__meta_kubernetes_pod_annotation_prometheus_io_path]
+            action: replace
+            target_label: __metrics_path__
+            regex: (.+)
+          - source_labels: [__address__, __meta_kubernetes_pod_annotation_prometheus_io_port]
+            action: replace
+            regex: ([^:]+)(?::\d+)?;(\d+)
+            replacement: $1:$2
+            target_label: __address__
+          - action: labelmap
+            regex: __meta_kubernetes_pod_label_(.+)
+          - source_labels: [__meta_kubernetes_namespace]
+            action: replace
+            target_label: kubernetes_namespace
+          - source_labels: [__meta_kubernetes_pod_name]
+            action: replace
+            target_label: kubernetes_pod_name
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRole
+metadata:
+  name: prometheus
+rules:
+- apiGroups: [""]
+  resources:
+  - nodes
+  - nodes/proxy
+  - services
+  - endpoints
+  - pods
+  - ingresses
+  verbs: ["get", "list", "watch"]
+- apiGroups:
+  - extensions
+  resources:
+  - ingresses
+  verbs: ["get", "list", "watch"]
+- nonResourceURLs: ["/metrics"]
+  verbs: ["get"]
+---
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: prometheus
+  namespace: monitoring
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRoleBinding
+metadata:
+  name: prometheus
+roleRef:
+  apiGroup: rbac.authorization.k8s.io
+  kind: ClusterRole
+  name: prometheus
+subjects:
+- kind: ServiceAccount
+  name: prometheus
+  namespace: monitoring
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: prometheus
+  namespace: monitoring
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: prometheus
+  template:
+    metadata:
+      labels:
+        app: prometheus
+    spec:
+      serviceAccountName: prometheus
+      containers:
+      - name: prometheus
+        image: prom/prometheus:v2.47.0
+        args:
+        - "--config.file=/etc/prometheus/prometheus.yml"
+        volumeMounts:
+        - name: prometheus-config
+          mountPath: /etc/prometheus/
+        ports:
+        - containerPort: 9090
+          name: web
+      volumes:
+      - name: prometheus-config
+        configMap:
+          name: prometheus-config
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: prometheus
+  namespace: monitoring
+spec:
+  selector:
+    app: prometheus
+  ports:
+  - protocol: TCP
+    port: 9090
+    targetPort: 9090
+    name: web
diff --git a/examples/token_ratelimit/token_ratelimit.yaml b/examples/token_ratelimit/token_ratelimit.yaml
index d3c6a20..60f6c86 100644
--- a/examples/token_ratelimit/token_ratelimit.yaml
+++ b/examples/token_ratelimit/token_ratelimit.yaml
@@ -39,7 +39,7 @@ spec:
         - headers:
             - type: Exact
               name: x-ai-eg-model
-              value: gpt-4o-mini
+              value: rate-limit-funky-model
       backendRefs:
         - name: envoy-ai-gateway-token-ratelimit-testupstream
   # The following metadata keys are used to store the costs from the LLM request.
diff --git a/internal/apischema/openai/openai.go b/internal/apischema/openai/openai.go
index 3065952..d3bb7ad 100644
--- a/internal/apischema/openai/openai.go
+++ b/internal/apischema/openai/openai.go
@@ -409,10 +409,10 @@ type ChatCompletionResponseFormat struct {
 }
 
 type ChatCompletionResponseFormatJSONSchema struct {
-	Name        string         `json:"name"`
-	Description string         `json:"description,omitempty"`
-	Schema      json.Marshaler `json:"schema"`
-	Strict      bool           `json:"strict"`
+	Name        string `json:"name"`
+	Description string `json:"description,omitempty"`
+	Schema      any    `json:"schema"` // See detail in https://github.com/openai/openai-go/blob/28c93a9fa58bb622b5d23b3262af7d4fdd2ebde9/shared/shared.go#L519C6-L519C30
+	Strict      bool   `json:"strict"`
 }
 
 // ChatCompletionRequest represents a request structure for chat completion API.
diff --git a/internal/controller/ai_gateway_route.go b/internal/controller/ai_gateway_route.go
index e85be34..ef76b0a 100644
--- a/internal/controller/ai_gateway_route.go
+++ b/internal/controller/ai_gateway_route.go
@@ -509,6 +509,12 @@ func (c *AIGatewayRouteController) annotateExtProcPods(ctx context.Context, aiGa
 	return nil
 }
 
+var extProcPrometheusAnnotations = map[string]string{
+	"prometheus.io/scrape": "true",
+	"prometheus.io/port":   "9190",
+	"prometheus.io/path":   "/metrics",
+}
+
 // syncExtProcDeployment syncs the external processor's Deployment and Service.
 func (c *AIGatewayRouteController) syncExtProcDeployment(ctx context.Context, aiGatewayRoute *aigv1a1.AIGatewayRoute) error {
 	name := extProcName(aiGatewayRoute)
@@ -526,14 +532,17 @@ func (c *AIGatewayRouteController) syncExtProcDeployment(ctx context.Context, ai
 				Spec: appsv1.DeploymentSpec{
 					Selector: &metav1.LabelSelector{MatchLabels: labels},
 					Template: corev1.PodTemplateSpec{
-						ObjectMeta: metav1.ObjectMeta{Labels: labels},
+						ObjectMeta: metav1.ObjectMeta{Labels: labels, Annotations: extProcPrometheusAnnotations},
 						Spec: corev1.PodSpec{
 							Containers: []corev1.Container{
 								{
 									Name:            name,
 									Image:           c.extProcImage,
 									ImagePullPolicy: c.extProcImagePullPolicy,
-									Ports:           []corev1.ContainerPort{{Name: "grpc", ContainerPort: 1063}},
+									Ports: []corev1.ContainerPort{
+										{Name: "grpc", ContainerPort: 1063},
+										{Name: "metrics", ContainerPort: 9190},
+									},
 									Args: []string{
 										"-configPath", "/etc/ai-gateway/extproc/" + expProcConfigFileName,
 										"-logLevel", c.extProcLogLevel,


Test patch:
diff --git a/cmd/aigw/testdata/translate_basic.out.yaml b/cmd/aigw/testdata/translate_basic.out.yaml
index e564b65..9726a75 100644
--- a/cmd/aigw/testdata/translate_basic.out.yaml
+++ b/cmd/aigw/testdata/translate_basic.out.yaml
@@ -178,6 +178,10 @@ spec:
   strategy: {}
   template:
     metadata:
+      annotations:
+        "prometheus.io/path": "/metrics"
+        "prometheus.io/port": "9190"
+        "prometheus.io/scrape": "true"
       labels:
         app: ai-eg-route-extproc-envoy-ai-gateway-basic
         app.kubernetes.io/managed-by: envoy-ai-gateway
@@ -194,6 +198,8 @@ spec:
           ports:
             - containerPort: 1063
               name: grpc
+            - containerPort: 9190
+              name: metrics
           resources: {}
           volumeMounts:
             - mountPath: /etc/ai-gateway/extproc
diff --git a/internal/apischema/openai/openai_test.go b/internal/apischema/openai/openai_test.go
index e0b621d..5d46ee3 100644
--- a/internal/apischema/openai/openai_test.go
+++ b/internal/apischema/openai/openai_test.go
@@ -253,6 +253,39 @@ func TestOpenAIChatCompletionMessageUnmarshal(t *testing.T) {
                         "messages": [{"role": "some-funky", "content": [{"text": "what do you see in this image", "type": "text"}]}]}`),
 			expErr: "unknown ChatCompletionMessageParam type: some-funky",
 		},
+		{
+			name: "response_format",
+			in:   []byte(`{ "model": "azure.gpt-4o", "messages": [ { "role": "user", "content": "Tell me a story" } ], "response_format": { "type": "json_schema", "json_schema": { "name": "math_response", "schema": { "type": "object", "properties": { "step": "test_step" }, "required": [ "steps"], "additionalProperties": false }, "strict": true } } }`),
+			out: &ChatCompletionRequest{
+				Model: "azure.gpt-4o",
+				Messages: []ChatCompletionMessageParamUnion{
+					{
+						Value: ChatCompletionUserMessageParam{
+							Role: ChatMessageRoleUser,
+							Content: StringOrUserRoleContentUnion{
+								Value: "Tell me a story",
+							},
+						},
+						Type: ChatMessageRoleUser,
+					},
+				},
+				ResponseFormat: &ChatCompletionResponseFormat{
+					Type: "json_schema",
+					JSONSchema: &ChatCompletionResponseFormatJSONSchema{
+						Name:   "math_response",
+						Strict: true,
+						Schema: map[string]interface{}{
+							"additionalProperties": false,
+							"type":                 "object",
+							"properties": map[string]interface{}{
+								"step": "test_step",
+							},
+							"required": []interface{}{"steps"},
+						},
+					},
+				},
+			},
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			var chatCompletion ChatCompletionRequest
diff --git a/tests/e2e/e2e_test.go b/tests/e2e/e2e_test.go
index 69c4692..3cdd6a8 100644
--- a/tests/e2e/e2e_test.go
+++ b/tests/e2e/e2e_test.go
@@ -87,6 +87,11 @@ func TestMain(m *testing.M) {
 		panic(err)
 	}
 
+	if err := initPrometheus(ctx); err != nil {
+		cancel()
+		panic(err)
+	}
+
 	code := m.Run()
 	cancel()
 
@@ -226,6 +231,21 @@ func initTestupstream(ctx context.Context) (err error) {
 	return kubectlWaitForDeploymentReady("default", "testupstream")
 }
 
+func initPrometheus(ctx context.Context) (err error) {
+	initLog("Installing Prometheus")
+	start := time.Now()
+	defer func() {
+		elapsed := time.Since(start)
+		initLog(fmt.Sprintf("\tdone (took %.2fs in total)\n", elapsed.Seconds()))
+	}()
+	initLog("\tapplying manifests")
+	if err = kubectlApplyManifest(ctx, "../../examples/monitoring/monitoring.yaml"); err != nil {
+		return
+	}
+	initLog("\twaiting for deployment")
+	return kubectlWaitForDeploymentReady("monitoring", "prometheus")
+}
+
 func kubectl(ctx context.Context, args ...string) *exec.Cmd {
 	cmd := exec.CommandContext(ctx, "kubectl", args...)
 	cmd.Stdout = os.Stdout
diff --git a/tests/e2e/token_ratelimit_test.go b/tests/e2e/token_ratelimit_test.go
index 5d9647c..0a61e71 100644
--- a/tests/e2e/token_ratelimit_test.go
+++ b/tests/e2e/token_ratelimit_test.go
@@ -13,6 +13,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"strconv"
 	"strings"
 	"testing"
@@ -31,12 +32,12 @@ func Test_Examples_TokenRateLimit(t *testing.T) {
 	const egSelector = "gateway.envoyproxy.io/owning-gateway-name=envoy-ai-gateway-token-ratelimit"
 	requireWaitForPodReady(t, egNamespace, egSelector)
 
+	const modelName = "rate-limit-funky-model"
 	makeRequest := func(usedID string, input, output, total int, expStatus int) {
 		fwd := requireNewHTTPPortForwarder(t, egNamespace, egSelector, egDefaultPort)
 		defer fwd.kill()
 
-		requestBody := fmt.Sprintf(`{"messages":[{"role":"user","content":"Say this is a test"}],"model":"gpt-4o-mini"}`)
-
+		requestBody := fmt.Sprintf(`{"messages":[{"role":"user","content":"Say this is a test"}],"model":"%s"}`, modelName)
 		const fakeResponseBodyTemplate = `{"choices":[{"message":{"content":"This is a test.","role":"assistant"}}],"stopReason":null,"usage":{"prompt_tokens":%d,"completion_tokens":%d,"total_tokens":%d}}`
 		fakeResponseBody := fmt.Sprintf(fakeResponseBodyTemplate, input, output, total)
 
@@ -98,4 +99,48 @@ func Test_Examples_TokenRateLimit(t *testing.T) {
 	makeRequest(usedID, 3, 0, 0, 200)
 	// Any request with the same user ID should be rejected.
 	makeRequest(usedID, 0, 0, 0, 429)
+
+	require.Eventually(t, func() bool {
+		fwd := requireNewHTTPPortForwarder(t, "monitoring", "app=prometheus", 9090)
+		defer fwd.kill()
+		const query = `sum(gen_ai_client_token_usage_sum{app = "ai-eg-route-extproc-envoy-ai-gateway-token-ratelimit"}) by (gen_ai_request_model, gen_ai_token_type)`
+		req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("%s/api/v1/query?query=%s", fwd.address(), url.QueryEscape(query)), nil)
+		require.NoError(t, err)
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			t.Logf("Failed to query Prometheus: %v", err)
+			return false
+		}
+		defer func() { _ = resp.Body.Close() }()
+		body, err := io.ReadAll(resp.Body)
+		t.Logf("Response: status=%d, body=%s", resp.StatusCode, string(body))
+		if resp.StatusCode != http.StatusOK {
+			t.Logf("Failed to query Prometheus: status=%s", resp.Status)
+			return false
+		}
+		type prometheusResponse struct {
+			Status string `json:"status"`
+			Data   struct {
+				ResultType string `json:"resultType"`
+				Result     []struct {
+					Metric map[string]string `json:"metric"`
+					Value  []interface{}     `json:"value"`
+				}
+			}
+		}
+		var pr prometheusResponse
+		require.NoError(t, json.Unmarshal(body, &pr))
+		require.Equal(t, "success", pr.Status)
+		require.Equal(t, "vector", pr.Data.ResultType)
+		var actualTypes []string
+		for _, result := range pr.Data.Result {
+			require.Equal(t, modelName, result.Metric["gen_ai_request_model"])
+			typ := result.Metric["gen_ai_token_type"]
+			actualTypes = append(actualTypes, typ)
+			t.Logf("Type: %s, Value: %v", typ, result.Value)
+		}
+		// We should see input, output, and total token types.
+		require.ElementsMatch(t, []string{"input", "output", "total"}, actualTypes)
+		return true
+	}, 20*time.Second, 1*time.Second)
 }