Streaming with AI Gateway

In an LLM (Large Language Model) inference request, AI Gateway uses the upstream AI Model Provider’s REST API to generate the next chat message from the caller.

Normally, this request is processed and completely buffered by the LLM before being sent back to AI Gateway and then to the caller in a single large JSON block. This process can be time-consuming, depending on the max_tokens, other request parameters, and the complexity of the request sent to the LLM model. Request streaming in AI Gateway uses the AI Model entity.

To avoid making the user wait for their chat response with a loading animation, most models can stream each word (or sets of words and tokens) back to the client. This allows the chat response to be rendered in real time.

For example, a client could set up their streaming request using the OpenAI Python SDK like this:

from openai import OpenAI

client = OpenAI(
    base_url="http://127.0.0.1:8000/12/openai",
    api_key="none"
)

stream = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Tell me the history of Kong Inc."}],
    stream=True,
)

for chunk in stream:
    if chunk.choices and chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

A client configured to use streaming won’t have to wait for the entire response. Instead, tokens will appear as they come in.

In streaming mode, a client can set "stream": true in their request, and the LLM server will stream each part of the response text (usually token-by-token) as a server-sent event. AI Gateway captures each batch of events and translates them into the AI Gateway inference format. This ensures that all providers are compatible with the same framework including OpenAI-compatible SDKs or similar.

In a standard LLM transaction, requests proxied directly to the LLM look like this:

 
sequenceDiagram
  actor Client
  participant AI Gateway
  Client->>+AI Gateway:
  destroy AI Gateway
  AI Gateway->>+Cloud LLM: Sends proxy request information
  Cloud LLM->>+Client: Sends chunk to client
  

When streaming is requested, requests proxied directly to the LLM look like this:

 
flowchart LR
  A(client)
  B(AI Gateway)
  C(Cloud LLM)
  D[[transform frame]]
  E[[read frame]]

subgraph main
direction LR
  subgraph 1
  A
  end
  subgraph 3
  C
  end
  subgraph 2
  D
  E
  end
  A --> B --request--> C
  C --response--> B
  B --> D-->E
  E --> B
  B --> A
end

  linkStyle 2,3,4,5,6 stroke:#b6d7a8,color:#b6d7a8
  style 1 color:#fff,stroke:#fff
  style 2 color:#fff,stroke:#fff
  style 3 color:#fff,stroke:#fff
  style main color:#fff,stroke:#fff
  

The streaming framework captures each event, sends the chunk back to the client, and then exits early.

It also estimates tokens for LLM services that decided to not stream back the token use counts when the message is completed.

Keep the following limitations in mind when you configure streaming for the AI Gateway:

• Multiple AI features shouldn’t be expected to be applied and work simultaneously.
• You can’t add AI Policies that use the Response Transformer Policy or otherwise trigger in the response phase when streaming is configured.
• The AI Request Transformer Policy will work, but the AI Response Transformer Policy will not. This is because AI Gateway can’t check every single response token against a separate system.

Streaming is already enabled on AI Gateway; all you have to do is add streaming to your request.

The following is an example llm/v1/completions route streaming request:

{
  "prompt": "What is the theory of relativity?",
  "stream": true
}

You should receive each batch of tokens as HTTP chunks, each containing one or many server-sent events.

Token usage in streaming responses

You can receive token usage statistics in an SSE streaming response. Set the following parameter in the request JSON:

{
  "stream_options": {
    "include_usage": true
  }
}

When you set this parameter, the usage object appears in the final SSE frame, before the [DONE] terminator. This object contains token count statistics for the request.

The following example shows how to request and process token usage statistics in a streaming response:

from openai import OpenAI

client = OpenAI(
    base_url="http://127.0.0.1:8000/openai",
    api_key="none"
)

stream = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Tell me the history of Kong Inc."}],
    stream=True,
    stream_options={"include_usage": True}
)

for chunk in stream:
    if chunk.choices and chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)
    if chunk.usage:
        print("\nDONE. Usage stats:\n")
        print(chunk.usage)

This feature works with any AI Model Provider and AI Model when formats is set to openai mode.

See the OpenAI API Documentation for more information on stream options.

Response streaming configuration parameters

In the AI Model configuration, you can set an optional field config.response_streaming to one of three values: