Add a ML-based classifier for important messages

Signed-off-by: Christoph Wurst <christoph@winzerhof-wurst.at>

1 files changed, 10 insertions, 4 deletions

diff --git a/lib/Listener/NewMessageClassificationListener.php b/lib/Listener/NewMessageClassificationListener.php
index a78e0298c..85c10aaa8 100644
--- a/lib/Listener/NewMessageClassificationListener.php
+++ b/lib/Listener/NewMessageClassificationListener.php
@@ -27,7 +27,7 @@ namespace OCA\Mail\Listener;
 
 use Horde_Imap_Client;
 use OCA\Mail\Events\NewMessagesSynchronized;
-use OCA\Mail\Service\Classification\MessageClassifier;
+use OCA\Mail\Service\Classification\ImportanceClassifier;
 use OCP\EventDispatcher\Event;
 use OCP\EventDispatcher\IEventListener;
 
@@ -40,10 +40,10 @@ class NewMessageClassificationListener implements IEventListener {
 		Horde_Imap_Client::SPECIALUSE_TRASH,
 	];
 
-	/** @var MessageClassifier */
+	/** @var ImportanceClassifier */
 	private $classifier;
 
-	public function __construct(MessageClassifier $classifier) {
+	public function __construct(ImportanceClassifier $classifier) {
 		$this->classifier = $classifier;
 	}
 
@@ -59,8 +59,14 @@ class NewMessageClassificationListener implements IEventListener {
 			}
 		}
 
+		$predictions = $this->classifier->classifyImportance(
+			$event->getAccount(),
+			$event->getMailbox(),
+			$event->getMessages()
+		);
+
 		foreach ($event->getMessages() as $message) {
-			if ($this->classifier->isImportant($event->getAccount(), $event->getMailbox(), $message)) {
+			if ($predictions[$message->getUid()] ?? false) {
 				$message->setFlagImportant(true);
 			}
 		}