From 812d9102128e86789bee2eb588ab44bccbd42ccb Mon Sep 17 00:00:00 2001
From: Azalea <22280294+hykilpikonna@users.noreply.github.com>
Date: Thu, 26 Dec 2024 07:55:54 -0500
Subject: [PATCH] [F] Fix safety :(

---
 src/main/java/icu/samnyan/aqua/net/Safety.kt | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/main/java/icu/samnyan/aqua/net/Safety.kt b/src/main/java/icu/samnyan/aqua/net/Safety.kt
index da6636c2..ecf6bf6a 100644
--- a/src/main/java/icu/samnyan/aqua/net/Safety.kt
+++ b/src/main/java/icu/samnyan/aqua/net/Safety.kt
@@ -54,6 +54,7 @@ class AquaNetSafetyService(
      */
     suspend fun isSafeBatch(rawContents: List<String>): List<Boolean> {
         val contents = rawContents.map { Normalizer.normalize(it, Normalizer.Form.NFKC) }
+        val origMap = safety.findAll().associateBy { it.content }.toMutableMap()
         val map = safety.findAll().associateBy { it.content.lowercase().trim() }.toMutableMap()
 
         // Process unseen content with OpenAI
@@ -72,8 +73,9 @@ class AquaNetSafetyService(
             }
         }
         if (news.isNotEmpty()) safety.saveAll(news)
+        news.associateByTo(origMap) { it.content }
         news.associateByTo(map) { it.content.lowercase().trim() }
 
-        return contents.map { map[it.lowercase().trim()]!!.safe }
+        return contents.map { map[it.lowercase().trim()]?.safe ?: origMap[it]?.safe ?: true }
     }
 }