backport: completion stream + metrics and assistant content (#497)

**Commit Message**

PR to backport `mockChatCompletionMetrics`, chat completion stream fix,
and openai content type.

Including:
- #459 (468 uses mock components introduced here)
- #468 
- #486

---------

Signed-off-by: Huamin Chen <[email protected]>
Signed-off-by: Ignasi Barrera <[email protected]>
Signed-off-by: Takeshi Yoneda <[email protected]>
Signed-off-by: Aaron Choo <[email protected]>
Co-authored-by: Ignasi Barrera <[email protected]>
Co-authored-by: Takeshi Yoneda <[email protected]>
Co-authored-by: Dan Sun <[email protected]>

Author: 4 people

cmd/extproc/mainlib/main.go

@@ -13,17 +13,24 @@ import (
 	"log"
 	"log/slog"
 	"net"
+	"net/http"
 	"os"
 	"os/signal"
 	"strings"
 	"syscall"
 	"time"
 
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	otelprom "go.opentelemetry.io/otel/exporters/prometheus"
+	"go.opentelemetry.io/otel/metric"
+	metricsdk "go.opentelemetry.io/otel/sdk/metric"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health/grpc_health_v1"
 
 	"github.com/envoyproxy/ai-gateway/internal/extproc"
+	"github.com/envoyproxy/ai-gateway/internal/metrics"
 	"github.com/envoyproxy/ai-gateway/internal/version"
 )
 
@@ -32,6 +39,7 @@ type extProcFlags struct {
 	configPath  string     // path to the configuration file.
 	extProcAddr string     // gRPC address for the external processor.
 	logLevel    slog.Level // log level for the external processor.
+	metricsAddr string     // HTTP address for the metrics server.
 }
 
 // parseAndValidateFlags parses and validates the flas passed to the external processor.
@@ -51,13 +59,14 @@ func parseAndValidateFlags(args []string) (extProcFlags, error) {
 	fs.StringVar(&flags.extProcAddr,
 		"extProcAddr",
 		":1063",
-		"gRPC address for the external processor. For example, :1063 or unix:///tmp/ext_proc.sock",
+		"gRPC address for the external processor. For example, :1063 or unix:///tmp/ext_proc.sock.",
 	)
 	logLevelPtr := fs.String(
 		"logLevel",
 		"info",
 		"log level for the external processor. One of 'debug', 'info', 'warn', or 'error'.",
 	)
+	fs.StringVar(&flags.metricsAddr, "metricsAddr", ":9190", "HTTP address for the metrics server.")
 
 	if err := fs.Parse(args); err != nil {
 		return extProcFlags{}, fmt.Errorf("failed to parse extProcFlags: %w", err)
@@ -102,11 +111,13 @@ func Main() {
 		log.Fatalf("failed to listen: %v", err)
 	}
 
+	metricsServer, meter := startMetricsServer(flags.metricsAddr, l)
+
 	server, err := extproc.NewServer(l)
 	if err != nil {
 		log.Fatalf("failed to create external processor server: %v", err)
 	}
-	server.Register("/v1/chat/completions", extproc.NewChatCompletionProcessor)
+	server.Register("/v1/chat/completions", extproc.ChatCompletionProcessorFactory(metrics.NewChatCompletion(meter)))
 	server.Register("/v1/models", extproc.NewModelsProcessor)
 
 	if err := extproc.StartConfigWatcher(ctx, flags.configPath, server, l, time.Second*5); err != nil {
@@ -116,10 +127,18 @@ func Main() {
 	s := grpc.NewServer()
 	extprocv3.RegisterExternalProcessorServer(s, server)
 	grpc_health_v1.RegisterHealthServer(s, server)
+
 	go func() {
 		<-ctx.Done()
 		s.GracefulStop()
+
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := metricsServer.Shutdown(shutdownCtx); err != nil {
+			l.Error("Failed to shutdown metrics server gracefully", "error", err)
+		}
 	}()
+
 	_ = s.Serve(lis)
 }
 
@@ -130,3 +149,46 @@ func listenAddress(addrFlag string) (string, string) {
 	}
 	return "tcp", addrFlag
 }
+
+// startMetricsServer starts the HTTP server for Prometheus metrics.
+func startMetricsServer(addr string, logger *slog.Logger) (*http.Server, metric.Meter) {
+	registry := prometheus.NewRegistry()
+	exporter, err := otelprom.New(otelprom.WithRegisterer(registry))
+	if err != nil {
+		log.Fatal("failed to create metrics exporter")
+	}
+	provider := metricsdk.NewMeterProvider(metricsdk.WithReader(exporter))
+	meter := provider.Meter("envoyproxy/ai-gateway")
+
+	// Create a new HTTP server for metrics.
+	mux := http.NewServeMux()
+
+	// Register the metrics handler.
+	mux.Handle("/metrics", promhttp.HandlerFor(
+		registry,
+		promhttp.HandlerOpts{
+			EnableOpenMetrics: true,
+		},
+	))
+
+	// Add a simple health check endpoint.
+	mux.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("OK"))
+	})
+
+	server := &http.Server{
+		Addr:              addr,
+		Handler:           mux,
+		ReadHeaderTimeout: 5 * time.Second,
+	}
+
+	go func() {
+		logger.Info("Starting metrics server", "address", addr)
+		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			logger.Error("Metrics server failed", "error", err)
+		}
+	}()
+
+	return server, meter
+}

