feat: end-to-end OpenAI Image Generation support

.env.ollama

@@ -2,3 +2,4 @@ CHAT_MODEL=qwen2.5:0.5b
 THINKING_MODEL=qwen3:1.7b
 COMPLETION_MODEL=qwen2.5:0.5b
 EMBEDDINGS_MODEL=all-minilm:33m
+IMAGE_GENERATION_MODEL=dall-e-2

.gitignore

@@ -49,3 +49,4 @@ inference-extension-conformance-test-report.yaml
 .mcp.json
 
 .goose
+/aigw

cmd/aigw/README.md

@@ -56,12 +56,38 @@ Here are values we use for Ollama:
      ```
 
    - MCP (Model Context Protocol) tool call:
+
      ```bash
      docker compose run --rm mcp
      ```
+
      This calls the kiwi MCP server through aigw's MCP Gateway at `/mcp`.
 
-4. **Shutdown the example stack**:
+   - Image generation:
+     - Using service:
+       ```bash
+       docker compose run --rm image-generation
+       ```
+     - Using curl (save to file):
+       ```bash
+       curl -s \
+         -X POST http://localhost:1975/v1/images/generations \
+         -H "Authorization: Bearer unused" \
+         -H "Content-Type: application/json" \
+         -d '{"model":"gpt-image-1-mini","prompt":"A watercolor painting of a red fox in a birch forest","size":"1024x1024","quality":"low"}' \
+         | jq -r '.data[0].b64_json' | base64 -d > image.png
+       ```
+
+4. **Create embeddings**:
+
+   The `embeddings` service uses `curl` to send an embeddings request
+   to the AI Gateway CLI (aigw) which routes it to Ollama.
+
+   ```bash
+   docker compose run --rm embeddings
+   ```
+
+5. **Shutdown the example stack**:
 
    `down` stops the containers and removes the volumes used by the stack.
 
@@ -156,8 +182,9 @@ This configures the OTLP endpoint to otel-tui on port 4318.
    ```bash
    COMPOSE_PROFILES="<profile>" docker compose -f docker-compose-otel.yaml run --build --rm chat-completion
    COMPOSE_PROFILES="<profile>" docker compose -f docker-compose-otel.yaml run --build --rm create-embeddings
-   COMPOSE_PROFILES="<profile>" docker compose -f docker-compose-otel.yaml run --build --rm completion
    COMPOSE_PROFILES="<profile>" docker compose -f docker-compose-otel.yaml run --build --rm mcp
