[Feature] Implement manual ANN train mode

diff --git a/src/plugins/lua/fann_redis.lua b/src/plugins/lua/fann_redis.lua
index 65d354b21540b70e743ad63c04e2cdf4f90bd92e..2ae9c19d84b20599ac848c3715199be16d52530b 100644 (file)
--- a/src/plugins/lua/fann_redis.lua
+++ b/src/plugins/lua/fann_redis.lua
@@ -456,15 +456,28 @@ local function fann_train_callback(rule, task, score, required_score, id)
 
   local learn_spam, learn_ham
 
-  if train_opts['spam_score'] then
-    learn_spam = score >= train_opts['spam_score']
-  else
-    learn_spam = score >= required_score
-  end
-  if train_opts['ham_score'] then
-    learn_ham = score <= train_opts['ham_score']
+  if rule.autotrain then
+    if train_opts['spam_score'] then
+      learn_spam = score >= train_opts['spam_score']
+    else
+      learn_spam = score >= required_score
+    end
+    if train_opts['ham_score'] then
+      learn_ham = score <= train_opts['ham_score']
+    else
+      learn_ham = score < 0
+    end
   else
-    learn_ham = score < 0
+    -- Train by request header
+    local hdr = task:get_request_header('ANN-Train')
+
+    if hdr then
+      if hdr:lower() == 'spam' then
+        learn_spam = true
+      elseif hdr:lower() == 'ham' then
+        learn_ham = true
+      end
+    end
   end
 
   if learn_spam or learn_ham then