cmd/extproc/mainlib/main_test.go

@@ -6,7 +6,9 @@
 package mainlib
 
 import (
+	"io"
 	"log/slog"
+	"net/http"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -104,3 +106,19 @@ func TestListenAddress(t *testing.T) {
 		})
 	}
 }
+
+func TestStartMetricsServer(t *testing.T) {
+	s, m := startMetricsServer("127.0.0.1:", slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{})))
+	t.Cleanup(func() { _ = s.Shutdown(t.Context()) })
+
+	require.NotNil(t, s)
+	require.NotNil(t, m)
+
+	require.HTTPStatusCode(t, s.Handler.ServeHTTP, http.MethodGet, "/", nil, http.StatusNotFound)
+
+	require.HTTPSuccess(t, s.Handler.ServeHTTP, http.MethodGet, "/health", nil)
+	require.HTTPBodyContains(t, s.Handler.ServeHTTP, http.MethodGet, "/health", nil, "OK")
+
+	require.HTTPSuccess(t, s.Handler.ServeHTTP, http.MethodGet, "/metrics", nil)
+	require.HTTPBodyContains(t, s.Handler.ServeHTTP, http.MethodGet, "/metrics", nil, "target_info{")
+}

go.mod

@@ -13,8 +13,13 @@ require (
 	github.com/go-logr/logr v1.4.2
 	github.com/google/cel-go v0.23.2
 	github.com/google/go-cmp v0.7.0
-	github.com/openai/openai-go v0.1.0-alpha.59
+	github.com/openai/openai-go v0.1.0-alpha.61
+	github.com/prometheus/client_golang v1.20.5
 	github.com/stretchr/testify v1.10.0
+	go.opentelemetry.io/otel v1.34.0
+	go.opentelemetry.io/otel/exporters/prometheus v0.56.0
+	go.opentelemetry.io/otel/metric v1.34.0
+	go.opentelemetry.io/otel/sdk/metric v1.34.0
 	go.uber.org/goleak v1.3.0
 	go.uber.org/zap v1.27.0
 	golang.org/x/exp v0.0.0-20250128182459-e0ece0dbea4c
@@ -271,7 +276,6 @@ require (
 	github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/polyfloyd/go-errorlint v1.7.1 // indirect
-	github.com/prometheus/client_golang v1.20.5 // indirect
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/common v0.62.0 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect
@@ -343,8 +347,7 @@ require (
 	go-simpler.org/sloglint v0.9.0 // indirect
 	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.58.0 // indirect
-	go.opentelemetry.io/otel v1.34.0 // indirect
-	go.opentelemetry.io/otel/metric v1.34.0 // indirect
+	go.opentelemetry.io/otel/sdk v1.34.0 // indirect
 	go.opentelemetry.io/otel/trace v1.34.0 // indirect
 	go.uber.org/automaxprocs v1.6.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect

go.sum

@@ -632,8 +632,8 @@ github.com/onsi/ginkgo/v2 v2.22.2 h1:/3X8Panh8/WwhU/3Ssa6rCKqPLuAkVY2I0RoyDLySlU
 github.com/onsi/ginkgo/v2 v2.22.2/go.mod h1:oeMosUL+8LtarXBHu/c0bx2D/K9zyQ6uX3cTyztHwsk=
 github.com/onsi/gomega v1.36.2 h1:koNYke6TVk6ZmnyHrCXba/T/MoLBXFjeC1PtvYgw0A8=
 github.com/onsi/gomega v1.36.2/go.mod h1:DdwyADRjrc825LhMEkD76cHR5+pUnjhUN8GlHlRPHzY=
-github.com/openai/openai-go v0.1.0-alpha.59 h1:T3IYwKSCezfIlL9Oi+CGvU03fq0RoH33775S78Ti48Y=
-github.com/openai/openai-go v0.1.0-alpha.59/go.mod h1:3SdE6BffOX9HPEQv8IL/fi3LYZ5TUpRYaqGQZbyk11A=
+github.com/openai/openai-go v0.1.0-alpha.61 h1:dLJW1Dk15VAwm76xyPsiPt/Ky94NNGoMLETAI1ISoBY=
+github.com/openai/openai-go v0.1.0-alpha.61/go.mod h1:3SdE6BffOX9HPEQv8IL/fi3LYZ5TUpRYaqGQZbyk11A=
 github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.1.0 h1:8SG7/vwALn54lVB/0yZ/MMwhFrPYtpEHQb2IpWsCzug=

internal/apischema/openai/openai.go

@@ -10,6 +10,7 @@ package openai
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"strconv"
 	"strings"
@@ -149,6 +150,28 @@ func (c *ChatCompletionContentPartUserUnionParam) UnmarshalJSON(data []byte) err
 	return nil
 }
 
+type StringOrAssistantRoleContentUnion struct {
+	Value interface{}
+}
+
+func (s *StringOrAssistantRoleContentUnion) UnmarshalJSON(data []byte) error {
+	var str string
+	err := json.Unmarshal(data, &str)
+	if err == nil {
+		s.Value = str
+		return nil
+	}
+
+	var content ChatCompletionAssistantMessageParamContent
+	err = json.Unmarshal(data, &content)
+	if err == nil {
+		s.Value = content
+		return nil
+	}
+
+	return errors.New("cannot unmarshal JSON data as string or assistant content parts")
+}
+
 type StringOrArray struct {
 	Value interface{}
 }
@@ -335,7 +358,7 @@ type ChatCompletionAssistantMessageParam struct {
 	Audio ChatCompletionAssistantMessageParamAudio `json:"audio,omitempty"`
 	// The contents of the assistant message. Required unless `tool_calls` or
 	// `function_call` is specified.
-	Content ChatCompletionAssistantMessageParamContent `json:"content"`
+	Content StringOrAssistantRoleContentUnion `json:"content"`
 	// An optional name for the participant. Provides the model information to
 	// differentiate between participants of the same role.
 	Name string `json:"name,omitempty"`

internal/apischema/openai/openai_test.go

@@ -153,7 +153,35 @@ func TestOpenAIChatCompletionMessageUnmarshal(t *testing.T) {
 					{
 						Value: ChatCompletionAssistantMessageParam{
 							Role:    ChatMessageRoleAssistant,
-							Content: ChatCompletionAssistantMessageParamContent{Text: ptr.To("you are a helpful assistant")},
+							Content: StringOrAssistantRoleContentUnion{Value: ChatCompletionAssistantMessageParamContent{Text: ptr.To("you are a helpful assistant")}},
+						},
+						Type: ChatMessageRoleAssistant,
+					},
+				},
+			},
+		},
+		{
+			name: "assistant message string",
+			in: []byte(`{"model": "gpu-o4",
+                        "messages": [
+                         {"role": "assistant", "content": "you are a helpful assistant"},
+                         {"role": "assistant", "content": "{'text': 'you are a helpful assistant'}"}
+						 ]}
+`),
+			out: &ChatCompletionRequest{
+				Model: "gpu-o4",
+				Messages: []ChatCompletionMessageParamUnion{
+					{
+						Value: ChatCompletionAssistantMessageParam{
+							Role:    ChatMessageRoleAssistant,
+							Content: StringOrAssistantRoleContentUnion{Value: "you are a helpful assistant"},
+						},
+						Type: ChatMessageRoleAssistant,
+					},
+					{
+						Value: ChatCompletionAssistantMessageParam{
+							Role:    ChatMessageRoleAssistant,
+							Content: StringOrAssistantRoleContentUnion{Value: "{'text': 'you are a helpful assistant'}"},
 						},
 						Type: ChatMessageRoleAssistant,
 					},