From 34cd042f097702bf2551b372664d2e440e1962eb Mon Sep 17 00:00:00 2001
From: Jack Tysoe <jack.tysoe@konghq.com>
Date: Tue, 26 Nov 2024 23:20:37 +0000
Subject: [PATCH] fix(AG-178): properly support all streaming content types

---
 kong/llm/drivers/shared.lua                        | 6 ++++++
 kong/llm/plugin/shared-filters/parse-sse-chunk.lua | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/kong/llm/drivers/shared.lua b/kong/llm/drivers/shared.lua
index 0a1ec6d6496..03c00bbcddb 100644
--- a/kong/llm/drivers/shared.lua
+++ b/kong/llm/drivers/shared.lua
@@ -78,6 +78,12 @@ _M._CONST = {
   ["AWS_STREAM_CONTENT_TYPE"] = "application/vnd.amazon.eventstream",
 }
 
+_M._SUPPORTED_STREAMING_CONTENT_TYPES = {
+  ["text/event-stream"] = true,
+  ["application/vnd.amazon.eventstream"] = true,
+  ["application/json"] = true,
+}
+
 _M.streaming_has_token_counts = {
   ["cohere"] = true,
   ["llama2"] = true,
diff --git a/kong/llm/plugin/shared-filters/parse-sse-chunk.lua b/kong/llm/plugin/shared-filters/parse-sse-chunk.lua
index 219a5017259..b0c9195a8e2 100644
--- a/kong/llm/plugin/shared-filters/parse-sse-chunk.lua
+++ b/kong/llm/plugin/shared-filters/parse-sse-chunk.lua
@@ -20,7 +20,7 @@ local function handle_streaming_frame(conf, chunk, finished)
 
   local content_type = kong.service.response.get_header("Content-Type")
   local normalized_content_type = content_type and content_type:sub(1, (content_type:find(";") or 0) - 1)
-  if normalized_content_type and normalized_content_type ~= "text/event-stream" and normalized_content_type ~= ai_shared._CONST.AWS_STREAM_CONTENT_TYPE then
+  if normalized_content_type and (not ai_shared._SUPPORTED_STREAMING_CONTENT_TYPES[normalized_content_type]) then
     return true
   end