fix(ai-proxy): gemini streaming transformer bug

Kong · Jun 28, 2024 · d312ce6 · d312ce6
1 parent 26202ae
commit d312ce6
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 21 deletions.
diff --git a/kong/llm/drivers/gemini.lua b/kong/llm/drivers/gemini.lua
@@ -136,11 +136,9 @@ local function to_gemini_chat_openai(request_table, model_info, route_type)
         }
       end
     end
-
+    
     new_r.generationConfig = to_gemini_generation_config(request_table)
 
-    kong.log.debug(cjson.encode(new_r))
-
     return new_r, "application/json", nil
   end
 

diff --git a/kong/llm/drivers/shared.lua b/kong/llm/drivers/shared.lua
@@ -258,29 +258,31 @@ function _M.frame_to_events(frame, raw_json_mode)
     local event_lines = split(frame, "\n")
     local struct = { event = nil, id = nil, data = nil }
 
-    -- test for truncated chunk on the last line (no trailing \r\n\r\n)
-    if #dat > 0 and #event_lines == i then
-      ngx.log(ngx.DEBUG, "[ai-proxy] truncated sse frame head")
-      kong.ctx.plugin.truncated_frame = dat
-      break  -- stop parsing immediately, server has done something wrong
-    end
+    for i, dat in ipairs(event_lines) do
+      if #dat < 1 then
+        events[#events + 1] = struct
+        struct = { event = nil, id = nil, data = nil }
+      end
 
-    -- test for abnormal start-of-frame (truncation tail)
-    if kong and kong.ctx.plugin.truncated_frame then
-      -- this is the tail of a previous incomplete chunk
-      ngx.log(ngx.DEBUG, "[ai-proxy] truncated sse frame tail")
-      dat = fmt("%s%s", kong.ctx.plugin.truncated_frame, dat)
-      kong.ctx.plugin.truncated_frame = nil
-    end
+      -- test for truncated chunk on the last line (no trailing \r\n\r\n)
+      if #dat > 0 and #event_lines == i then
+        ngx.log(ngx.DEBUG, "[ai-proxy] truncated sse frame head")
+        kong.ctx.plugin.truncated_frame = dat
+        break  -- stop parsing immediately, server has done something wrong
+      end
 
-    local s1, _ = str_find(dat, ":") -- find where the cut point is
+      -- test for abnormal start-of-frame (truncation tail)
+      if kong and kong.ctx.plugin.truncated_frame then
+        -- this is the tail of a previous incomplete chunk
+        ngx.log(ngx.DEBUG, "[ai-proxy] truncated sse frame tail")
+        dat = fmt("%s%s", kong.ctx.plugin.truncated_frame, dat)
+        kong.ctx.plugin.truncated_frame = nil
+      end
 
-    if s1 and s1 ~= 1 then
-      local field = str_sub(dat, 1, s1-1) -- returns "data" from data: hello world
-      local value = str_ltrim(str_sub(dat, s1+1)) -- returns "hello world" from data: hello world
+      local s1, _ = str_find(dat, ":") -- find where the cut point is
 
       if s1 and s1 ~= 1 then
-        local field = str_sub(dat, 1, s1-1) -- returns "data " from data: hello world
+        local field = str_sub(dat, 1, s1-1) -- returns "data" from data: hello world
         local value = str_ltrim(str_sub(dat, s1+1)) -- returns "hello world" from data: hello world
 
         -- for now not checking if the value is already been set