ai-proxy buffers streamed responses

[agarza22]

Is there an existing issue for this?

• I have searched the existing issues

Kong version ($ kong version)

3.6

Current Behavior

When using the ai-proxy plugin, streamed responses are buffered by Kong before being returned to the client.

Expected Behavior

Streamed responses should not be buffered, instead being streamed back to the client.

Steps To Reproduce

Running Kong 3.6 with the ai-proxy plugin enabled, make a streaming API call (in my case to AzureOpenAI).

Response chunks will be buffered and then returned to client all at once.

Anything else?

No response

[chobits]

This plugin is introduced by #12323,

Hi @tysoekong, could you take a preliminary look at the code, and the non-stream output seems to be as expetect. It seems that the response content in the code is completely rebuilt.

[subnetmarco]

@agarza22 we are planning to introduce support for streaming in the next minor release of Kong Gateway (3.7).

[agarza22]

@subnetmarco That's great to hear! Do you have a ballpark on when we'll see that release? Does this also mean the ai-proxy plugin will get an update to support the streaming use case?

[subnetmarco]

@agarza22 in May most likely.

[ttyS0e]

@agarza22 @chobits Sorry for direct mention, but I see your interest in this feature.

We have added the streaming support, which currently is in review. Code is subject to change slightly.

You can package the streaming-enabled ai-proxy plugin into the existing Kong 3.6.1 image, using (for example) this builder:

FROM kong:3.6.1 as builder

USER root
WORKDIR /builder

RUN apt update && \
    apt install -y zip unzip git

RUN git clone -b 'feat/KAG-4126-ai-proxy-streaming' https://github.com/Kong/kong.git

#---#

FROM kong:3.6.1

USER root

COPY --from=builder --chown=1001:1001 \
     /builder/kong/kong/plugins/ai-proxy \
     /usr/local/share/lua/5.1/kong/plugins/ai-proxy

COPY --from=builder --chown=1001:1001 \
     /builder/kong/kong/llm \
     /usr/local/share/lua/5.1/kong/llm

USER kong

Then you simply add "stream": true to the JSON in your request, and it should print tokens (SSE events) as they are transmitted.

Hope this helps you to start testing it out?