fix(ai-proxy): (Gemini)(AG-154): content-safety failure no longer blo…

…cks gemini in ai-proxy only

changelog/unreleased/kong/ai-gemini-blocks-content-safety.yml

@@ -0,0 +1,3 @@
+message: "**ai-proxy**: Fixed a bug where Gemini provider would return an error if content safety failed in AI Proxy."
+type: bugfix
+scope: Plugin

kong/llm/drivers/gemini.lua

@@ -32,6 +32,14 @@ local function to_gemini_generation_config(request_table)
   }
 end
 
+local function is_content_safety_failure(content)
+  return content
+          and content.candidates
+          and #content.candidates > 0
+          and content.candidates[1].finishReason
+          and content.candidates[1].finishReason == "SAFETY"
+end
+
 local function is_response_content(content)
   return content
         and content.candidates
@@ -214,8 +222,6 @@ local function to_gemini_chat_openai(request_table, model_info, route_type)
     new_r.tools = request_table.tools and to_tools(request_table.tools)
   end
 
-  kong.log.warn(cjson.encode(new_r))
-
   return new_r, "application/json", nil
 end
 
@@ -236,7 +242,16 @@ local function from_gemini_chat_openai(response, model_info, route_type)
   messages.choices = {}
 
   if response.candidates and #response.candidates > 0 then
-    if is_response_content(response) then
+    -- for transformer plugins only
+    if model_info.source
+          and (model_info.source == "ai-request-transformer" or model_info.source == "ai-response-transformer")
+          and is_content_safety_failure(response) then
+      local err = "transformation generation candidate breached Gemini content safety"
+      ngx.log(ngx.ERR, err)
+
+      return nil, err
+
+    elseif is_response_content(response) then
       messages.choices[1] = {
         index = 0,
         message = {
@@ -277,14 +292,6 @@ local function from_gemini_chat_openai(response, model_info, route_type)
       }
     end
 
-  elseif response.candidates
-           and #response.candidates > 0
-           and response.candidates[1].finishReason
-           and response.candidates[1].finishReason == "SAFETY" then
-    local err = "transformation generation candidate breached Gemini content safety"
-    ngx.log(ngx.ERR, err)
-    return nil, err
-
   else -- probably a server fault or other unexpected response
     local err = "no generation candidates received from Gemini, or max_tokens too short"
     ngx.log(ngx.ERR, err)

kong/plugins/ai-request-transformer/handler.lua

@@ -65,6 +65,7 @@ function _M:access(conf)
   local http_opts = create_http_opts(conf)
   conf.llm.__plugin_id = conf.__plugin_id
   conf.llm.__key__ = conf.__key__
+  conf.llm.model.source = "ai-request-transformer"
   local ai_driver, err = llm.new_driver(conf.llm, http_opts, identity_interface)
 
   if not ai_driver then

kong/plugins/ai-response-transformer/handler.lua

@@ -124,6 +124,7 @@ function _M:access(conf)
   local http_opts = create_http_opts(conf)
   conf.llm.__plugin_id = conf.__plugin_id
   conf.llm.__key__ = conf.__key__
+  conf.llm.model.source = "ai-response-transformer"
   local ai_driver, err = llm.new_driver(conf.llm, http_opts, identity_interface)
 
   if not ai_driver then