Fixed missed comma

trustgraph-ai · Oct 9, 2024 · 6bfa01d · 6bfa01d
1 parent b98e553
commit 6bfa01d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py
@@ -55,7 +55,7 @@
         "openai",
         "neo4j",
         "tiktoken",
-        "google-generativeai"
+        "google-generativeai",
     ],
     scripts=[
         "scripts/chunker-recursive",

diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py
@@ -74,12 +74,14 @@ def __init__(self, **params):
             "response_mime_type": "text/plain",
         }
 
+        block_level = HarmBlockThreshold.BLOCK_ONLY_HIGH
+
         self.safety_settings={
-            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-            HarmCategory.HARM_CATEGORY_DANGEROUS: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-            HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+            HarmCategory.HARM_CATEGORY_HATE_SPEECH: block_level,
+            HarmCategory.HARM_CATEGORY_HARASSMENT: block_level,
+            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: block_level,
+            HarmCategory.HARM_CATEGORY_DANGEROUS: block_level,
+            HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY: block_level,
         }
 
         self.llm = genai.GenerativeModel(