From 8c4610f4eebfd5c92207ef04f59c5fea7ed49e6a Mon Sep 17 00:00:00 2001
From: Gayathri <142381512+gaya3-zipstack@users.noreply.github.com>
Date: Mon, 11 Nov 2024 11:55:50 +0530
Subject: [PATCH] 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 <142381512+gaya3-zipstack@users.noreply.github.com>
---
 src/unstract/sdk/__init__.py                      |  2 +-
 .../llm/vertex_ai/src/static/json_schema.json     | 15 ++++++++++-----
 .../sdk/adapters/llm/vertex_ai/src/vertex_ai.py   |  1 +
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/unstract/sdk/__init__.py b/src/unstract/sdk/__init__.py
index 07ecf5ee..e3f0e0c3 100644
--- 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
index 470df588..78076135 100644
--- 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
index 5fb30382..e333fc8f 100644
--- 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,
 }