phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(utils.py): add 'enforce_validation' param

Browse source
This commit is contained in:
Krrish Dholakia 2024-06-29 19:12:00 -07:00
parent 69d06cfbcd
commit 7670c5bd13
2 changed files with 14 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/tests/test_amazing_vertex_completion.py
View file

@ -1003,9 +1003,13 @@ def vertex_httpx_mock_post_invalid_schema_response(*args, **kwargs):
"invalid_response", "invalid_response",
[True, False], [True, False],
) )
@pytest.mark.parametrize(
"enforce_validation",
[True, False],
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_gemini_pro_json_schema_args_sent_httpx( async def test_gemini_pro_json_schema_args_sent_httpx(
model, supports_response_schema, invalid_response model, supports_response_schema, invalid_response, enforce_validation
): ):
load_vertex_ai_credentials() load_vertex_ai_credentials()
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@ -1043,10 +1047,11 @@ async def test_gemini_pro_json_schema_args_sent_httpx(
response_format={ response_format={
"type": "json_object", "type": "json_object",
"response_schema": response_schema, "response_schema": response_schema,
"enforce_validation": enforce_validation,
}, },
client=client, client=client,
) )
if invalid_response is True: if invalid_response is True and enforce_validation is True:
pytest.fail("Expected this to fail") pytest.fail("Expected this to fail")
except litellm.JSONSchemaValidationError as e: except litellm.JSONSchemaValidationError as e:
if invalid_response is False: if invalid_response is False:

8
litellm/utils.py
View file

@ -621,8 +621,14 @@ def client(original_function):
], ],
dict, 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( litellm.litellm_core_utils.json_validation_rule.validate_schema(
schema=optional_params["response_format"][ schema=optional_params["response_format"][
"response_schema" "response_schema"