complete migration of all shield results

genia-dev · Dec 17, 2023 · b725e84 · b725e84
1 parent 9368f04
commit b725e84
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 3 deletions.
diff --git a/vibraniumdome-shields/tests/shields/test_vibranium_sensitive_shield.py b/vibraniumdome-shields/tests/shields/test_vibranium_sensitive_shield.py
@@ -29,9 +29,8 @@ def create_interaction_from_string(self, str):
 
     def test_no_sensitive(self):
         test = "Is there anything else I can help you with?"
-        expected = []
         result = self._sensitive_shield.deflect(self.create_interaction_from_string(test), self._policy, self._scan_id, self._full_policy)
-        self.assertEqual(result, expected)
+        self.assertEqual(result[0].risk, 0)
 
     def test_sensitive_date(self):
         test = "Good morning!"

diff --git a/vibraniumdome-shields/vibraniumdome_shields/shields/input/semantic_similarity_shield.py b/vibraniumdome-shields/vibraniumdome_shields/shields/input/semantic_similarity_shield.py
@@ -43,4 +43,8 @@ def deflect(self, llm_interaction: LLMInteraction, shield_policy_config: dict, s
         except Exception as err:
             self.logger.exception("Failed to perform vector shield, scan_id=%d", scan_id)
             raise err
+
+        if len(shield_matches) == 0:
+            shield_matches.append(SemanticSimilarityShieldDeflectionResult())
+
         return shield_matches
diff --git a/vibraniumdome-shields/vibraniumdome_shields/shields/regex_shield_base.py b/vibraniumdome-shields/vibraniumdome_shields/shields/regex_shield_base.py
@@ -61,4 +61,7 @@ def deflect(self, llm_interaction: LLMInteraction, shield_policy_config: dict, s
                 except Exception:
                     self.logger.exception("Failed to evaluate regex pattern %s", pattern)
 
+        if len(shield_matches) == 0:
+            shield_matches.append(RegexShieldDeflectionResult(pattern="any"))
+
         return shield_matches
diff --git a/...aniumdome-shields/vibraniumdome_shields/shields/sensitive_information_disclosoure_base.py b/...aniumdome-shields/vibraniumdome_shields/shields/sensitive_information_disclosoure_base.py
@@ -42,8 +42,14 @@ def deflect(self, llm_interaction: LLMInteraction, shield_policy_config: dict, s
                 result = analyzer.analyze(text=message, entities=self._entites, language="en")
                 if result:
                     shield_matches.append(
-                        SensitiveShieldDeflectionResult(result=[{key: value for key, value in r.to_dict().items() if key != "recognition_metadata"} for r in result])
+                        SensitiveShieldDeflectionResult(
+                            result=[{key: value for key, value in r.to_dict().items() if key != "recognition_metadata"} for r in result]
+                        )
                     )
         except Exception:
             self._logger.exception("Presidio Analyzer error, scan_id: %s", scan_id)
+
+        if len(shield_matches) == 0:
+            shield_matches.append(SensitiveShieldDeflectionResult(result=[]))
+
         return shield_matches