+   # Image generation
+   COMPOSE_PROFILES="<profile>" docker compose -f docker-compose-otel.yaml run --build --rm image-generation
    ```
 
 3. **Check telemetry output**:

cmd/aigw/config_test.go

@@ -84,6 +84,7 @@ func TestReadConfig(t *testing.T) {
 			// Clear any existing env vars
 			t.Setenv("OPENAI_API_KEY", "")
 			t.Setenv("OPENAI_BASE_URL", "")
+			t.Setenv("AZURE_OPENAI_API_KEY", "")
 
 			for k, v := range tt.envVars {
 				t.Setenv(k, v)
@@ -104,6 +105,9 @@ func TestReadConfig(t *testing.T) {
 	}
 
 	t.Run("error when file and no OPENAI_API_KEY", func(t *testing.T) {
+		// Ensure both OpenAI and Azure keys are unset so readConfig errors
+		t.Setenv("OPENAI_API_KEY", "")
+		t.Setenv("AZURE_OPENAI_API_KEY", "")
 		_, err := readConfig("", nil, false)
 		require.Error(t, err)
 		require.EqualError(t, err, "you must supply at least OPENAI_API_KEY or AZURE_OPENAI_API_KEY or a config file path")

cmd/aigw/docker-compose-otel.yaml

@@ -123,24 +123,6 @@ services:
       - OPENAI_BASE_URL=http://aigw:1975/v1
       - OPENAI_API_KEY=unused
 
-  # completion is the standard OpenAI client (`openai` in pip), instrumented
-  # with the following OpenTelemetry instrumentation libraries:
-  # - openinference-instrumentation-openai (completions spans)
-  # - opentelemetry-instrumentation-httpx (HTTP client spans and trace headers)
-  completion:
-    build:
-      context: ../../tests/internal/testopeninference
-      dockerfile: Dockerfile.openai_client
-      target: completion
-    container_name: completion
-    profiles: ["test"]
-    env_file:
-      - ../../.env.ollama
-      - .env.otel.${COMPOSE_PROFILES:-console}
-    environment:
-      - OPENAI_BASE_URL=http://aigw:1975/v1
-      - OPENAI_API_KEY=unused
-
   # mcp is a test client for calling MCP tools through aigw.
   # TODO: add client tracing + mcp propagation to this
   mcp:
@@ -162,3 +144,21 @@ services:
       - flyTo=LAX
       - --tool-arg
       - departureDate=15/12/2025
+
+  # image-generation is a simple curl-based test client for sending image
+  # generation requests to aigw.
+  image-generation:
+    image: golang:1.25
+    container_name: image-generation
+    profiles: ["test"]
+    env_file:
+      - ../../.env.ollama
+    command:
+      - sh
+      - -c
+      - |
+        curl -s -w %{http_code} \
+          -X POST http://aigw:1975/v1/images/generations \
+          -H "Authorization: Bearer unused" \
+          -H "Content-Type: application/json" \
+          -d "{\"model\":\"$$IMAGE_GENERATION_MODEL\",\"prompt\":\"A watercolor painting of a red fox in a birch forest\",\"size\":\"1024x1024\",\"quality\":\"low\"}"

cmd/aigw/docker-compose.yaml

@@ -135,3 +135,23 @@ services:
       - flyTo=LAX
       - --tool-arg
       - departureDate=15/12/2025
+
+  # image-generation is a simple curl-based test client for sending image
+  # generation requests to aigw.
+  image-generation:
+    image: golang:1.25
+    container_name: image-generation
+    profiles: ["test"]
+    env_file:
+      - ../../.env.ollama
+    command:
+      - sh
+      - -c
+      - |
+        curl -s -w %{http_code} \
+          -X POST http://aigw:1975/v1/images/generations \
+          -H "Authorization: Bearer unused" \
+          -H "Content-Type: application/json" \
+          -d "{\"model\":\"$$IMAGE_GENERATION_MODEL\",\"prompt\":\"A watercolor painting of a red fox in a birch forest\",\"size\":\"1024x1024\",\"quality\":\"low\"}"
+    extra_hosts: # localhost:host-gateway trick doesn't work with aigw
+      - "host.docker.internal:host-gateway"

cmd/aigw/main_test.go

@@ -71,6 +71,7 @@ Flags:
 		{
 			name:         "run no arg",
 			args:         []string{"run"},
+			env:          map[string]string{"OPENAI_API_KEY": "", "AZURE_OPENAI_API_KEY": ""},
 			rf:           func(context.Context, cmdRun, runOpts, io.Writer, io.Writer) error { return nil },
 			expPanicCode: ptr.To(80),
 		},
@@ -190,6 +191,9 @@ func TestCmdRun_Validate(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			// Ensure a clean environment for validation.
+			t.Setenv("OPENAI_API_KEY", "")
+			t.Setenv("AZURE_OPENAI_API_KEY", "")
 			for k, v := range tt.envVars {
 				t.Setenv(k, v)
 			}

cmd/extproc/mainlib/main.go

@@ -234,6 +234,8 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) {
 	messagesMetrics := metrics.NewMessagesFactory(meter, metricsRequestHeaderAttributes)
 	completionMetrics := metrics.NewCompletionFactory(meter, metricsRequestHeaderAttributes)
 	embeddingsMetrics := metrics.NewEmbeddingsFactory(meter, metricsRequestHeaderAttributes)
+	imageGenerationMetrics := metrics.NewImageGenerationFactory(meter, metricsRequestHeaderAttributes)()
+
 	mcpMetrics := metrics.NewMCP(meter, metricsRequestHeaderAttributes)
 
 	tracing, err := tracing.NewTracingFromEnv(ctx, os.Stdout, spanRequestHeaderAttributes)
@@ -248,6 +250,7 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) {
 	server.Register(path.Join(flags.rootPrefix, "/v1/chat/completions"), extproc.ChatCompletionProcessorFactory(chatCompletionMetrics))
 	server.Register(path.Join(flags.rootPrefix, "/v1/completions"), extproc.CompletionsProcessorFactory(completionMetrics))
 	server.Register(path.Join(flags.rootPrefix, "/v1/embeddings"), extproc.EmbeddingsProcessorFactory(embeddingsMetrics))
+	server.Register(path.Join(flags.rootPrefix, "/v1/images/generations"), extproc.ImageGenerationProcessorFactory(imageGenerationMetrics))
 	server.Register(path.Join(flags.rootPrefix, "/v1/models"), extproc.NewModelsProcessor)
 	server.Register(path.Join(flags.rootPrefix, "/anthropic/v1/messages"), extproc.MessagesProcessorFactory(messagesMetrics))
 

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.0
 	github.com/a8m/envsubst v1.4.3
 	github.com/alecthomas/kong v1.12.1
+	github.com/andybalholm/brotli v1.2.0
 	github.com/anthropics/anthropic-sdk-go v1.14.0
 	github.com/aws/aws-sdk-go-v2 v1.39.3
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.2

go.sum

@@ -41,6 +41,8 @@ github.com/alecthomas/kong v1.12.1 h1:iq6aMJDcFYP9uFrLdsiZQ2ZMmcshduyGv4Pek0MQPW
 github.com/alecthomas/kong v1.12.1/go.mod h1:p2vqieVMeTAnaC83txKtXe8FLke2X07aruPWXyMPQrU=
 github.com/alecthomas/repr v0.4.0 h1:GhI2A8MACjfegCPVq9f1FLvIBS+DrQ2KQBFZP1iFzXc=
 github.com/alecthomas/repr v0.4.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4=
+github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ=
+github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
 github.com/anthropics/anthropic-sdk-go v1.14.0 h1:EzNQvnZlaDHe2UPkoUySDz3ixRgNbwKdH8KtFpv7pi4=
 github.com/anthropics/anthropic-sdk-go v1.14.0/go.mod h1:WTz31rIUHUHqai2UslPpw5CwXrQP3geYBioRV4WOLvE=
 github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ=
@@ -448,6 +450,8 @@ github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
 github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
 github.com/xiang90/probing v0.0.0-20221125231312-a49e3df8f510 h1:S2dVYn90KE98chqDkyE9Z4N61UnQd+KOfgp5Iu53llk=
 github.com/xiang90/probing v0.0.0-20221125231312-a49e3df8f510/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU=
+github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
+github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
 github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

internal/apischema/openai/openai.go

@@ -57,6 +57,10 @@ const (
 
 	// ModelTextEmbedding3Small is the cheapest model usable with /embeddings.
 	ModelTextEmbedding3Small = "text-embedding-3-small"
+
+	// ModelGPTImage1Mini is the smallest/cheapest Images model usable with
+	// /v1/images/generations. Use with size "1024x1024" and quality "low".
+	ModelGPTImage1Mini = "gpt-image-1-mini"
 )
 
 // ChatCompletionContentPartRefusalType The type of the content part.