Merge branch 'main' into gemini-usage

Author: hustxiayang

examples/inference-pool/README.md

@@ -2,10 +2,19 @@
 
 This example demonstrates how to use AI Gateway with the InferencePool feature, which enables intelligent request routing across multiple inference endpoints with load balancing and health checking capabilities.
 
+The setup includes **three distinct backends**:
+
+- Two `InferencePool` resources for LLMs (`Llama-3.1-8B-Instruct` and `Mistral`)
+- One standard `Backend` for non-InferencePool traffic
+
+Routing is controlled by the `x-ai-eg-model` HTTP header.
+
 ## Files in This Directory
 
 - **`envoy-gateway-values-addon.yaml`**: Envoy Gateway values addon for InferencePool support. Combine with `../../manifests/envoy-gateway-values.yaml`.
-- **`base.yaml`**: Complete example that includes Gateway, AIServiceBackend, InferencePool CRDs, and a sample application deployment.
+- **`base.yaml`**: Deploys all inference backends and supporting resources using the **standard approach documented in the official guide**. This includes:
+  - A `mistral` backend with custom Endpoint Picker configuration
+  - A standard fallback backend (`envoy-ai-gateway-basic-testupstream`) for non-InferencePool routing
 - **`aigwroute.yaml`**: Example AIGatewayRoute that uses InferencePool as a backend.
 - **`httproute.yaml`**: Example HTTPRoute for traditional HTTP routing to InferencePool endpoints.
 - **`with-annotations.yaml`**: Advanced example showing InferencePool with Kubernetes annotations for fine-grained control.
@@ -27,16 +36,63 @@ This example demonstrates how to use AI Gateway with the InferencePool feature,
 
    ```bash
    kubectl apply -f base.yaml
+   kubectl apply -f aigwroute.yaml
    ```
 
