Support OFF as a threshold value for Safety settings in Vertex AI (#114)

* Safety threshold changes for Verted toswitch OFF * Add OFF for all safety settings --------- Signed-off-by: Gayathri <[email protected]>
Zipstack · Nov 11, 2024 · 8c4610f · 8c4610f
1 parent 2d82112
commit 8c4610f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/src/unstract/sdk/__init__.py b/src/unstract/sdk/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.52.1"
+__version__ = "0.53.0"
 
 
 def get_sdk_version():

diff --git a/src/unstract/sdk/adapters/llm/vertex_ai/src/static/json_schema.json b/src/unstract/sdk/adapters/llm/vertex_ai/src/static/json_schema.json
@@ -61,7 +61,8 @@
             "BLOCK_LOW_AND_ABOVE",
             "BLOCK_MEDIUM_AND_ABOVE",
             "BLOCK_ONLY_HIGH",
-            "BLOCK_NONE"
+            "BLOCK_NONE",
+            "OFF"
           ],
           "default": "BLOCK_ONLY_HIGH",
           "description": "Settings for HARM_CATEGORY_DANGEROUS_CONTENT"
@@ -74,7 +75,8 @@
             "BLOCK_LOW_AND_ABOVE",
             "BLOCK_MEDIUM_AND_ABOVE",
             "BLOCK_ONLY_HIGH",
-            "BLOCK_NONE"
+            "BLOCK_NONE",
+            "OFF"
           ],
           "default": "BLOCK_ONLY_HIGH",
           "description": "Settings for HARM_CATEGORY_HATE_SPEECH"
@@ -87,7 +89,8 @@
             "BLOCK_LOW_AND_ABOVE",
             "BLOCK_MEDIUM_AND_ABOVE",
             "BLOCK_ONLY_HIGH",
-            "BLOCK_NONE"
+            "BLOCK_NONE",
+            "OFF"
           ],
           "default": "BLOCK_ONLY_HIGH",
           "description": "Settings for HARM_CATEGORY_HARASSMENT"
@@ -100,7 +103,8 @@
             "BLOCK_LOW_AND_ABOVE",
             "BLOCK_MEDIUM_AND_ABOVE",
             "BLOCK_ONLY_HIGH",
-            "BLOCK_NONE"
+            "BLOCK_NONE",
+            "OFF"
           ],
           "default": "BLOCK_ONLY_HIGH",
           "description": "Settings for HARM_CATEGORY_SEXUAL_CONTENT"
@@ -113,7 +117,8 @@
             "BLOCK_LOW_AND_ABOVE",
             "BLOCK_MEDIUM_AND_ABOVE",
             "BLOCK_ONLY_HIGH",
-            "BLOCK_NONE"
+            "BLOCK_NONE",
+            "OFF"
           ],
           "default": "BLOCK_ONLY_HIGH",
           "description": "Settings for HARM_CATEGORY_OTHER"

diff --git a/src/unstract/sdk/adapters/llm/vertex_ai/src/vertex_ai.py b/src/unstract/sdk/adapters/llm/vertex_ai/src/vertex_ai.py
@@ -45,6 +45,7 @@ class SafetySettingsConstants:
     "BLOCK_MEDIUM_AND_ABOVE": HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
     "BLOCK_ONLY_HIGH": HarmBlockThreshold.BLOCK_ONLY_HIGH,
     "BLOCK_NONE": HarmBlockThreshold.BLOCK_NONE,
+    "OFF": HarmBlockThreshold.OFF,
 }