]> git.ipfire.org Git - thirdparty/rspamd.git/commitdiff
[Feature] GPT: Try harder to find JSON in NN reply
authorVsevolod Stakhov <vsevolod@rspamd.com>
Tue, 10 Dec 2024 15:09:07 +0000 (15:09 +0000)
committerVsevolod Stakhov <vsevolod@rspamd.com>
Tue, 10 Dec 2024 15:09:07 +0000 (15:09 +0000)
src/plugins/lua/gpt.lua

index 10cff0fb561542b4f5146b8cc158892198dcf9d6..36938c0d187c50dfdbcb9e94eabc1944724b01d3 100644 (file)
@@ -153,6 +153,40 @@ local function default_condition(task)
   return true, sel_part:get_content_oneline()
 end
 
+local function maybe_extract_json(str)
+  -- Find the first opening brace
+  local startPos = str:find("{")
+  if not startPos then
+    return nil
+  end
+
+  local openBraces = 0
+  local endPos = startPos
+  local len = #str
+
+  -- Iterate through the string to find matching braces
+  for i = startPos, len do
+    local char = str:sub(i, i)
+    if char == "{" then
+      openBraces = openBraces + 1
+    elseif char == "}" then
+      openBraces = openBraces - 1
+      -- When we find the matching closing brace
+      if openBraces == 0 then
+        endPos = i
+        break
+      end
+    end
+  end
+
+  -- If we found a complete JSON-like structure
+  if openBraces == 0 then
+    return str:sub(startPos, endPos)
+  end
+
+  return nil
+end
+
 local function default_conversion(task, input)
   local parser = ucl.parser()
   local res, err = parser:parse_string(input)
@@ -178,6 +212,9 @@ local function default_conversion(task, input)
     return
   end
 
+  -- Apply heuristic to extract JSON
+  first_message = maybe_extract_json(first_message) or first_message
+
   parser = ucl.parser()
   res, err = parser:parse_string(first_message)
   if not res then