+   > Note: The `aigwroute.yaml` file defines the InferencePool and routing logic, but does not deploy the actual inference backend (e.g., the vLLM server for Llama-3.1-8B-Instruct).
+   > You must deploy the backend separately by following [Step 3: Deploy Inference Backends](https://aigateway.envoyproxy.io/docs/capabilities/inference/aigatewayroute-inferencepool#step-3-deploy-inference-backends)
+
 3. Test the setup:
 
-   ```bash
-   GATEWAY_HOST=$(kubectl get gateway/ai-gateway -o jsonpath='{.status.addresses[0].value}')
-   curl -X POST "http://${GATEWAY_HOST}/v1/chat/completions" \
-     -H "Content-Type: application/json" \
-     -d '{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello!"}]}'
-   ```
+You can access the gateway in two ways, depending on your environment.
+
+✅ Option A: Using External IP (e.g., cloud LoadBalancer, MetalLB)
+If your cluster assigns an external address to the Gateway:
+
+```bash
+GATEWAY_HOST=$(kubectl get gateway/inference-pool-with-aigwroute -n default -o jsonpath='{.status.addresses[0].value}')
+echo "Gateway available at: http://${GATEWAY_HOST}"
+```
+
+Then send a request:
+
+```bash
+curl -X POST "http://${GATEWAY_HOST}/v1/chat/completions" \
+  -H "x-ai-eg-model: meta-llama/Llama-3.1-8B-Instruct" \
+  -H "Authorization: sk-abcdefghijklmnopqrstuvwxyz" \
+  -H "Content-Type: application/json" \
+  -d '{"model": "meta-llama/Llama-3.1-8B-Instruct", "messages": [{"role": "user", "content": "Hello!"}]}'
+```
+
+✅ Option B: Using kubectl port-forward (ideal for local clusters like Minikube/Kind)
+In one terminal, forward the gateway service:
+
+```bash
+kubectl port-forward svc/envoy-default-inference-pool-with-aigwroute-d416582c 8080:80 -n envoy-gateway-system
+```
+
+In another terminal, send requests to localhost:8080:
+
+```bash
+# Route to Llama (InferencePool)
+curl -X POST "http://localhost:8080/v1/chat/completions" \
+  -H "x-ai-eg-model: meta-llama/Llama-3.1-8B-Instruct" \
+  -H "Authorization: sk-abcdefghijklmnopqrstuvwxyz" \
+  -H "Content-Type: application/json" \
+  -d '{"model": "meta-llama/Llama-3.1-8B-Instruct", "messages": [{"role": "user", "content": "Hello!"}]}'
+
+# Route to Mistral (InferencePool)
+curl -X POST "http://localhost:8080/v1/chat/completions" \
+  -H "x-ai-eg-model: mistral:latest" \
+  -H "Content-Type: application/json" \
+  -d '{"model": "mistral:latest", "messages": [{"role": "user", "content": "Hello!"}]}'
+
+# Route to fallback backend (Standard Backend)
+curl -X POST "http://localhost:8080/v1/chat/completions" \
+  -H "x-ai-eg-model: some-cool-self-hosted-model" \
+  -H "Content-Type: application/json" \
+  -d '{"model": "some-cool-self-hosted-model", "messages": [{"role": "user", "content": "Hello!"}]}'
+```
 
 ### Combining with Other Features
 

internal/extproc/backendauth/anthropicapikey.go renamed to internal/backendauth/anthropicapikey.go

@@ -9,10 +9,8 @@ import (
 	"context"
 	"strings"
 
-	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
-	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
-
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
+	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 )
 
 type anthropicAPIKeyHandler struct {
@@ -27,13 +25,7 @@ func newAnthropicAPIKeyHandler(auth *filterapi.AnthropicAPIKeyAuth) (Handler, er
 // Anthropic uses "x-api-key" header instead of "Authorization: Bearer".
 //
 // https://docs.claude.com/en/api/overview#authentication
-func (a *anthropicAPIKeyHandler) Do(_ context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, _ *extprocv3.BodyMutation) error {
+func (a *anthropicAPIKeyHandler) Do(_ context.Context, requestHeaders map[string]string, _ []byte) ([]internalapi.Header, error) {
 	requestHeaders["x-api-key"] = a.apiKey
-	headerMut.SetHeaders = append(headerMut.SetHeaders, &corev3.HeaderValueOption{
-		Header: &corev3.HeaderValue{
-			Key:      "x-api-key",
-			RawValue: []byte(a.apiKey),
-		},
-	})
-	return nil
+	return []internalapi.Header{{"x-api-key", a.apiKey}}, nil
 }

internal/extproc/backendauth/anthropicapikey_test.go renamed to internal/backendauth/anthropicapikey_test.go

@@ -9,7 +9,6 @@ import (
 	"context"
 	"testing"
 
-	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"github.com/stretchr/testify/require"
 
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
@@ -21,30 +20,31 @@ func TestAnthropicAPIKeyHandler(t *testing.T) {
 		require.NoError(t, err)
 
 		headers := make(map[string]string)
-		headerMut := &extprocv3.HeaderMutation{}
 
-		err = handler.Do(context.Background(), headers, headerMut, nil)
+		hders, err := handler.Do(context.Background(), headers, nil)
 		require.NoError(t, err)
 
 		// Verify header in map
 		require.Equal(t, "test-azure-key", headers["x-api-key"])
 
 		// Verify header in mutation
-		require.Len(t, headerMut.SetHeaders, 1)
-		require.Equal(t, "x-api-key", headerMut.SetHeaders[0].Header.Key)
-		require.Equal(t, "test-azure-key", string(headerMut.SetHeaders[0].Header.RawValue))
+		require.Len(t, hders, 1)
+		require.Equal(t, "x-api-key", hders[0][0])
+		require.Equal(t, "test-azure-key", hders[0][1])
 	})
 
 	t.Run("trims whitespace", func(t *testing.T) {
 		handler, err := newAnthropicAPIKeyHandler(&filterapi.AnthropicAPIKeyAuth{Key: "  key-with-spaces  "})
 		require.NoError(t, err)
 
 		headers := make(map[string]string)
-		headerMut := &extprocv3.HeaderMutation{}
 
-		err = handler.Do(context.Background(), headers, headerMut, nil)
+		hdrs, err := handler.Do(context.Background(), headers, nil)
 		require.NoError(t, err)
 
 		require.Equal(t, "key-with-spaces", headers["x-api-key"])
+		require.Len(t, hdrs, 1)
+		require.Equal(t, "x-api-key", hdrs[0][0])
+		require.Equal(t, "key-with-spaces", hdrs[0][1])
 	})
 }

internal/extproc/backendauth/api_key.go renamed to internal/backendauth/api_key.go

@@ -10,10 +10,8 @@ import (
 	"fmt"
 	"strings"
 
-	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
-	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
-
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
+	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 )
 
 // apiKeyHandler implements [Handler] for api key authz.
@@ -28,10 +26,7 @@ func newAPIKeyHandler(auth *filterapi.APIKeyAuth) (Handler, error) {
 // Do implements [Handler.Do].
 //
 // Extracts the api key from the local file and set it as an authorization header.
-func (a *apiKeyHandler) Do(_ context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, _ *extprocv3.BodyMutation) error {
+func (a *apiKeyHandler) Do(_ context.Context, requestHeaders map[string]string, _ []byte) ([]internalapi.Header, error) {
 	requestHeaders["Authorization"] = fmt.Sprintf("Bearer %s", a.apiKey)
-	headerMut.SetHeaders = append(headerMut.SetHeaders, &corev3.HeaderValueOption{
-		Header: &corev3.HeaderValue{Key: "Authorization", RawValue: []byte(requestHeaders["Authorization"])},
-	})
-	return nil
+	return []internalapi.Header{{"Authorization", fmt.Sprintf("Bearer %s", a.apiKey)}}, nil
 }

internal/extproc/backendauth/api_key_test.go renamed to internal/backendauth/api_key_test.go

@@ -8,8 +8,6 @@ package backendauth
 import (
 	"testing"
 
-	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
-	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"github.com/stretchr/testify/require"
 
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
@@ -30,28 +28,15 @@ func TestApiKeyHandler_Do(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, handler)
 
-	requestHeaders := map[string]string{":method": "POST"}
-	headerMut := &extprocv3.HeaderMutation{
-		SetHeaders: []*corev3.HeaderValueOption{
-			{Header: &corev3.HeaderValue{
-				Key:   ":path",
-				Value: "/model/some-random-model/converse",
-			}},
-		},
-	}
-	bodyMut := &extprocv3.BodyMutation{
-		Mutation: &extprocv3.BodyMutation_Body{
-			Body: []byte(`{"messages": [{"role": "user", "content": [{"text": "Say this is a test!"}]}]}`),
-		},
-	}
-	err = handler.Do(t.Context(), requestHeaders, headerMut, bodyMut)
+	requestHeaders := map[string]string{":method": "POST", ":path": "/model/some-random-model/converse"}
+	hdrs, err := handler.Do(t.Context(), requestHeaders, nil)
 	require.NoError(t, err)
 
 	bearerToken, ok := requestHeaders["Authorization"]
 	require.True(t, ok)
 	require.Equal(t, "Bearer test", bearerToken)
 
-	require.Len(t, headerMut.SetHeaders, 2)
-	require.Equal(t, "Authorization", headerMut.SetHeaders[1].Header.Key)
-	require.Equal(t, []byte("Bearer test"), headerMut.SetHeaders[1].Header.GetRawValue())
+	require.Len(t, hdrs, 1)
+	require.Equal(t, "Authorization", hdrs[0][0])
+	require.Equal(t, "Bearer test", hdrs[0][1])
 }

internal/extproc/backendauth/auth.go renamed to internal/backendauth/auth.go

@@ -9,17 +9,17 @@ import (
 	"context"
 	"errors"
 
-	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
-
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
+	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 )
 
 // Handler is the interface that deals with the backend auth for a specific backend.
 //
 // TODO: maybe this can be just "post-transformation" handler, as it is not really only about auth.
 type Handler interface {
-	// Do performs the backend auth, and make changes to the request headers and body mutations.
-	Do(ctx context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error
+	// Do performs the backend auth, and make changes to the request headers passed in as `requestHeaders`.
+	// It also returns a list of headers that were added or modified as a slice of key-value pairs.
+	Do(ctx context.Context, requestHeaders map[string]string, mutatedBody []byte) ([]internalapi.Header, error)
 }
 
 // NewHandler returns a new implementation of [Handler] based on the configuration.

internal/extproc/backendauth/auth_test.go renamed to internal/backendauth/auth_test.go

@@ -15,15 +15,13 @@ import (
 	"os"
 	"strings"
 	"time"
-	"unsafe"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
 	"github.com/aws/aws-sdk-go-v2/config"
-	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
-	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
+	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 )
 
 // awsHandler implements [Handler] for AWS Bedrock authz.
@@ -84,34 +82,21 @@ func newAWSHandler(ctx context.Context, awsAuth *filterapi.AWSAuth) (Handler, er
 //
 // This assumes that during the transformation, the path is set in the header mutation as well as
 // the body in the body mutation.
-func (a *awsHandler) Do(ctx context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error {
+func (a *awsHandler) Do(ctx context.Context, requestHeaders map[string]string, mutatedBody []byte) ([]internalapi.Header, error) {
 	method := requestHeaders[":method"]
-	path := ""
-	if headerMut.SetHeaders != nil {
-		for _, h := range headerMut.SetHeaders {
-			if h.Header.Key == ":path" {
-				if len(h.Header.Value) > 0 {
-					path = h.Header.Value
-				} else {
-					rv := h.Header.RawValue
-					path = unsafe.String(&rv[0], len(rv))
-				}
-				break
-			}
-		}
-	}
+	path := requestHeaders[":path"]
 
 	var body []byte
-	if _body := bodyMut.GetBody(); len(_body) > 0 {
-		body = _body
+	if len(mutatedBody) > 0 {
+		body = mutatedBody
 	}
 
 	payloadHash := sha256.Sum256(body)
 	req, err := http.NewRequest(method,
 		fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com%s", a.region, path),
 		bytes.NewReader(body))
 	if err != nil {
-		return fmt.Errorf("cannot create request: %w", err)
+		return nil, fmt.Errorf("cannot create request: %w", err)
 	}
 	// By setting the content length to -1, we can avoid the inclusion of the `Content-Length` header in the signature.
 	// https://github.com/aws/aws-sdk-go-v2/blob/755839b2eebb246c7eec79b65404aee105196d5b/aws/signer/v4/v4.go#L427-L431
@@ -124,21 +109,21 @@ func (a *awsHandler) Do(ctx context.Context, requestHeaders map[string]string, h
 
 	credentials, err := a.credentialsProvider.Retrieve(ctx)
 	if err != nil {
-		return fmt.Errorf("cannot retrieve AWS credentials: %w", err)
+		return nil, fmt.Errorf("cannot retrieve AWS credentials: %w", err)
 	}
 
 	err = a.signer.SignHTTP(ctx, credentials, req,
 		hex.EncodeToString(payloadHash[:]), "bedrock", a.region, time.Now())
 	if err != nil {
-		return fmt.Errorf("cannot sign request: %w", err)
+		return nil, fmt.Errorf("cannot sign request: %w", err)
 	}
 
+	var headers []internalapi.Header
 	for key, hdr := range req.Header {
 		if key == "Authorization" || strings.HasPrefix(key, "X-Amz-") {
-			headerMut.SetHeaders = append(headerMut.SetHeaders, &corev3.HeaderValueOption{
-				Header: &corev3.HeaderValue{Key: key, RawValue: []byte(hdr[0])}, // Assume aws-go-sdk always returns a single value.
-			})
+			headers = append(headers, internalapi.Header{key, hdr[0]})
+			requestHeaders[key] = hdr[0]
 		}
 	}
-	return nil
+	return headers, nil
 }