From 3f9859e638e4ea92a22293f0d8669854bb24d9b4 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 29 Jun 2024 19:57:41 -0700
Subject: [PATCH] fix(utils.py): add 'enforce_validation' param

---
 .../tests/test_amazing_vertex_completion.py   | 41 +++++++++----------
 litellm/utils.py                              |  8 +++-
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py
index 505933fc24d..807fb31318f 100644
--- a/litellm/tests/test_amazing_vertex_completion.py
+++ b/litellm/tests/test_amazing_vertex_completion.py
@@ -1003,9 +1003,17 @@ def vertex_httpx_mock_post_invalid_schema_response(*args, **kwargs):
     "invalid_response",
     [True, False],
 )
+@pytest.mark.parametrize(
+    "enforce_validation",
+    [True, False],
+)
 @pytest.mark.asyncio
 async def test_gemini_pro_json_schema_args_sent_httpx(
-    model, supports_response_schema, vertex_location, invalid_response
+    model,
+    supports_response_schema,
+    vertex_location,
+    invalid_response,
+    enforce_validation,
 ):
     load_vertex_ai_credentials()
     os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@@ -1015,27 +1023,17 @@ async def test_gemini_pro_json_schema_args_sent_httpx(
     messages = [{"role": "user", "content": "List 5 cookie recipes"}]
     from litellm.llms.custom_httpx.http_handler import HTTPHandler
 
-    # response_schema = {
-    #     "type": "array",
-    #     "items": {
-    #         "type": "object",
-    #         "properties": {
-    #             "recipe_name": {
-    #                 "type": "string",
-    #             },
-    #         },
-    #         "required": ["recipe_name"],
-    #     },
-    # }
     response_schema = {
-        "type": "object",
-        "properties": {
-            "recipe_name": {"type": "string"},
-            "ingredients": {"type": "array", "items": {"type": "string"}},
-            "prep_time": {"type": "number"},
-            "difficulty": {"type": "string", "enum": ["easy", "medium", "hard"]},
+        "type": "array",
+        "items": {
+            "type": "object",
+            "properties": {
+                "recipe_name": {
+                    "type": "string",
+                },
+            },
+            "required": ["recipe_name"],
         },
-        "required": ["recipe_name", "ingredients", "prep_time"],
     }
 
     client = HTTPHandler()
@@ -1053,11 +1051,12 @@ async def test_gemini_pro_json_schema_args_sent_httpx(
                 response_format={
                     "type": "json_object",
                     "response_schema": response_schema,
+                    "enforce_validation": enforce_validation,
                 },
                 vertex_location=vertex_location,
                 client=client,
             )
-            if invalid_response is True:
+            if invalid_response is True and enforce_validation is True:
                 pytest.fail("Expected this to fail")
         except litellm.JSONSchemaValidationError as e:
             if invalid_response is False and "claude-3" not in model:
diff --git a/litellm/utils.py b/litellm/utils.py
index 87768541678..2598aa46b14 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -621,8 +621,14 @@ def client(original_function):
                                         ],
                                         dict,
                                     )
+                                    and "enforce_validation"
+                                    in optional_params["response_format"]
+                                    and optional_params["response_format"][
+                                        "enforce_validation"
+                                    ]
+                                    is True
                                 ):
-                                    # schema given, json response expected
+                                    # schema given, json response expected, and validation enforced
                                     litellm.litellm_core_utils.json_validation_rule.validate_schema(
                                         schema=optional_params["response_format"][
                                             "response_schema"