Skip to content

Commit

Permalink
apply suggestions
Browse files Browse the repository at this point in the history
  • Loading branch information
Revolyssup committed Mar 5, 2025
1 parent 88f9996 commit ac03b08
Showing 1 changed file with 31 additions and 13 deletions.
44 changes: 31 additions & 13 deletions apisix/plugins/ai-prompt-guard.lua
Original file line number Diff line number Diff line change
Expand Up @@ -82,20 +82,16 @@ end
local function get_content_to_check(conf, messages)
local contents = {}
if conf.match_all_conversation_history then
for _, msg in ipairs(messages) do
if msg.content then
core.table.insert(contents, msg.content)
end
end
return messages
else
if #messages > 0 then
local last_msg = messages[#messages]
if last_msg.content then
core.table.insert(contents, last_msg.content)
if last_msg then
core.table.insert(contents, last_msg)
end
end
end
return table.concat(contents, " ")
return contents
end

function _M.access(conf, ctx)
Expand All @@ -111,12 +107,34 @@ function _M.access(conf, ctx)
end

local messages = json_body.messages or {}
if not conf.match_all_roles and #messages > 0 and messages[#messages].role ~= "user" then
return
messages = get_content_to_check(conf, messages)
if not conf.match_all_roles then
-- filter to only user messages
local new_messages = {}
for _, msg in ipairs(messages) do
if not msg then
return 400, {message = "request doesn't contain messages"}
end
if msg.role == "user" then
core.table.insert(new_messages, msg)
end
end
messages = new_messages
end

local content_to_check = get_content_to_check(conf, messages)

if #messages == 0 then --nothing to check
return 200
end
-- extract only messages
local content = {}
for _, msg in ipairs(messages) do
if not msg then
return 400, {message = "request doesn't contain messages"}
end
if msg.content then
core.table.insert(content, msg.content)
end
end
local content_to_check = table.concat(content, " ")
-- Allow patterns check
if #conf.allow_patterns > 0 then
local any_allowed = false
Expand Down

0 comments on commit ac03b08

Please # to comment.