fix(ai-proxy): bedrock streaming format detection failure

Kong · Jul 8, 2024 · 2dd382b · 2dd382b
1 parent a660900
commit 2dd382b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 31 deletions.
diff --git a/kong/llm/drivers/bedrock.lua b/kong/llm/drivers/bedrock.lua
@@ -103,26 +103,11 @@ local function handle_stream_event(event_t, model_info, route_type)
     }
 
   elseif event_type == "metadata" then
-    local function dump(o)
-      if type(o) == 'table' then
-         local s = '{ '
-         for k,v in pairs(o) do
-            if type(k) ~= 'number' then k = '"'..k..'"' end
-            s = s .. '['..k..'] = ' .. dump(v) .. ','
-         end
-         return s .. '} '
-      else
-         return tostring(o)
-      end
-    end
-    kong.log.warn(dump(body))
     metadata = {
       prompt_tokens = body.usage and body.usage.inputTokens or 0,
       completion_tokens = body.usage and body.usage.outputTokens or 0,
     }
 
-    kong.log.warn(dump(metadata))
-
     new_event = "[DONE]"
 
   elseif event_type == "contentBlockStop" then
@@ -184,8 +169,6 @@ local function to_bedrock_chat_openai(request_table, model_info, route_type)
 
   new_r.inferenceConfig = to_bedrock_generation_config(request_table)
 
-  kong.log.debug(new_r.inferenceConfig.maxTokens)
-
   return new_r, "application/json", nil
 end
 

diff --git a/kong/llm/drivers/shared.lua b/kong/llm/drivers/shared.lua
@@ -277,8 +277,8 @@ function _M.frame_to_events(frame, provider)
       }
     end
 
-  -- standard SSE parser
   else
+    -- standard SSE parser
     local event_lines = split(frame, "\n")
     local struct = { event = nil, id = nil, data = nil }
 
@@ -287,9 +287,6 @@ function _M.frame_to_events(frame, provider)
         events[#events + 1] = struct
         struct = { event = nil, id = nil, data = nil }
       end
-    else
-      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
@@ -312,17 +309,12 @@ function _M.frame_to_events(frame, provider)
         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
 
-        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
-
-          -- for now not checking if the value is already been set
-          if     field == "event" then struct.event = value
-          elseif field == "id"    then struct.id = value
-          elseif field == "data"  then struct.data = value
-          end -- if
+        -- for now not checking if the value is already been set
+        if     field == "event" then struct.event = value
+        elseif field == "id"    then struct.id = value
+        elseif field == "data"  then struct.data = value
         end -- if
-      end
+      end -- if
     end
   end