Skip to content

Write an AI policy for the AI Gateway

An AI policy inspects, controls, and modifies the traffic going to and from a large language model (LLM) provider such as OpenAI or Anthropic.

AI policies use the same Policy interface as standard gateway policies. For full API details, see the policy SDK reference.

The key difference is how you handle LLM request and response bodies, especially:

  • JSON responses
  • Streaming responses (SSE)

How it works

Every request and response that flows through the gateway passes through a policy chain. Each policy declares which phases it participates in, and the kernel calls the appropriate hook for each phase:

Incoming Request
       │
       ▼
  Request Headers  ──► OnRequestHeaders()
       │
       ▼
  Request Body     ──► OnRequestBody()  (or OnRequestBodyChunk() for streaming)
       │
       ▼
   Upstream LLM
       │
       ▼
  Response Headers ──► OnResponseHeaders()
       │
       ▼
  Response Body    ──► OnResponseBody() (or OnResponseBodyChunk() for streaming)
       │
       ▼
  Downstream Client

Note

A policy chain is an ordered sequence of policies that the gateway runs on every request and response for a given LLM provider or App LLM proxy. Policies run in the order the runtime configuration lists them, and each policy sees the modifications the ones before it made.

Key idea

LLM responses come in two formats:

Mode Format
Non-streaming Single JSON object
Streaming SSE events (data: {...})

Your policy must handle both formats.

Which interfaces to implement

Choose based on what your policy needs to do:

Goal Interface Mode setting
Inspect prompt / model RequestPolicy RequestBodyMode: BodyModeBuffer
Inspect headers (auth, routing) RequestHeaderPolicy RequestHeaderMode: HeaderModeProcess
Inspect or modify buffered (in-memory) response ResponsePolicy ResponseBodyMode: BodyModeBuffer
Inspect or modify streaming response StreamingResponsePolicy (embeds ResponsePolicy) ResponseBodyMode: BodyModeStream

How to write an AI policy

Step 1: Create the policy

Each policy lives in its own Go module. Create a "policies" directory inside your gateway:

/policies/my-ai-policy/
 ├── go.mod
 ├── my_ai_policy.go
 └── policy-definition.yaml

Step 2: Implement the mode

Mode() declares which phases this policy participates in and how bodies are handled. The kernel reads this once at startup — there is no per-request overhead.

package myaipolicy

import (
    "context"

    policy "github.com/wso2/api-platform/sdk/core/policy/v1alpha2"
)

type MyAIPolicy struct{}

func (p *MyAIPolicy) Mode() policy.ProcessingMode {
    return policy.ProcessingMode{
        RequestBodyMode:  policy.BodyModeBuffer,
        ResponseBodyMode: policy.BodyModeStream,
    }
}

Tip

If your policy doesn't need to inspect a phase, set that phase to HeaderModeSkip or BodyModeSkip.

Step 3: Implement request inspection

OnRequestBody is called once the request body is fully buffered. Use it to inspect the model name, messages, or parameters before the request reaches the LLM provider.

func (p *MyAIPolicy) OnRequestBody(
    ctx context.Context,
    reqCtx *policy.RequestContext,
    params map[string]interface{},
) policy.RequestAction {
    // Inspect model + messages
    return nil
}

Step 4: Implement response handling

For most AI policies, implement both:

  • ResponsePolicy: handles buffered responses, where the entire response is available at once. That's either a non-streaming JSON response or concatenated SSE events.
  • StreamingResponsePolicy: handles streaming responses, either JSON or SSE events.

Tip

The gateway automatically chooses which handler to call. The gateway calls OnResponseBodyChunk only when the entire policy chain is streaming-compatible. If any policy in the chain doesn't support streaming, the gateway falls back to OnResponseBody, so implement both even when streaming is your primary target.

// Streaming Response Handling
func (p *MyAIPolicy) OnResponseBodyChunk(
    ctx context.Context,
    respCtx *policy.ResponseStreamContext,
    chunk *policy.StreamBody,
    params map[string]interface{},
) policy.StreamingResponseAction {
    // Accumulate + process
    return policy.ForwardResponseChunk{}
}

// Gating response chunks before processing response
func (p *MyAIPolicy) NeedsMoreResponseData(_ []byte) bool {
    return false
}

// Buffered Fallback
func (p *MyAIPolicy) OnResponseBody(
    ctx context.Context,
    respCtx *policy.ResponseContext,
    params map[string]interface{},
) policy.ResponseAction {
    // Same logic as streaming, applied to the full body
    return nil
}

Gate-then-stream pattern

A common pattern for AI guardrails is to accumulate chunks until you have a complete SSE event to inspect, then switch to pass-through:

// Buffer until we can parse a complete SSE event, then stream freely
func (p *MyAIPolicy) NeedsMoreResponseData(accumulated []byte) bool {
    return !bytes.Contains(accumulated, []byte("\n\n"))
}

Step 5: Factory function

Initialize your policy and validate parameters:

func GetPolicy(
    metadata policy.PolicyMetadata,
    params map[string]interface{},
) (policy.Policy, error) {

    threshold, ok := params["blockThreshold"].(float64)
    if !ok {
        return nil, fmt.Errorf("invalid blockThreshold")
    }

    return &MyAIPolicy{blockThreshold: threshold}, nil
}

Step 6: Define parameters

Create a policy-definition.yaml in your policy directory:

name: my-ai-policy
displayName: my ai policy
version: v1.0.0

parameters:
  type: object
  properties:
    blockThreshold:
      type: number
      default: 0.8

Step 7: Share data between phases

Use the Metadata map to pass data between request and response phases — for example, the model name read from the request, used later to apply model-specific logic in the response phase:

// In request phase
reqCtx.Metadata["model"] = model

// In response phase
model := respCtx.Metadata["model"]

Step 8: Register and build

Add your policy to the gateway folder's build.yaml under policies: using filePath for local development:

policies:
  - name: my-ai-policy
    filePath: ./policies/my-ai-policy

For published policies (production), use the module reference instead:

policies:
  - name: my-ai-policy
    gomodule: github.com/abc/policy-repo/policies/my-ai-policy@v1

Best practices

  • Always handle both streaming and non-streaming. The gateway falls back to buffered mode when any policy in the chain doesn't support streaming.
  • Use the Metadata map to share state. Pass data between the request and response phases through it.
  • Implement streaming and the buffered fallback. Your policy then works whether the chain runs in streaming or buffered mode.
  • Parse SSE incrementally. When gating on streaming responses, buffer only until you have a complete SSE event, terminated by \n\n, rather than the entire response. This keeps latency low.

Next steps